def test_segmentationFromKeypoints(self): from supermariopy.crf import SegmentationFromKeypoints from supermariopy import imageutils from skimage import data n_keypoints = 10 var = 0.05 keypoints = np.stack( [np.linspace(-1, 1, n_keypoints), np.linspace(-1, 1, n_keypoints)], axis=1) img = data.astronaut() segmentation_algorithm = SegmentationFromKeypoints(var) labels = segmentation_algorithm(img, keypoints) labels_rgb = imageutils.make_colors(n_keypoints + 1)[labels] heatmaps = imageutils.keypoints_to_heatmaps(img.shape[:2], keypoints, var) heatmaps_rgb = imageutils.colorize_heatmaps( heatmaps[np.newaxis, ...], imageutils.make_colors(n_keypoints)) fig, axes = plt.subplots(1, 2, figsize=(8, 4)) axes[0].imshow(labels_rgb) axes[0].set_axis_off() axes[1].imshow(np.squeeze(heatmaps_rgb)) axes[1].set_axis_off() return fig
def batched_keypoints_to_segments( img: np.ndarray, keypoints: np.ndarray, segmentation_algorithm: SegmentationAlgorithm, ) -> Union[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: n_keypoints = keypoints.shape[0] MAP = segmentation_algorithm(img, keypoints) MAP_colorized = imageutils.make_colors(n_keypoints + 1, with_background=True, background_id=0)[MAP] heatmaps = imageutils.keypoints_to_heatmaps(img.shape[:2], keypoints, segmentation_algorithm.var) heatmaps *= heatmaps > 0.8 heatmaps_rgb = imageutils.colorize_heatmaps( heatmaps[np.newaxis, ...], imageutils.make_colors(n_keypoints)) img_resized = cv2.resize(img, (256, 256), cv2.INTER_LINEAR) img_resized = imageutils.convert_range(img_resized, [0, 255], [0, 1]) im_with_keypoints = imageutils.draw_keypoint_markers( img_resized, keypoints, marker_list=[str(i) for i in range(10)] + ["x", "o", "v", "<", ">", "*"], font_scale=1, thickness=4, ) im_with_keypoints = cv2.resize(im_with_keypoints, (img.shape[1], img.shape[1]), cv2.INTER_LINEAR) return MAP, MAP_colorized, heatmaps_rgb, im_with_keypoints
def argmax_rgb(m, cmap=plt.cm.viridis): """Take argmax of m along dimension 1 and apply RGB colorcode on it Parameters ---------- m : [type] [description] Returns ------- np.array RGB mask tensor shaped [B, 3, H, W] """ B, P, H, W = ptutils.nn.shape_as_list(m) max_values, argmax_map = torch.max(m, dim=1) # if m.is_cuda: # dtype = torch.cuda.FloatTensor # else: # dtype = torch.FloatTensor colors = imageutils.make_colors(P, cmap=cmap) colors = colors.astype(np.float32) colors = torch.from_numpy(colors) colors = colors.to(m.device) m_one_hot = ptnn.to_one_hot(argmax_map, P) m_one_hot = m_one_hot.permute(0, 3, 1, 2) mask_rgb = torch.einsum("bphw,pc->bchw", m_one_hot, colors) return mask_rgb
def test_make_colors_bytes(self): from supermariopy.imageutils import make_colors colors = make_colors(10, bytes=True) expected = [0, 255] assert isinstance(colors.flatten()[0], np.integer) assert colors.min() >= expected[0] and colors.max() <= expected[1] assert colors.shape == (10, 3)
def test_make_colors_01(self): from supermariopy.imageutils import make_colors colors = make_colors(10) expected = [0, 1] assert isinstance(colors.flatten()[0], np.floating) assert colors.min() >= expected[0] and colors.max() <= expected[1] assert colors.shape == (10, 3)
def batched_keypoints_to_segments(img, keypoints, segmentation_algorithm): n_keypoints = keypoints.shape[0] MAP = segmentation_algorithm(img, keypoints) MAP_colorized = imageutils.make_colors(n_keypoints + 1, with_background=True, background_id=0)[MAP] heatmaps = imageutils.keypoints_to_heatmaps(img.shape[:2], keypoints, segmentation_algorithm.var) heatmaps *= heatmaps > 0.8 heatmaps_rgb = imageutils.colorize_heatmaps( heatmaps[np.newaxis, ...], imageutils.make_colors(n_keypoints)) im_with_keypoints = imageutils.draw_keypoint_markers( imageutils.convert_range(img, [0, 255], [0, 1]), keypoints, marker_list=[str(i) for i in range(10)] + ["x", "o", "v", "<", ">", "*"], font_scale=0.5, ) return MAP, MAP_colorized, heatmaps_rgb, im_with_keypoints
def main(iuv_path): I, u, v = denseposelib.load_iuv(iuv_path) base_dir = os.path.dirname(iuv_path) out_name = "{}_IRGB.png".format(os.path.splitext(iuv_path)[0]) colors = imageutils.make_colors( len(denseposelib.PART_LIST), cmap=plt.cm.coolwarm, with_background=True, background_id=0, ) I_colors = imageutils.convert_range(colors[I], [0, 1], [0, 255]) cv2.imwrite(os.path.join(base_dir, out_name), I_colors)
def test_colorize_heatmaps(self): from supermariopy.imageutils import ( keypoints_to_heatmaps, make_colors, colorize_heatmaps, ) from matplotlib import pyplot as plt keypoints = np.stack([np.linspace(-1, 1, 10), np.linspace(-1, 1, 10)], axis=1) heatmaps = keypoints_to_heatmaps((512, 512), keypoints) colors = make_colors(keypoints.shape[0]) img_marked = colorize_heatmaps(heatmaps[np.newaxis, ...], colors) fig, ax = plt.subplots(1, 1, figsize=(6, 6)) ax.imshow(np.squeeze(img_marked)) return fig
def argmax_rgb(m, cmap=plt.cm.viridis): """Take argmax of m along dimension 1 and apply RGB colorcode on it Parameters ---------- m : tf.Tensor Tensorflor tensor or numpy array as result of eager execution Returns ------- tf.Tensor RGB mask tensor shaped [B, H, W, 3] """ B, H, W, P = nn.shape_as_list(m) argmax_map = tf.arg_max(m, dimension=-1) colors = imageutils.make_colors(P, cmap=cmap) colors = colors.astype(np.float32) colors = tf.convert_to_tensor(colors) m_one_hot = tf.one_hot(argmax_map, P, axis=-1) mask_rgb = tf.einsum("bhwp,pc->bhwc", m_one_hot, colors) return mask_rgb
def show_active_parts( parts, active_idx, inactive_parts_color=np.array([0.8, 0.8, 0.8]), background_color=np.array([1.0, 1.0, 1.0]), cmap=plt.cm.viridis, ): """Highlight active parts in color and set color of other parts to a gray value Parameters ---------- parts : np.ndarray tensor of binary masks indicating mutually exclusive parts or segments of a segmentation. Shaped [N, H, W, P] active_idx : np.array or list list of indices indicating which part (indexing into last dimension of parts) are active. inactive_parts_color : np.ndarray, default np.array([0.8, 0.8, 0.8]) Color value in range [0, 1] for inactive parts. background_color : np.ndarray, default np.array([1.0, 1.0, 1.0]) Color value in range [0, 1] for background. Note that this only makes sense if parts does not provide a background itself. If part contains a background part, it will be assigned a background_color specified by the color map. """ if not any( [isinstance(active_idx, list), isinstance(active_idx, np.ndarray)]): raise TypeError() N, H, W, P = parts.shape colors = imageutils.make_colors(P, cmap) all_part_idx = set(range(P)) inactive_part_idx = all_part_idx.difference(set(active_idx)) colors[ np.array(list(inactive_part_idx)), :] = inactive_parts_color.reshape( (1, 3)) background_part = 1.0 - np.clip(np.sum(parts, axis=-1), 0, 1) colors = np.insert(colors, 0, background_color.reshape(1, 3), axis=0) parts = np.insert(parts, 0, background_part, axis=-1) parts_rgb = np.einsum("nhwp,pc->nhwc", parts, colors) return parts_rgb
def argmax_rgb(m, cmap=plt.cm.viridis): """Take argmax of m along dimension 1 and apply RGB colorcode on it Parameters ---------- m : np.ndarray tensor representing list of binary masks, shaped [B, H, W, P] Returns ------- np.array RGB mask tensor shaped [B, H, W, 3] See also -------- .. tfutils.viz.argmax_rgb """ B, H, W, P = m.shape argmax_map = np.argmax(m, axis=-1) colors = imageutils.make_colors(P, cmap=cmap) colors = colors.astype(np.float32) m_one_hot = npu.one_hot(argmax_map, P, axis=-1) mask_rgb = np.einsum("bhwp,pc->bhwc", m_one_hot, colors) return mask_rgb
def main(infer_dir, output_folder, run_crf_config, n_processes): os.makedirs(output_folder, exist_ok=True) with open(run_crf_config, "r") as f: config = yaml.load(f) segmentation_algorithm_args = config["segmentation_algorithm_args"] npz_files = glob.glob(os.path.join(infer_dir, "*.npz")) npz_files = sorted(npz_files) print("Using files :") print(npz_files) segmentation_algorithm = crf.SegmentationFromKeypoints( **segmentation_algorithm_args) data = [] with closing(Pool(n_processes)) as p: for outputs in tqdm.tqdm(p.imap(load_npz, npz_files)): data.append(outputs) data = list_of_dicts2dict_of_lists(data) data = {k: np.concatenate(data[k]) for k in ["image", "gauss_yx"]} data["gauss_yx"] = data["gauss_yx"][..., ::-1] process_func = functools.partial( process_batches, **{ "segmentation_algorithm": segmentation_algorithm, }) tuples = list( zip(np.array_split(data["image"], n_processes, 0), np.array_split(data["gauss_yx"], n_processes, 0))) processed_data = [] with closing(Pool(n_processes)) as p: for outputs in tqdm.tqdm(p.imap(process_func, tuples)): processed_data.append(outputs) labels = np.concatenate([p["labels"] for p in processed_data], 0) labels_rgb = np.concatenate([p["labels_rgb"] for p in processed_data], 0) heatmaps = np.concatenate([p["heatmaps"] for p in processed_data], 0) ims_with_keypoints = np.concatenate( [p["ims_with_keypoints"] for p in processed_data], 0) target_dir = os.path.join(output_folder, "01_keypoints") os.makedirs(target_dir, exist_ok=True) write_rgb(ims_with_keypoints, target_dir, n_processes) target_dir = os.path.join(output_folder, "02_heatmaps") os.makedirs(target_dir, exist_ok=True) write_rgb(heatmaps, target_dir, n_processes) target_dir = os.path.join(output_folder, "03_labels_rgb") os.makedirs(target_dir, exist_ok=True) write_rgb(labels_rgb, target_dir, n_processes) densepose_csv_path = config["densepose_csv_path"] data_root = config["data_root"] fname_col = config["data_fname_col"] iuv_files = get_iuv_files(densepose_csv_path, data_root, len(labels), fname_col) iuvs = [cv2.imread(x, -1) for x in iuv_files] iuvs = [ denseposelib.resize_labels(i[..., 0], labels.shape[1:]) for i in iuvs ] iuvs = np.stack(iuvs, axis=0) dp_semantic_remap_dict = config["dp_semantic_remap_dict"] dp_new_part_list = sorted(list(dp_semantic_remap_dict.keys())) dp_remap_dict = denseposelib.semantic_remap_dict2remap_dict( dp_semantic_remap_dict, dp_new_part_list) remapped_gt_segmentation, remapped_inferred = denseposelib.get_best_segmentation( iuvs, labels, dp_remap_dict) df = pd.DataFrame(columns=["batch_idx"] + dp_new_part_list) df = denseposelib.calculate_iou_df(remapped_inferred, remapped_gt_segmentation, dp_new_part_list) df.to_csv(os.path.join(output_folder, "part_ious.csv"), index=False, header=True) df_mean = denseposelib.calculate_overall_iou_from_df(df) with open(os.path.join(output_folder, "mean_part_ios.csv"), "w") as f: print( tabulate(df_mean, headers="keys", tablefmt="psql", showindex="never"), file=f, ) target_dir = os.path.join(output_folder, "04_compare") os.makedirs(target_dir, exist_ok=True) background_color = np.array([1, 1, 1]) colors1 = imageutils.make_colors(config["n_inferred_parts"] + 1, with_background=True, background_id=0) colors2 = imageutils.make_colors( len(dp_new_part_list), with_background=True, background_id=dp_new_part_list.index("background")) for i, (im1, im2, im3) in enumerate( zip(labels, remapped_inferred, remapped_gt_segmentation)): canvas = np.concatenate([colors1[im1], colors2[im2], colors2[im3]], 1).astype(np.float32) canvas = cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR) fname = os.path.join(target_dir, "{:06d}.png".format(i)) cv2.imwrite(fname, imageutils.convert_range(canvas, [0, 1], [0, 255])) # batches.plot_batch( # imageutils.convert_range(canvas, [0, 1], [-1, 1]), fname, cols=3 # ) target_dir = os.path.join(output_folder, "05_remapped_inferred") os.makedirs(target_dir, exist_ok=True) write_labels(remapped_inferred, target_dir, colors2, n_processes) target_dir = os.path.join(output_folder, "06_remapped_labels") os.makedirs(target_dir, exist_ok=True) write_labels(remapped_gt_segmentation, target_dir, colors2, n_processes)
def draw_keypoint_markers( img: np.ndarray, keypoints: np.ndarray, font_scale: float = 0.5, thickness: int = 2, font=cv2.FONT_HERSHEY_SIMPLEX, marker_list=["o", "v", "x", "+", "<", "-", ">", "c"], ) -> np.ndarray: """ Draw keypoints on image with markers Parameters ---------- img : np.ndarray shaped [H, W, 3] array in range [0, 1] keypoints : np.ndarray shaped [kp, 2] - array giving keypoint positions in range [-1, 1] for x and y. keypoints[:, 0] is x-coordinate (horizontal). font_scale : int, optional openCV font scale passed to 'cv2.putText', by default 1 thickness : int, optional openCV font thickness passed to 'cv2.putText', by default 2 font : cv2.FONT_xxx, optional openCV font, by default cv2.FONT_HERSHEY_SIMPLEX Examples -------- from skimage import data astronaut = data.astronaut() keypoints = np.stack([np.linspace(-1, 1, 10), np.linspace(-1, 1, 10)], axis=1) img_marked = draw_keypoint_markers(astronaut, keypoints, font_scale=2, thickness=3) plt.imshow(img_marked) """ if not imageutils.is_in_range(img, [0, 1]): raise RangeError(img, "img", [0, 1]) if img.shape[0] != img.shape[1]: raise ValueError("only square images are supported currently") img_marked = img.copy() keypoints = imageutils.convert_range(keypoints, [-1, 1], [0, img.shape[0] - 1]) colors = imageutils.make_colors( keypoints.shape[0], bytes=False, cmap=plt.cm.inferno ) for i, kp in enumerate(keypoints): text = marker_list[i % len(marker_list)] (label_width, label_height), baseline = cv2.getTextSize( text, font, font_scale, thickness ) textX = kp[0] textY = kp[1] font_color = colors[i] text_position = ( textX - label_width / 2.0 - baseline, textY - label_height / 2.0 + baseline, ) text_position = tuple([int(x) for x in text_position]) img_marked = cv2.putText( img_marked, text, text_position, font, font_scale, font_color, thickness=thickness, ) return img_marked