def __getitem__()

in c3dm/dataset/keypoints_dataset.py [0:0]

• 74 lines of code
• 7 McCabe index (conditional complexity)

	def __getitem__(self, index):

		assert index < len(self.db), \
			'index %d out of range (%d)' % (index, len(self.db))

		entry = copy.deepcopy(self.db[index])

		if self.image_root is not None and 'image_path' in entry:
			entry['image_path'] = os.path.join(self.image_root,entry['image_path'])
		if self.mask_root is not None and 'mask_path' in entry:
			entry['mask_path'] = os.path.join(self.mask_root,entry['mask_path'])
		if self.depth_root is not None and 'depth_path' in entry:
			entry['depth_path'] = os.path.join(self.depth_root,entry['depth_path'])

		if self.load_images:
			entry['images'] = self.load_image(entry)
			entry['orig_image_size'] = list(entry['images'].shape[1:])

		if self.load_depths:
			entry['depths'] = load_depth(entry)

		if self.load_masks:
			entry['masks'] = load_mask(entry)
			if entry['masks'] is None:
				entry['masks'] = np.zeros(entry['images'].shape[1:3] \
										  )[None].astype(np.float32)
			else:
				# assert entry['masks'].shape[1:3]==entry['images'].shape[1:3]
				if self.load_images and \
				   entry['masks'].shape[1:3] != entry['images'].shape[1:3]:
					# print(entry['mask_path'])
					# print(entry['image_path'])
					# import pdb; pdb.set_trace()
					print('bad mask size!!!!')
					# print(entry['image_path'])
					# print(entry['mask_path'])
					# import pdb; pdb.set_trace()
					entry['masks'] = np.zeros(entry['images'].shape[1:3] \
										  )[None].astype(np.float32)

		# convert to torch Tensors where possible
		for fld in ( 'kp_loc', 'kp_vis', 'kp_loc_3d', 
					'class_mask', 'kp_defined', 'images', 
					'orig_image_size', 'masks', 'K', 'depths', 'bbox',
					'kp_conf', 'R', 'T'):
			if fld in entry:
				entry[fld] = torch.FloatTensor(entry[fld])

		# first crop if needed, then resize
		if self.box_crop and self.load_images:
			entry = self.crop_around_box(entry, self.box_crop_context)

		if 'sfm_model' not in entry:
			entry['sfm_model'] = '<NO_MODEL>'
		
		entry['K_orig'] = entry['K'].clone()

		if self.load_images:
			# resize image
			entry['images'], scale = self.resize_image(entry['images'], 
													   mode='bilinear')	
			for fld in ('kp_loc', 'kp_loc_3d', 'K'):
				if fld in entry:
					entry[fld] *= scale
				if fld=='K':
					entry[fld][2,2] = 1. 
		else:
			scale = 1.

		if self.load_masks:
			entry['masks'], _ = self.resize_image(entry['masks'], 
												  mode='nearest')			
			if self.dilate_masks > 0:
				#print('mask dilation')
				entry['masks'] = torch.nn.functional.max_pool2d(
									entry['masks'],
									self.dilate_masks*2+1, 
									stride=1, 
									padding=self.dilate_masks )
			elif self.dilate_masks < 0:
				imask_dil = torch.nn.functional.max_pool2d(
									1-entry['masks'],
									abs(self.dilate_masks)*2+1, 
									stride=1, 
									padding=abs(self.dilate_masks) )
				entry['masks'] = torch.clamp(entry['masks'] - imask_dil, 0.)
				
		if self.load_depths:
			entry['depths'], _ = self.resize_image(entry['depths'], 
												   mode='nearest')
			entry['depths'] *= scale

		if 'p3d_info' in entry:  # filter the kp out of bbox
			bbox = torch.FloatTensor(entry['p3d_info']['bbox'])
			bbox_vis, bbox_err = bbox_kp_visibility( \
									bbox, entry['kp_loc'], entry['kp_vis'])
			entry['kp_vis'] = entry['kp_vis'] * bbox_vis.float()

		# mask out invisible
		entry['kp_loc'] = entry['kp_loc'] * entry['kp_vis'][None]

		return entry