def __getitem__(self, index): img_orig_PIL = self.get_PIL_img(index) img_orig_rgb = np.array(img_orig_PIL) # sample pixels uv_orig = self.get_random_pixels(width=img_orig_rgb.shape[1], height=img_orig_rgb.shape[0], num_samples=self.num_matching_pixels*4) keypoints_on_images_orig = [ia.KeypointsOnImage.from_coords_array(coords=uv_orig, shape=img_orig_rgb.shape)] # augment image to generate a pair of images for training and find correspondences aug_seq_det_a = self.aug_seq.to_deterministic() img_a_rgb = aug_seq_det_a.augment_image(img_orig_rgb) keypoints_a = aug_seq_det_a.augment_keypoints(keypoints_on_images_orig)[0] uv_a = np.rint(keypoints_a.get_coords_array()).astype(int) aug_seq_det_b = self.aug_seq.to_deterministic() img_b_rgb = aug_seq_det_b.augment_image(img_orig_rgb) keypoints_b = aug_seq_det_b.augment_keypoints(keypoints_on_images_orig)[0] uv_b = np.rint(keypoints_b.get_coords_array()).astype(int) # remove pixels outside frame within_a = np.logical_and(uv_a >= [0, 0], uv_a < [img_orig_rgb.shape[1], img_orig_rgb.shape[0]]) within_b = np.logical_and(uv_b >= [0, 0], uv_b < [img_orig_rgb.shape[1], img_orig_rgb.shape[0]]) valid_ids = np.where(np.logical_and(within_a, within_b).all(axis=1))[0] uv_orig = uv_orig[valid_ids][:self.num_matching_pixels] uv_a = uv_a[valid_ids][:self.num_matching_pixels] uv_b = uv_b[valid_ids][:self.num_matching_pixels] if self.debug: self.aug_seq.show_grid(img_orig_rgb, cols=4, rows=4) img_a_cv = cv2.cvtColor(img_a_rgb, cv2.COLOR_RGB2BGR) img_b_cv = cv2.cvtColor(img_b_rgb, cv2.COLOR_RGB2BGR) for pix_a, pix_b in zip(uv_a, uv_b): color = cv2.cvtColor(np.array([[[np.random.randint(256), 255, 255]]], dtype=np.uint8), cv2.COLOR_HSV2RGB)[0][0].astype(int) cv2.circle(img_a_cv, tuple(pix_a), 5, color, -1) cv2.circle(img_b_cv, tuple(pix_b), 5, color, -1) cv2.imshow('a', img_a_cv) cv2.imshow('b', img_b_cv) k = cv2.waitKey(0) # convert to torch Tensor uv_a = (torch.from_numpy(uv_a[:, 0]).type(torch.LongTensor), torch.from_numpy(uv_a[:, 1]).type(torch.LongTensor)) uv_b = (torch.from_numpy(uv_b[:, 0]).type(torch.FloatTensor), torch.from_numpy(uv_b[:, 1]).type(torch.FloatTensor)) # find non_correspondences uv_b_non_matches = correspondence_finder.create_non_correspondences(uv_b, img_b_rgb.shape, num_non_matches_per_match=self.num_non_matches_per_match) # convert PIL.Image to torch.FloatTensor image_a_rgb = self.rgb_image_to_tensor(img_a_rgb) image_b_rgb = self.rgb_image_to_tensor(img_b_rgb) image_width = img_orig_rgb.shape[1] image_height = img_orig_rgb.shape[0] matches_a = SaladDataset.flatten_uv_tensor(uv_a, image_width) matches_b = SaladDataset.flatten_uv_tensor(uv_b, image_width) uv_a_long, uv_b_non_matches_long = self.create_non_matches(uv_a, uv_b_non_matches, self.num_non_matches_per_match) non_matches_a = SaladDataset.flatten_uv_tensor(uv_a_long, image_width).squeeze(1) non_matches_b = SaladDataset.flatten_uv_tensor(uv_b_non_matches_long, image_width).squeeze(1) # 5 is custom data type for distortion image training return 5, image_a_rgb, image_b_rgb, matches_a, matches_b, non_matches_a, non_matches_b
def __getitem__(self, index): """ The method through which the dataset is accessed for training. The index param is not currently used, and instead each dataset[i] is the result of a random sampling over: - random scene - random rgbd frame from that scene - random rgbd frame (different enough pose) from that scene - various randomization in the match generation and non-match generation procedure returns a large amount of variables, separated by commas. 0th return arg: the type of data sampled (this can be used as a flag for different loss functions) 0th rtype: string 1st, 2nd return args: image_a_rgb, image_b_rgb 1st, 2nd rtype: 3-dimensional torch.FloatTensor of shape (image_height, image_width, 3) 3rd, 4th return args: matches_a, matches_b 3rd, 4th rtype: 1-dimensional torch.LongTensor of shape (num_matches) 5th, 6th return args: non_matches_a, non_matches_b 5th, 6th rtype: 1-dimensional torch.LongTensor of shape (num_non_matches) Return values 3,4,5,6 are all in the "single index" format for pixels. That is (u,v) --> n = u + image_width * v """ # stores metadata about this data metadata = dict() # pick a scene scene_name = self.get_random_scene_name() metadata['scene_name'] = scene_name # image a image_a_idx = self.get_random_image_index(scene_name) image_a_rgb, image_a_depth, image_a_mask, image_a_pose = self.get_rgbd_mask_pose(scene_name, image_a_idx) metadata['image_a_idx'] = image_a_idx # image b image_b_idx = self.get_img_idx_with_different_pose(scene_name, image_a_pose, num_attempts=50) # image_b_idx = self.get_next_img_idx(scene_name, image_a_idx) metadata['image_b_idx'] = image_b_idx if image_b_idx is None: logging.info("no frame with sufficiently different pose found, returning") # TODO: return something cleaner than no-data image_a_rgb_tensor = self.rgb_image_to_tensor(image_a_rgb) return self.return_empty_data(image_a_rgb_tensor, image_a_rgb_tensor) image_b_rgb, image_b_depth, image_b_mask, image_b_pose = self.get_rgbd_mask_pose(scene_name, image_b_idx) image_a_depth_numpy = np.asarray(image_a_depth) image_b_depth_numpy = np.asarray(image_b_depth) # find correspondences uv_a, uv_b = correspondence_finder.batch_find_pixel_correspondences(image_a_depth_numpy, image_a_pose, image_b_depth_numpy, image_b_pose, num_attempts=self.num_matching_attempts, img_a_mask=np.asarray(image_a_mask)) if uv_a is None: logging.info("no matches found, returning") image_a_rgb_tensor = self.rgb_image_to_tensor(image_a_rgb) return self.return_empty_data(image_a_rgb_tensor, image_a_rgb_tensor) if self.debug: # downsample so can plot num_matches_to_plot = 10 indexes_to_keep = (torch.rand(num_matches_to_plot)*len(uv_a[0])).floor().type(torch.LongTensor) uv_a = (torch.index_select(uv_a[0], 0, indexes_to_keep), torch.index_select(uv_a[1], 0, indexes_to_keep)) uv_b = (torch.index_select(uv_b[0], 0, indexes_to_keep), torch.index_select(uv_b[1], 0, indexes_to_keep)) # data augmentation if self._domain_randomize: image_a_rgb = correspondence_augmentation.random_domain_randomize_background(image_a_rgb, image_a_mask) image_b_rgb = correspondence_augmentation.random_domain_randomize_background(image_b_rgb, image_b_mask) if not self.debug: [image_a_rgb], uv_a = correspondence_augmentation.random_image_and_indices_mutation([image_a_rgb], uv_a) [image_b_rgb, image_b_mask], uv_b = correspondence_augmentation.random_image_and_indices_mutation([image_b_rgb, image_b_mask], uv_b) else: # also mutate depth just for plotting [image_a_rgb, image_a_depth], uv_a = correspondence_augmentation.random_image_and_indices_mutation([image_a_rgb, image_a_depth], uv_a) [image_b_rgb, image_b_depth, image_b_mask], uv_b = correspondence_augmentation.random_image_and_indices_mutation([image_b_rgb, image_b_depth, image_b_mask], uv_b) image_a_depth_numpy = np.asarray(image_a_depth) image_b_depth_numpy = np.asarray(image_b_depth) # find non_correspondences if index%2: metadata['non_match_type'] = 'masked' logging.debug("masking non-matches") image_b_mask = torch.from_numpy(np.asarray(image_b_mask)).type(torch.FloatTensor) else: metadata['non_match_type'] = 'non_masked' logging.debug("not masking non-matches") image_b_mask = None image_b_shape = image_b_depth_numpy.shape image_width = image_b_shape[1] image_height = image_b_shape[1] uv_b_non_matches = correspondence_finder.create_non_correspondences(uv_b, image_b_shape, num_non_matches_per_match=self.num_non_matches_per_match, img_b_mask=image_b_mask) if self.debug: # only want to bring in plotting code if in debug mode import correspondence_plotter # Just show all images uv_a_long = (torch.t(uv_a[0].repeat(self.num_non_matches_per_match, 1)).contiguous().view(-1,1), torch.t(uv_a[1].repeat(self.num_non_matches_per_match, 1)).contiguous().view(-1,1)) uv_b_non_matches_long = (uv_b_non_matches[0].view(-1,1), uv_b_non_matches[1].view(-1,1) ) # Show correspondences if uv_a is not None: fig, axes = correspondence_plotter.plot_correspondences_direct(image_a_rgb, image_a_depth_numpy, image_b_rgb, image_b_depth_numpy, uv_a, uv_b, show=False) correspondence_plotter.plot_correspondences_direct(image_a_rgb, image_a_depth_numpy, image_b_rgb, image_b_depth_numpy, uv_a_long, uv_b_non_matches_long, use_previous_plot=(fig,axes), circ_color='r') # image_a_rgb, image_b_rgb = self.both_to_tensor([image_a_rgb, image_b_rgb]) # convert PIL.Image to torch.FloatTensor image_a_rgb = self.rgb_image_to_tensor(image_a_rgb) image_b_rgb = self.rgb_image_to_tensor(image_b_rgb) uv_a_long = (torch.t(uv_a[0].repeat(self.num_non_matches_per_match, 1)).contiguous().view(-1,1), torch.t(uv_a[1].repeat(self.num_non_matches_per_match, 1)).contiguous().view(-1,1)) uv_b_non_matches_long = (uv_b_non_matches[0].view(-1,1), uv_b_non_matches[1].view(-1,1) ) # flatten correspondences and non_correspondences matches_a = uv_a[1].long()*image_width+uv_a[0].long() matches_b = uv_b[1].long()*image_width+uv_b[0].long() non_matches_a = uv_a_long[1].long()*image_width+uv_a_long[0].long() non_matches_a = non_matches_a.squeeze(1) non_matches_b = uv_b_non_matches_long[1].long()*image_width+uv_b_non_matches_long[0].long() non_matches_b = non_matches_b.squeeze(1) return "matches", image_a_rgb, image_b_rgb, matches_a, matches_b, non_matches_a, non_matches_b, metadata
def __getitem__(self, index): """ The method through which the dataset is accessed for training. The index param is not currently used, and instead each dataset[i] is the result of a random sampling over: - random scene - random rgbd frame from that scene - random rgbd frame (different enough pose) from that scene - various randomization in the match generation and non-match generation procedure returns a large amount of variables, separated by commas. 0th return arg: the type of data sampled (this can be used as a flag for different loss functions) 0th rtype: string 1st, 2nd return args: image_a_rgb, image_b_rgb 1st, 2nd rtype: 3-dimensional torch.FloatTensor of shape (image_height, image_width, 3) 3rd, 4th return args: matches_a, matches_b 3rd, 4th rtype: 1-dimensional torch.LongTensor of shape (num_matches) 5th, 6th return args: non_matches_a, non_matches_b 5th, 6th rtype: 1-dimensional torch.LongTensor of shape (num_non_matches) Return values 3,4,5,6 are all in the "single index" format for pixels. That is (u,v) --> n = u + image_width * v """ # stores metadata about this data metadata = dict() # pick a scene scene_name = self.get_random_scene_name() metadata['scene_name'] = scene_name # image a image_a_idx = self.get_random_image_index(scene_name) image_a_rgb, image_a_depth, image_a_mask, image_a_pose = self.get_rgbd_mask_pose(scene_name, image_a_idx) metadata['image_a_idx'] = image_a_idx # image b image_b_idx = self.get_img_idx_with_different_pose(scene_name, image_a_pose, num_attempts=50) metadata['image_b_idx'] = image_b_idx if image_b_idx is None: logging.info("no frame with sufficiently different pose found, returning") # TODO: return something cleaner than no-data image_a_rgb_tensor = self.rgb_image_to_tensor(image_a_rgb) return self.return_empty_data(image_a_rgb_tensor, image_a_rgb_tensor) image_b_rgb, image_b_depth, image_b_mask, image_b_pose = self.get_rgbd_mask_pose(scene_name, image_b_idx) image_a_depth_numpy = np.asarray(image_a_depth) image_b_depth_numpy = np.asarray(image_b_depth) # find correspondences uv_a, uv_b = correspondence_finder.batch_find_pixel_correspondences(image_a_depth_numpy, image_a_pose, image_b_depth_numpy, image_b_pose, num_attempts=self.num_matching_attempts, img_a_mask=np.asarray(image_a_mask)) if uv_a is None: logging.info("no matches found, returning") image_a_rgb_tensor = self.rgb_image_to_tensor(image_a_rgb) return self.return_empty_data(image_a_rgb_tensor, image_a_rgb_tensor) if self.debug: # downsample so can plot num_matches_to_plot = 10 indexes_to_keep = (torch.rand(num_matches_to_plot)*len(uv_a[0])).floor().type(torch.LongTensor) uv_a = (torch.index_select(uv_a[0], 0, indexes_to_keep), torch.index_select(uv_a[1], 0, indexes_to_keep)) uv_b = (torch.index_select(uv_b[0], 0, indexes_to_keep), torch.index_select(uv_b[1], 0, indexes_to_keep)) # data augmentation if self._domain_randomize: image_a_rgb = correspondence_augmentation.random_domain_randomize_background(image_a_rgb, image_a_mask) image_b_rgb = correspondence_augmentation.random_domain_randomize_background(image_b_rgb, image_b_mask) if not self.debug: [image_a_rgb], uv_a = correspondence_augmentation.random_image_and_indices_mutation([image_a_rgb], uv_a) [image_b_rgb, image_b_mask], uv_b = correspondence_augmentation.random_image_and_indices_mutation([image_b_rgb, image_b_mask], uv_b) else: # also mutate depth just for plotting [image_a_rgb, image_a_depth], uv_a = correspondence_augmentation.random_image_and_indices_mutation([image_a_rgb, image_a_depth], uv_a) [image_b_rgb, image_b_depth, image_b_mask], uv_b = correspondence_augmentation.random_image_and_indices_mutation([image_b_rgb, image_b_depth, image_b_mask], uv_b) image_a_depth_numpy = np.asarray(image_a_depth) image_b_depth_numpy = np.asarray(image_b_depth) # find non_correspondences if index%2: metadata['non_match_type'] = 'masked' logging.debug("masking non-matches") image_b_mask = torch.from_numpy(np.asarray(image_b_mask)).type(torch.FloatTensor) else: metadata['non_match_type'] = 'non_masked' logging.debug("not masking non-matches") image_b_mask = None image_b_shape = image_b_depth_numpy.shape image_width = image_b_shape[1] image_height = image_b_shape[1] uv_b_non_matches = correspondence_finder.create_non_correspondences(uv_b, image_b_shape, num_non_matches_per_match=self.num_non_matches_per_match, img_b_mask=image_b_mask) if self.debug: # only want to bring in plotting code if in debug mode import correspondence_plotter # Just show all images uv_a_long = (torch.t(uv_a[0].repeat(self.num_non_matches_per_match, 1)).contiguous().view(-1,1), torch.t(uv_a[1].repeat(self.num_non_matches_per_match, 1)).contiguous().view(-1,1)) uv_b_non_matches_long = (uv_b_non_matches[0].view(-1,1), uv_b_non_matches[1].view(-1,1) ) # Show correspondences if uv_a is not None: fig, axes = correspondence_plotter.plot_correspondences_direct(image_a_rgb, image_a_depth_numpy, image_b_rgb, image_b_depth_numpy, uv_a, uv_b, show=False) correspondence_plotter.plot_correspondences_direct(image_a_rgb, image_a_depth_numpy, image_b_rgb, image_b_depth_numpy, uv_a_long, uv_b_non_matches_long, use_previous_plot=(fig,axes), circ_color='r') # image_a_rgb, image_b_rgb = self.both_to_tensor([image_a_rgb, image_b_rgb]) # convert PIL.Image to torch.FloatTensor image_a_rgb = self.rgb_image_to_tensor(image_a_rgb) image_b_rgb = self.rgb_image_to_tensor(image_b_rgb) uv_a_long = (torch.t(uv_a[0].repeat(self.num_non_matches_per_match, 1)).contiguous().view(-1,1), torch.t(uv_a[1].repeat(self.num_non_matches_per_match, 1)).contiguous().view(-1,1)) uv_b_non_matches_long = (uv_b_non_matches[0].view(-1,1), uv_b_non_matches[1].view(-1,1) ) # flatten correspondences and non_correspondences matches_a = uv_a[1].long()*image_width+uv_a[0].long() matches_b = uv_b[1].long()*image_width+uv_b[0].long() non_matches_a = uv_a_long[1].long()*image_width+uv_a_long[0].long() non_matches_a = non_matches_a.squeeze(1) non_matches_b = uv_b_non_matches_long[1].long()*image_width+uv_b_non_matches_long[0].long() non_matches_b = non_matches_b.squeeze(1) return "matches", image_a_rgb, image_b_rgb, matches_a, matches_b, non_matches_a, non_matches_b, metadata