class HeatmapVisualization(object): """ Launches a live interactive heatmap visualization. Edit config/dense_correspondence/heatmap_vis/heatmap.yaml to specify which networks to visualize. Specifically add the network you want to visualize to the "networks" list. Make sure that this network appears in the file pointed to by EVAL_CONFIG Usage: Launch this file with python after sourcing the environment with `use_pytorch_dense_correspondence` Then `python live_heatmap_visualization.py`. """ def __init__(self, config): self._config = config self._dce = DenseCorrespondenceEvaluation(EVAL_CONFIG) self._load_networks() self._reticle_color = COLOR_GREEN self.load_specific_dataset( ) # uncomment if you want to load a specific dataset def _load_networks(self): # we will use the dataset for the first network in the series self._dcn_dict = dict() self._dataset = None self._network_reticle_color = dict() for idx, network_name in enumerate(self._config["networks"]): dcn = self._dce.load_network_from_config(network_name) dcn.eval() self._dcn_dict[network_name] = dcn if len(self._config["networks"]) == 1: self._network_reticle_color[network_name] = COLOR_RED else: self._network_reticle_color[network_name] = label_colors[idx] if self._dataset is None: self._dataset = dcn.load_training_dataset() def load_specific_dataset(self): dataset_config_filename = os.path.join( utils.getDenseCorrespondenceSourceDir(), 'config', 'dense_correspondence', 'dataset', 'composite', 'rope_nonrigid_412vert_only.yaml') dataset_config = utils.getDictFromYamlFilename(dataset_config_filename) self._dataset = SpartanDataset(config=dataset_config) #self._dataset.get_knots_info('rope_nonrigid_412vert') def get_random_image_pair(self): object_id = self._dataset.get_random_object_id() scene_name_a = self._dataset.get_random_single_object_scene_name( object_id) scene_name_b = self._dataset.get_random_single_object_scene_name( object_id) image_a_idx = self._dataset.get_random_image_index(scene_name_a) image_b_idx = self._dataset.get_random_image_index(scene_name_b) return scene_name_a, scene_name_b, image_a_idx, image_b_idx def _get_new_images(self): """ Gets a new pair of images :return: :rtype: """ if random.random() <= 1.0: self._dataset.set_train_mode() else: self._dataset.set_test_mode() if self._config["same_object"]: print "getting random image pair" scene_name_1, scene_name_2, image_1_idx, image_2_idx = self.get_random_image_pair( ) else: raise ValueError( "At least one of the image types must be set tot True") print "got pair" self.img1_pil = self._dataset.get_rgb_image_from_scene_name_and_idx( scene_name_1, image_1_idx) print "got img1" self.img1_knots = self._dataset._knots_info[scene_name_1][image_1_idx] print "got img1 knots" self.img2_pil = self._dataset.get_rgb_image_from_scene_name_and_idx( scene_name_2, image_2_idx) self.img2_knots = self._dataset._knots_info[scene_name_2][image_2_idx] def _get_task_images(self): self.img1_pil = self._dataset.get_rgb_image('./images/000025_rgb.png') img1_mask = self._dataset.get_mask_image( './image_masks/000025_mask.png') self.img2_pil = self._dataset.get_rgb_image('./images/000018_rgb.png') pixs = correspondence_finder.random_sample_from_masked_image_torch( np.asarray(img1_mask), 25) self.img1_knots = list(zip(pixs[0], pixs[1])) #self.img1_knots = utils.getDictFromJSONFilename('./images_start/knots_info.json')["0"][0] #self.img2_knots = utils.getDictFromJSONFilename('./images_goal/knots_info.json')["0"][0] def _compute_descriptors(self, knot_idx): """ Computes the descriptors for image 1 and image 2 for each network :return: :rtype: """ self.img1 = pil_image_to_cv2(self.img1_pil) self.img2 = pil_image_to_cv2(self.img2_pil) self.rgb_1_tensor = self._dataset.rgb_image_to_tensor(self.img1_pil) self.rgb_2_tensor = self._dataset.rgb_image_to_tensor(self.img2_pil) self.img1_gray = cv2.cvtColor(self.img1, cv2.COLOR_RGB2GRAY) / 255.0 self.img2_gray = cv2.cvtColor(self.img2, cv2.COLOR_RGB2GRAY) / 255.0 self._res_a = dict() self._res_b = dict() for network_name, dcn in self._dcn_dict.iteritems(): self._res_a[network_name] = dcn.forward_single_image_tensor( self.rgb_1_tensor).data.cpu().numpy() self._res_b[network_name] = dcn.forward_single_image_tensor( self.rgb_2_tensor).data.cpu().numpy() print(self.img1_knots) u, v = self.img1_knots[knot_idx] source, blended, target, p = self.find_best_match( None, u, v, None, None) return (source, blended, target, p) def scale_norm_diffs_to_make_heatmap(self, norm_diffs, threshold): """ Scales the norm diffs to make a heatmap. This will be scaled between 0 and 1. 0 corresponds to a match, 1 to non-match :param norm_diffs: The norm diffs :type norm_diffs: numpy.array [H,W] :return: :rtype: """ heatmap = np.copy(norm_diffs) greater_than_threshold = np.where(norm_diffs > threshold) heatmap = heatmap / threshold * self._config[ "heatmap_vis_upper_bound"] # linearly scale [0, threshold] to [0, 0.5] heatmap[ greater_than_threshold] = 1 # greater than threshold is set to 1 heatmap = heatmap.astype(self.img1_gray.dtype) return heatmap def find_best_match(self, event, u, v, flags, param): """ For each network, find the best match in the target image to point highlighted with reticle in the source image. Displays the result :return: :rtype: """ img_1_with_reticle = np.copy(self.img1) draw_reticle(img_1_with_reticle, u, v, self._reticle_color) source = img_1_with_reticle alpha = self._config["blend_weight_original_image"] beta = 1 - alpha img_2_with_reticle = np.copy(self.img2) self._res_uv = dict() for network_name in self._dcn_dict: res_a = self._res_a[network_name] res_b = self._res_b[network_name] best_match_uv, best_match_diff, norm_diffs = \ DenseCorrespondenceNetwork.find_best_match((u, v), res_a, res_b) print "network_name:", network_name self._res_uv[network_name] = dict() self._res_uv[network_name]['source'] = res_a[v, u, :].tolist() self._res_uv[network_name]['target'] = res_b[v, u, :].tolist() print "%s best match diff: %.3f" % (network_name, best_match_diff) threshold = self._config["norm_diff_threshold"] if network_name in self._config["norm_diff_threshold_dict"]: threshold = self._config["norm_diff_threshold_dict"][ network_name] heatmap = self.scale_norm_diffs_to_make_heatmap( norm_diffs, threshold) reticle_color = self._network_reticle_color[network_name] draw_reticle(heatmap, best_match_uv[0], best_match_uv[1], reticle_color) draw_reticle(img_2_with_reticle, best_match_uv[0], best_match_uv[1], reticle_color) blended = cv2.addWeighted(self.img2_gray, alpha, heatmap, beta, 0) target = img_2_with_reticle return (source, blended, target, [best_match_uv[0], best_match_uv[1]]) def run(self): self._get_task_images() pixels = [] # for i in range(len(utils.getDictFromJSONFilename('./images_start/knots_info.json')["0"][0])): for i in range(25): print "computing descriptors" source, blended, target, p = self._compute_descriptors(i) pixels.append(p) print "computed descriptors" vis = np.concatenate((source, target), axis=1) print "concatenated, writing image" cv2.imwrite("/home/priya/code/data_volume/annotated/%06d.png" % i, vis) np.savetxt('pixels_pred.txt', pixels)
class HeatmapVisualization(object): """ Launches a live interactive heatmap visualization. Edit config/dense_correspondence/heatmap_vis/heatmap.yaml to specify which networks to visualize. Specifically add the network you want to visualize to the "networks" list. Make sure that this network appears in the file pointed to by EVAL_CONFIG Usage: Launch this file with python after sourcing the environment with `use_pytorch_dense_correspondence` Then `python live_heatmap_visualization.py`. Keypresses: n: new set of images s: swap images p: pause/un-pause """ def __init__(self, config, eval_config): self._config = config self._dce = DenseCorrespondenceEvaluation(eval_config) self._load_networks() self._reticle_color = COLOR_GREEN self._paused = False if LOAD_SPECIFIC_DATASET: self.load_specific_dataset( ) # uncomment if you want to load a specific dataset def _load_networks(self): # we will use the dataset for the first network in the series self._dcn_dict = dict() self._dataset = None self._network_reticle_color = dict() for idx, network_name in enumerate(self._config["networks"]): dcn = self._dce.load_network_from_config(network_name) dcn.eval() self._dcn_dict[network_name] = dcn # self._network_reticle_color[network_name] = label_colors[idx] if len(self._config["networks"]) == 1: self._network_reticle_color[network_name] = COLOR_RED else: self._network_reticle_color[network_name] = label_colors[idx] if self._dataset is None: self._dataset = dcn.load_training_dataset() def load_specific_dataset(self): dataset_config_filename = os.path.join( utils.getDenseCorrespondenceSourceDir(), 'config', 'dense_correspondence', 'dataset', 'composite', 'hats_3_demo_composite.yaml') # dataset_config_filename = os.path.join(utils.getDenseCorrespondenceSourceDir(), 'config', # 'dense_correspondence', # 'dataset', 'composite', '4_shoes_all.yaml') dataset_config = utils.getDictFromYamlFilename(dataset_config_filename) self._dataset = SpartanDataset(config=dataset_config) def get_random_image_pair(self): """ Gets a pair of random images for different scenes of the same object """ object_id = self._dataset.get_random_object_id() # scene_name_a = "2018-04-10-16-02-59" # scene_name_b = scene_name_a scene_name_a = self._dataset.get_random_single_object_scene_name( object_id) scene_name_b = self._dataset.get_different_scene_for_object( object_id, scene_name_a) if self._config["randomize_images"]: image_a_idx = self._dataset.get_random_image_index(scene_name_a) image_b_idx = self._dataset.get_random_image_index(scene_name_b) else: image_a_idx = 0 image_b_idx = 0 return scene_name_a, scene_name_b, image_a_idx, image_b_idx def get_random_image_pair_across_object(self): """ Gets cross object image pairs :param randomize: :type randomize: :return: :rtype: """ object_id_a, object_id_b = self._dataset.get_two_different_object_ids() # object_id_a = "shoe_red_nike.yaml" # object_id_b = "shoe_gray_nike" # object_id_b = "shoe_green_nike" scene_name_a = self._dataset.get_random_single_object_scene_name( object_id_a) scene_name_b = self._dataset.get_random_single_object_scene_name( object_id_b) if self._config["randomize_images"]: image_a_idx = self._dataset.get_random_image_index(scene_name_a) image_b_idx = self._dataset.get_random_image_index(scene_name_b) else: image_a_idx = 0 image_b_idx = 0 return scene_name_a, scene_name_b, image_a_idx, image_b_idx def get_random_image_pair_multi_object_scenes(self): """ Gets cross object image pairs :param randomize: :type randomize: :return: :rtype: """ scene_name_a = self._dataset.get_random_multi_object_scene_name() scene_name_b = self._dataset.get_random_multi_object_scene_name() if self._config["randomize_images"]: image_a_idx = self._dataset.get_random_image_index(scene_name_a) image_b_idx = self._dataset.get_random_image_index(scene_name_b) else: image_a_idx = 0 image_b_idx = 0 return scene_name_a, scene_name_b, image_a_idx, image_b_idx def _get_new_images(self): """ Gets a new pair of images :return: :rtype: """ if random.random() < 0.5: self._dataset.set_train_mode() else: self._dataset.set_test_mode() if self._config["same_object"]: scene_name_1, scene_name_2, image_1_idx, image_2_idx = self.get_random_image_pair( ) elif self._config["different_objects"]: scene_name_1, scene_name_2, image_1_idx, image_2_idx = self.get_random_image_pair_across_object( ) elif self._config["multiple_object"]: scene_name_1, scene_name_2, image_1_idx, image_2_idx = self.get_random_image_pair_multi_object_scenes( ) else: raise ValueError( "At least one of the image types must be set tot True") # caterpillar # scene_name_1 = "2018-04-16-14-42-26" # scene_name_2 = "2018-04-16-14-25-19" # hats # scene_name_1 = "2018-05-15-22-01-44" # scene_name_2 = "2018-05-15-22-04-17" self.img1_pil = self._dataset.get_rgb_image_from_scene_name_and_idx( scene_name_1, image_1_idx) self.img2_pil = self._dataset.get_rgb_image_from_scene_name_and_idx( scene_name_2, image_2_idx) self._scene_name_1 = scene_name_1 self._scene_name_2 = scene_name_2 self._image_1_idx = image_1_idx self._image_2_idx = image_2_idx self._compute_descriptors() # self.rgb_1_tensor = self._dataset.rgb_image_to_tensor(img1_pil) # self.rgb_2_tensor = self._dataset.rgb_image_to_tensor(img2_pil) def _compute_descriptors(self): """ Computes the descriptors for image 1 and image 2 for each network :return: :rtype: """ self.img1 = pil_image_to_cv2(self.img1_pil) self.img2 = pil_image_to_cv2(self.img2_pil) self.rgb_1_tensor = self._dataset.rgb_image_to_tensor(self.img1_pil) self.rgb_2_tensor = self._dataset.rgb_image_to_tensor(self.img2_pil) self.img1_gray = cv2.cvtColor(self.img1, cv2.COLOR_RGB2GRAY) / 255.0 self.img2_gray = cv2.cvtColor(self.img2, cv2.COLOR_RGB2GRAY) / 255.0 cv2.imshow('source', self.img1) cv2.imshow('target', self.img2) self._res_a = dict() self._res_b = dict() for network_name, dcn in self._dcn_dict.items(): self._res_a[network_name] = dcn.forward_single_image_tensor( self.rgb_1_tensor).data.cpu().numpy() self._res_b[network_name] = dcn.forward_single_image_tensor( self.rgb_2_tensor).data.cpu().numpy() self.find_best_match(None, 0, 0, None, None) def scale_norm_diffs_to_make_heatmap(self, norm_diffs, threshold): """ TODO (@manuelli) scale with Gaussian kernel instead of linear Scales the norm diffs to make a heatmap. This will be scaled between 0 and 1. 0 corresponds to a match, 1 to non-match :param norm_diffs: The norm diffs :type norm_diffs: numpy.array [H,W] :return: :rtype: """ heatmap = np.copy(norm_diffs) greater_than_threshold = np.where(norm_diffs > threshold) heatmap = heatmap / threshold * self._config[ "heatmap_vis_upper_bound"] # linearly scale [0, threshold] to [0, 0.5] heatmap[ greater_than_threshold] = 1 # greater than threshold is set to 1 heatmap = heatmap.astype(self.img1_gray.dtype) return heatmap def find_best_match(self, event, u, v, flags, param): """ For each network, find the best match in the target image to point highlighted with reticle in the source image. Displays the result :return: :rtype: """ if self._paused: return img_1_with_reticle = np.copy(self.img1) draw_reticle(img_1_with_reticle, u, v, self._reticle_color) cv2.imshow("source", img_1_with_reticle) alpha = self._config["blend_weight_original_image"] beta = 1 - alpha img_2_with_reticle = np.copy(self.img2) print("\n\n") self._res_uv = dict() # self._res_a_uv = dict() # self._res_b_uv = dict() for network_name in self._dcn_dict: res_a = self._res_a[network_name] res_b = self._res_b[network_name] best_match_uv, best_match_diff, norm_diffs = \ DenseCorrespondenceNetwork.find_best_match((u, v), res_a, res_b) print("\n\n") print("network_name:", network_name) print("scene_name_1", self._scene_name_1) print("image_1_idx", self._image_1_idx) print("scene_name_2", self._scene_name_2) print("image_2_idx", self._image_2_idx) d = dict() d['scene_name'] = self._scene_name_1 d['image_idx'] = self._image_1_idx d['descriptor'] = res_a[v, u, :].tolist() d['u'] = u d['v'] = v print("\n-------keypoint info\n", d) print("\n--------\n") self._res_uv[network_name] = dict() self._res_uv[network_name]['source'] = res_a[v, u, :].tolist() self._res_uv[network_name]['target'] = res_b[v, u, :].tolist() print("res_a[v, u, :]:", res_a[v, u, :]) print("res_b[v, u, :]:", res_b[best_match_uv[1], best_match_uv[0], :]) print("%s best match diff: %.3f" % (network_name, best_match_diff)) print("res_a", self._res_uv[network_name]['source']) print("res_b", self._res_uv[network_name]['target']) threshold = self._config["norm_diff_threshold"] if network_name in self._config["norm_diff_threshold_dict"]: threshold = self._config["norm_diff_threshold_dict"][ network_name] heatmap_color = vis_utils.compute_gaussian_kernel_heatmap_from_norm_diffs( norm_diffs, self._config['kernel_variance']) reticle_color = self._network_reticle_color[network_name] draw_reticle(heatmap_color, best_match_uv[0], best_match_uv[1], reticle_color) draw_reticle(img_2_with_reticle, best_match_uv[0], best_match_uv[1], reticle_color) blended = cv2.addWeighted(self.img2, alpha, heatmap_color, beta, 0) cv2.imshow(network_name, blended) cv2.imshow("target", img_2_with_reticle) if event == cv2.EVENT_LBUTTONDOWN: utils.saveToYaml(self._res_uv, 'clicked_point.yaml') def run(self): self._get_new_images() cv2.namedWindow('target') cv2.setMouseCallback('source', self.find_best_match) self._get_new_images() while True: k = cv2.waitKey(20) & 0xFF if k == 27: break elif k == ord('n'): self._get_new_images() elif k == ord('s'): img1_pil = self.img1_pil img2_pil = self.img2_pil self.img1_pil = img2_pil self.img2_pil = img1_pil self._compute_descriptors() elif k == ord('p'): if self._paused: print("un pausing") self._paused = False else: print("pausing") self._paused = True
class HeatmapVisualization(object): def __init__(self, config): self._config = config self._dce = DenseCorrespondenceEvaluation(EVAL_CONFIG) self._load_networks() self._reticle_color = COLOR_GREEN # self.load_specific_dataset() # uncomment if you want to load a specific dataset def _load_networks(self): # we will use the dataset for the first network in the series self._dcn_dict = dict() self._dataset = None self._network_reticle_color = dict() for idx, network_name in enumerate(self._config["networks"]): dcn = self._dce.load_network_from_config(network_name) dcn.eval() self._dcn_dict[network_name] = dcn # self._network_reticle_color[network_name] = label_colors[idx] if len(self._config["networks"]) == 1: self._network_reticle_color[network_name] = COLOR_RED else: self._network_reticle_color[network_name] = label_colors[idx] if self._dataset is None: self._dataset = dcn.load_training_dataset() def load_specific_dataset(self): dataset_config_filename = os.path.join( utils.getDenseCorrespondenceSourceDir(), 'config', 'dense_correspondence', 'dataset', 'composite', 'hats_3_demo_composite.yaml') dataset_config_filename = os.path.join( utils.getDenseCorrespondenceSourceDir(), 'config', 'dense_correspondence', 'dataset', 'composite', '4_shoes_all.yaml') dataset_config = utils.getDictFromYamlFilename(dataset_config_filename) self._dataset = SpartanDataset(config=dataset_config) def get_random_image_pair(self): object_id = self._dataset.get_random_object_id() scene_name_a = self._dataset.get_random_single_object_scene_name( object_id) scene_name_b = self._dataset.get_different_scene_for_object( object_id, scene_name_a) if self._config["randomize_images"]: image_a_idx = self._dataset.get_random_image_index(scene_name_a) image_b_idx = self._dataset.get_random_image_index(scene_name_b) else: image_a_idx = 0 image_b_idx = 0 # image_b_idx = self._dataset.get_random_image_index(scene_name_b) return scene_name_a, scene_name_b, image_a_idx, image_b_idx def get_random_image_pair_across_object(self): """ Gets cross object image pairs :param randomize: :type randomize: :return: :rtype: """ object_id_a, object_id_b = self._dataset.get_two_different_object_ids() # object_id_a = "shoe_red_nike.yaml" # object_id_b = "shoe_gray_nike" # object_id_b = "shoe_green_nike" scene_name_a = self._dataset.get_random_single_object_scene_name( object_id_a) scene_name_b = self._dataset.get_random_single_object_scene_name( object_id_b) if self._config["randomize_images"]: image_a_idx = self._dataset.get_random_image_index(scene_name_a) image_b_idx = self._dataset.get_random_image_index(scene_name_b) else: image_a_idx = 0 image_b_idx = 0 return scene_name_a, scene_name_b, image_a_idx, image_b_idx def get_random_image_pair_multi_object_scenes(self): """ Gets cross object image pairs :param randomize: :type randomize: :return: :rtype: """ scene_name_a = self._dataset.get_random_multi_object_scene_name() scene_name_b = self._dataset.get_random_multi_object_scene_name() if self._config["randomize_images"]: image_a_idx = self._dataset.get_random_image_index(scene_name_a) image_b_idx = self._dataset.get_random_image_index(scene_name_b) else: image_a_idx = 0 image_b_idx = 0 return scene_name_a, scene_name_b, image_a_idx, image_b_idx def _get_new_images(self): """ Gets a new pair of images :return: :rtype: """ if random.random() < 0.5: self._dataset.set_train_mode() else: self._dataset.set_test_mode() if self._config["same_object"]: scene_name_1, scene_name_2, image_1_idx, image_2_idx = self.get_random_image_pair( ) elif self._config["different_objects"]: scene_name_1, scene_name_2, image_1_idx, image_2_idx = self.get_random_image_pair_across_object( ) elif self._config["multiple_object"]: scene_name_1, scene_name_2, image_1_idx, image_2_idx = self.get_random_image_pair_multi_object_scenes( ) else: raise ValueError( "At least one of the image types must be set tot True") self.img1_pil = self._dataset.get_rgb_image_from_scene_name_and_idx( scene_name_1, image_1_idx) self.img2_pil = self._dataset.get_rgb_image_from_scene_name_and_idx( scene_name_2, image_2_idx) self._compute_descriptors() # self.rgb_1_tensor = self._dataset.rgb_image_to_tensor(img1_pil) # self.rgb_2_tensor = self._dataset.rgb_image_to_tensor(img2_pil) def _compute_descriptors(self): """ Computes the descriptors for image 1 and image 2 for each network :return: :rtype: """ self.img1 = pil_image_to_cv2(self.img1_pil) self.img2 = pil_image_to_cv2(self.img2_pil) self.rgb_1_tensor = self._dataset.rgb_image_to_tensor(self.img1_pil) self.rgb_2_tensor = self._dataset.rgb_image_to_tensor(self.img2_pil) self.img1_gray = cv2.cvtColor(self.img1, cv2.COLOR_RGB2GRAY) / 255.0 self.img2_gray = cv2.cvtColor(self.img2, cv2.COLOR_RGB2GRAY) / 255.0 cv2.imshow('source', self.img1) cv2.imshow('target', self.img2) self._res_a = dict() self._res_b = dict() for network_name, dcn in self._dcn_dict.iteritems(): self._res_a[network_name] = dcn.forward_single_image_tensor( self.rgb_1_tensor).data.cpu().numpy() self._res_b[network_name] = dcn.forward_single_image_tensor( self.rgb_2_tensor).data.cpu().numpy() self.find_best_match(None, 0, 0, None, None) def scale_norm_diffs_to_make_heatmap(self, norm_diffs, threshold): """ Scales the norm diffs to make a heatmap. This will be scaled between 0 and 1. 0 corresponds to a match, 1 to non-match :param norm_diffs: The norm diffs :type norm_diffs: numpy.array [H,W] :return: :rtype: """ heatmap = np.copy(norm_diffs) greater_than_threshold = np.where(norm_diffs > threshold) heatmap = heatmap / threshold * self._config[ "heatmap_vis_upper_bound"] # linearly scale [0, threshold] to [0, 0.5] heatmap[ greater_than_threshold] = 1 # greater than threshold is set to 1 heatmap = heatmap.astype(self.img1_gray.dtype) return heatmap def find_best_match(self, event, u, v, flags, param): """ For each network, find the best match in the target image to point highlighted with reticle in the source image. Displays the result :return: :rtype: """ img_1_with_reticle = np.copy(self.img1) draw_reticle(img_1_with_reticle, u, v, self._reticle_color) cv2.imshow("source", img_1_with_reticle) alpha = self._config["blend_weight_original_image"] beta = 1 - alpha img_2_with_reticle = np.copy(self.img2) print "\n\n" self._res_uv = dict() # self._res_a_uv = dict() # self._res_b_uv = dict() for network_name in self._dcn_dict: res_a = self._res_a[network_name] res_b = self._res_b[network_name] best_match_uv, best_match_diff, norm_diffs = \ DenseCorrespondenceNetwork.find_best_match((u, v), res_a, res_b) print "\n\n" print "network_name:", network_name self._res_uv[network_name] = dict() self._res_uv[network_name]['source'] = res_a[v, u, :].tolist() self._res_uv[network_name]['target'] = res_b[v, u, :].tolist() # print "res_a[v, u, :]:", res_a[v, u, :] # print "res_b[v, u, :]:", res_b[v, u, :] print "%s best match diff: %.3f" % (network_name, best_match_diff) threshold = self._config["norm_diff_threshold"] if network_name in self._config["norm_diff_threshold_dict"]: threshold = self._config["norm_diff_threshold_dict"][ network_name] heatmap = self.scale_norm_diffs_to_make_heatmap( norm_diffs, threshold) reticle_color = self._network_reticle_color[network_name] draw_reticle(heatmap, best_match_uv[0], best_match_uv[1], reticle_color) draw_reticle(img_2_with_reticle, best_match_uv[0], best_match_uv[1], reticle_color) blended = cv2.addWeighted(self.img2_gray, alpha, heatmap, beta, 0) cv2.imshow(network_name, blended) cv2.imshow("target", img_2_with_reticle) if event == cv2.EVENT_LBUTTONDOWN: utils.saveToYaml(self._res_uv, 'clicked_point.yaml') def run(self): self._get_new_images() cv2.namedWindow('target') cv2.setMouseCallback('source', self.find_best_match) self._get_new_images() while True: k = cv2.waitKey(20) & 0xFF if k == 27: break elif k == ord('n'): print "HEY" self._get_new_images() elif k == ord('s'): print "HEY" img1_pil = self.img1_pil img2_pil = self.img2_pil self.img1_pil = img2_pil self.img2_pil = img1_pil self._compute_descriptors()
class HeatmapVisualization(object): def __init__(self, config): self._config = config self._dce = DenseCorrespondenceEvaluation(EVAL_CONFIG) self._load_networks() self._reticle_color = COLOR_GREEN # self.load_specific_dataset() # uncomment if you want to load a specific dataset def _load_networks(self): # we will use the dataset for the first network in the series self._dcn_dict = dict() self._dataset = None self._network_reticle_color = dict() for idx, network_name in enumerate(self._config["networks"]): dcn = self._dce.load_network_from_config(network_name) dcn.eval() self._dcn_dict[network_name] = dcn # self._network_reticle_color[network_name] = label_colors[idx] if len(self._config["networks"]) == 1: self._network_reticle_color[network_name] = COLOR_RED else: self._network_reticle_color[network_name] = label_colors[idx] if self._dataset is None: self._dataset = dcn.load_training_dataset() def load_specific_dataset(self): dataset_config_filename = os.path.join(utils.getDenseCorrespondenceSourceDir(), 'config', 'dense_correspondence', 'dataset', 'composite', 'hats_3_demo_composite.yaml') dataset_config_filename = os.path.join(utils.getDenseCorrespondenceSourceDir(), 'config', 'dense_correspondence', 'dataset', 'composite', '4_shoes_all.yaml') dataset_config = utils.getDictFromYamlFilename(dataset_config_filename) self._dataset = SpartanDataset(config=dataset_config) def get_random_image_pair(self): object_id = self._dataset.get_random_object_id() scene_name_a = self._dataset.get_random_single_object_scene_name(object_id) scene_name_b = self._dataset.get_different_scene_for_object(object_id, scene_name_a) if self._config["randomize_images"]: image_a_idx = self._dataset.get_random_image_index(scene_name_a) image_b_idx = self._dataset.get_random_image_index(scene_name_b) else: image_a_idx = 0 image_b_idx = 0 # image_b_idx = self._dataset.get_random_image_index(scene_name_b) return scene_name_a, scene_name_b, image_a_idx, image_b_idx def get_random_image_pair_across_object(self): """ Gets cross object image pairs :param randomize: :type randomize: :return: :rtype: """ object_id_a, object_id_b = self._dataset.get_two_different_object_ids() # object_id_a = "shoe_red_nike.yaml" # object_id_b = "shoe_gray_nike" # object_id_b = "shoe_green_nike" scene_name_a = self._dataset.get_random_single_object_scene_name(object_id_a) scene_name_b = self._dataset.get_random_single_object_scene_name(object_id_b) if self._config["randomize_images"]: image_a_idx = self._dataset.get_random_image_index(scene_name_a) image_b_idx = self._dataset.get_random_image_index(scene_name_b) else: image_a_idx = 0 image_b_idx = 0 return scene_name_a, scene_name_b, image_a_idx, image_b_idx def get_random_image_pair_multi_object_scenes(self): """ Gets cross object image pairs :param randomize: :type randomize: :return: :rtype: """ scene_name_a = self._dataset.get_random_multi_object_scene_name() scene_name_b = self._dataset.get_random_multi_object_scene_name() if self._config["randomize_images"]: image_a_idx = self._dataset.get_random_image_index(scene_name_a) image_b_idx = self._dataset.get_random_image_index(scene_name_b) else: image_a_idx = 0 image_b_idx = 0 return scene_name_a, scene_name_b, image_a_idx, image_b_idx def _get_new_images(self): """ Gets a new pair of images :return: :rtype: """ if random.random() < 0.5: self._dataset.set_train_mode() else: self._dataset.set_test_mode() if self._config["same_object"]: scene_name_1, scene_name_2, image_1_idx, image_2_idx = self.get_random_image_pair() elif self._config["different_objects"]: scene_name_1, scene_name_2, image_1_idx, image_2_idx = self.get_random_image_pair_across_object() elif self._config["multiple_object"]: scene_name_1, scene_name_2, image_1_idx, image_2_idx = self.get_random_image_pair_multi_object_scenes() else: raise ValueError("At least one of the image types must be set tot True") self.img1_pil = self._dataset.get_rgb_image_from_scene_name_and_idx(scene_name_1, image_1_idx) self.img2_pil = self._dataset.get_rgb_image_from_scene_name_and_idx(scene_name_2, image_2_idx) self._compute_descriptors() # self.rgb_1_tensor = self._dataset.rgb_image_to_tensor(img1_pil) # self.rgb_2_tensor = self._dataset.rgb_image_to_tensor(img2_pil) def _compute_descriptors(self): """ Computes the descriptors for image 1 and image 2 for each network :return: :rtype: """ self.img1 = pil_image_to_cv2(self.img1_pil) self.img2 = pil_image_to_cv2(self.img2_pil) self.rgb_1_tensor = self._dataset.rgb_image_to_tensor(self.img1_pil) self.rgb_2_tensor = self._dataset.rgb_image_to_tensor(self.img2_pil) self.img1_gray = cv2.cvtColor(self.img1, cv2.COLOR_RGB2GRAY) / 255.0 self.img2_gray = cv2.cvtColor(self.img2, cv2.COLOR_RGB2GRAY) / 255.0 cv2.imshow('source', self.img1) cv2.imshow('target', self.img2) self._res_a = dict() self._res_b = dict() for network_name, dcn in self._dcn_dict.iteritems(): self._res_a[network_name] = dcn.forward_single_image_tensor(self.rgb_1_tensor).data.cpu().numpy() self._res_b[network_name] = dcn.forward_single_image_tensor(self.rgb_2_tensor).data.cpu().numpy() self.find_best_match(None, 0, 0, None, None) def scale_norm_diffs_to_make_heatmap(self, norm_diffs, threshold): """ Scales the norm diffs to make a heatmap. This will be scaled between 0 and 1. 0 corresponds to a match, 1 to non-match :param norm_diffs: The norm diffs :type norm_diffs: numpy.array [H,W] :return: :rtype: """ heatmap = np.copy(norm_diffs) greater_than_threshold = np.where(norm_diffs > threshold) heatmap = heatmap / threshold * self._config["heatmap_vis_upper_bound"] # linearly scale [0, threshold] to [0, 0.5] heatmap[greater_than_threshold] = 1 # greater than threshold is set to 1 heatmap = heatmap.astype(self.img1_gray.dtype) return heatmap def find_best_match(self, event,u,v,flags,param): """ For each network, find the best match in the target image to point highlighted with reticle in the source image. Displays the result :return: :rtype: """ img_1_with_reticle = np.copy(self.img1) draw_reticle(img_1_with_reticle, u, v, self._reticle_color) cv2.imshow("source", img_1_with_reticle) alpha = self._config["blend_weight_original_image"] beta = 1 - alpha img_2_with_reticle = np.copy(self.img2) print "\n\n" self._res_uv = dict() # self._res_a_uv = dict() # self._res_b_uv = dict() for network_name in self._dcn_dict: res_a = self._res_a[network_name] res_b = self._res_b[network_name] best_match_uv, best_match_diff, norm_diffs = \ DenseCorrespondenceNetwork.find_best_match((u, v), res_a, res_b) print "\n\n" print "network_name:", network_name self._res_uv[network_name] = dict() self._res_uv[network_name]['source'] = res_a[v, u, :].tolist() self._res_uv[network_name]['target'] = res_b[v, u, :].tolist() # print "res_a[v, u, :]:", res_a[v, u, :] # print "res_b[v, u, :]:", res_b[v, u, :] print "%s best match diff: %.3f" %(network_name, best_match_diff) threshold = self._config["norm_diff_threshold"] if network_name in self._config["norm_diff_threshold_dict"]: threshold = self._config["norm_diff_threshold_dict"][network_name] heatmap = self.scale_norm_diffs_to_make_heatmap(norm_diffs, threshold) reticle_color = self._network_reticle_color[network_name] draw_reticle(heatmap, best_match_uv[0], best_match_uv[1], reticle_color) draw_reticle(img_2_with_reticle, best_match_uv[0], best_match_uv[1], reticle_color) blended = cv2.addWeighted(self.img2_gray, alpha, heatmap, beta, 0) cv2.imshow(network_name, blended) cv2.imshow("target", img_2_with_reticle) if event == cv2.EVENT_LBUTTONDOWN: utils.saveToYaml(self._res_uv, 'clicked_point.yaml') def run(self): self._get_new_images() cv2.namedWindow('target') cv2.setMouseCallback('source', self.find_best_match) self._get_new_images() while True: k = cv2.waitKey(20) & 0xFF if k == 27: break elif k == ord('n'): print "HEY" self._get_new_images() elif k == ord('s'): print "HEY" img1_pil = self.img1_pil img2_pil = self.img2_pil self.img1_pil = img2_pil self.img2_pil = img1_pil self._compute_descriptors()