def warp_crop_faces(self, save_cropped_path=None): """Get affine matrix, warp and cropped faces. Also get inverse affine matrix for post-processing. """ for idx, landmark in enumerate(self.all_landmarks_5): # use 5 landmarks to get affine matrix self.similarity_trans.estimate(landmark, self.face_template) affine_matrix = self.similarity_trans.params[0:2, :] self.affine_matrices.append(affine_matrix) # warp and crop faces cropped_face = cv2.warpAffine(self.input_img, affine_matrix, self.out_size) self.cropped_faces.append(cropped_face) # save the cropped face if save_cropped_path is not None: path, ext = os.path.splitext(save_cropped_path) save_path = f'{path}_{idx:02d}{ext}' mmcv.imwrite(mmcv.rgb2bgr(cropped_face), save_path) # get inverse affine matrix self.similarity_trans.estimate(self.face_template, landmark * self.upscale_factor) inverse_affine = self.similarity_trans.params[0:2, :] self.inverse_affine_matrices.append(inverse_affine)
def test_rgb2bgr(self): in_img = np.random.rand(10, 10, 3).astype(np.float32) out_img = mmcv.rgb2bgr(in_img) assert out_img.shape == in_img.shape assert_array_equal(out_img[..., 0], in_img[..., 2]) assert_array_equal(out_img[..., 1], in_img[..., 1]) assert_array_equal(out_img[..., 2], in_img[..., 0])
def paste_faces_to_input_image(self, save_path): # operate in the BGR order input_img = mmcv.rgb2bgr(self.input_img) h, w, _ = input_img.shape h_up, w_up = h * self.upscale_factor, w * self.upscale_factor # simply resize the background upsample_img = cv2.resize(input_img, (w_up, h_up)) for restored_face, inverse_affine in zip(self.restored_faces, self.inverse_affine_matrices): inv_restored = cv2.warpAffine(restored_face, inverse_affine, (w_up, h_up)) mask = np.ones((*self.out_size, 3), dtype=np.float32) inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up)) # remove the black borders inv_mask_erosion = cv2.erode( inv_mask, np.ones((2 * self.upscale_factor, 2 * self.upscale_factor), np.uint8)) inv_restored_remove_border = inv_mask_erosion * inv_restored total_face_area = np.sum(inv_mask_erosion) // 3 # compute the fusion edge based on the area of face w_edge = int(total_face_area**0.5) // 20 erosion_radius = w_edge * 2 inv_mask_center = cv2.erode( inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8)) blur_size = w_edge * 2 inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0) upsample_img = inv_soft_mask * inv_restored_remove_border + ( 1 - inv_soft_mask) * upsample_img mmcv.imwrite(upsample_img.astype(np.uint8), save_path)
def __call__(self, results): img = results['img'] # bgr img = Image.fromarray(mmcv.bgr2rgb(img)) # rgb img = np.array(self.color_jitter(img)) img = mmcv.rgb2bgr(img) results['img'] = img return results
def main(): bgr_img = mmcv.imread(image_path) h, w, _ = bgr_img.shape # convert color rgb_img = mmcv.bgr2rgb(bgr_img) # resize resize_img = mmcv.imresize(rgb_img, size=(256, 256)) # rotate rotate_img = mmcv.imrotate(rgb_img, angle=45) # flip flip_img = mmcv.imflip(rgb_img, direction='horizontal') # crop if h <= w: y_min, y_max = 0, h x_min = int((w - h) / 2) x_max = x_min + h else: x_min, x_max = 0, h y_min = int((h - w) / 2) y_max = y_min + w bbox = np.array([x_min, y_min, x_max, y_max]) crop_img = mmcv.imcrop(rgb_img, bbox) # padding max_size = max(h, w) pad_img = mmcv.impad(rgb_img, shape=(max_size, max_size), padding_mode='constant') mmcv.imshow(mmcv.rgb2bgr(pad_img))
def put_img_infos(self, img, infos, text_color='white', font_size=26, row_width=20, win_name='', show=True, wait_time=0, out_file=None): """Show image with extra information. Args: img (str | ndarray): The image to be displayed. infos (dict): Extra infos to display in the image. text_color (:obj:`mmcv.Color`/str/tuple/int/ndarray): Extra infos display color. Defaults to 'white'. font_size (int): Extra infos display font size. Defaults to 26. row_width (int): width between each row of results on the image. win_name (str): The image title. Defaults to '' show (bool): Whether to show the image. Defaults to True. wait_time (int): How many seconds to display the image. Defaults to 0. out_file (Optional[str]): The filename to write the image. Defaults to None. Returns: np.ndarray: The image with extra infomations. """ self.prepare() text_color = color_val_matplotlib(text_color) img = mmcv.imread(img).astype(np.uint8) x, y = 3, row_width // 2 img = mmcv.bgr2rgb(img) width, height = img.shape[1], img.shape[0] img = np.ascontiguousarray(img) # add a small EPS to avoid precision lost due to matplotlib's # truncation (https://github.com/matplotlib/matplotlib/issues/15363) dpi = self.fig_save.get_dpi() self.fig_save.set_size_inches((width + EPS) / dpi, (height + EPS) / dpi) for k, v in infos.items(): if isinstance(v, float): v = f'{v:.2f}' label_text = f'{k}: {v}' self._put_text(self.ax_save, label_text, x, y, text_color, font_size) if show and not self.is_inline: self._put_text(self.ax_show, label_text, x, y, text_color, font_size) y += row_width self.ax_save.imshow(img) stream, _ = self.fig_save.canvas.print_to_buffer() buffer = np.frombuffer(stream, dtype='uint8') img_rgba = buffer.reshape(height, width, 4) rgb, _ = np.split(img_rgba, [3], axis=2) img_save = rgb.astype('uint8') img_save = mmcv.rgb2bgr(img_save) if out_file is not None: mmcv.imwrite(img_save, out_file) ret = 0 if show and not self.is_inline: # Reserve some space for the tip. self.ax_show.set_title(win_name) self.ax_show.set_ylim(height + 20) self.ax_show.text(width // 2, height + 18, 'Press SPACE to continue.', ha='center', fontsize=font_size) self.ax_show.imshow(img) # Refresh canvas, necessary for Qt5 backend. self.fig_show.canvas.draw() ret = self.wait_continue(timeout=wait_time) elif (not show) and self.is_inline: # If use inline backend, we use fig_save to show the image # So we need to close it if users don't want to show. plt.close(self.fig_save) return ret, img_save
def imshow_det_bboxes(img, bboxes, labels, segms=None, class_names=None, score_thr=0, bbox_color='green', text_color='green', mask_color=None, thickness=2, font_size=13, win_name='', show=True, wait_time=0, out_file=None, visual_label_threshold=-1, visual_label_offset=0): """Draw bboxes and class labels (with scores) on an image. Args: img (str or ndarray): The image to be displayed. bboxes (ndarray): Bounding boxes (with scores), shaped (n, 4) or (n, 5). labels (ndarray): Labels of bboxes. segms (ndarray or None): Masks, shaped (n,h,w) or None class_names (list[str]): Names of each classes. score_thr (float): Minimum score of bboxes to be shown. Default: 0 bbox_color (str or tuple(int) or :obj:`Color`):Color of bbox lines. The tuple of color should be in BGR order. Default: 'green' text_color (str or tuple(int) or :obj:`Color`):Color of texts. The tuple of color should be in BGR order. Default: 'green' mask_color (str or tuple(int) or :obj:`Color`, optional): Color of masks. The tuple of color should be in BGR order. Default: None thickness (int): Thickness of lines. Default: 2 font_size (int): Font size of texts. Default: 13 show (bool): Whether to show the image. Default: True win_name (str): The window name. Default: '' wait_time (float): Value of waitKey param. Default: 0. out_file (str, optional): The filename to write the image. Default: None visual_label_threshold(int, optional): Bbox showed label limit. visual_label_offset(int, optional): Bbox showed label offset. Returns: ndarray: The image with bboxes drawn on it. """ assert bboxes.ndim == 2, \ f' bboxes ndim should be 2, but its ndim is {bboxes.ndim}.' assert labels.ndim == 1, \ f' labels ndim should be 1, but its ndim is {labels.ndim}.' assert bboxes.shape[0] == labels.shape[0], \ 'bboxes.shape[0] and labels.shape[0] should have the same length.' assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5, \ f' bboxes.shape[1] should be 4 or 5, but its {bboxes.shape[1]}.' img = mmcv.imread(img).astype(np.uint8) if score_thr > 0: assert bboxes.shape[1] == 5 scores = bboxes[:, -1] inds = scores > score_thr bboxes = bboxes[inds, :] labels = labels[inds] if segms is not None: segms = segms[inds, ...] mask_colors = [] if labels.shape[0] > 0: if mask_color is None: # random color np.random.seed(42) mask_colors = [ np.random.randint(0, 256, (1, 3), dtype=np.uint8) for _ in range(max(labels) + 1) ] else: # specify color mask_colors = [ np.array(mmcv.color_val(mask_color)[::-1], dtype=np.uint8) ] * (max(labels) + 1) bbox_color = color_val_matplotlib(bbox_color) text_color = color_val_matplotlib(text_color) img = mmcv.bgr2rgb(img) width, height = img.shape[1], img.shape[0] img = np.ascontiguousarray(img) fig = plt.figure(win_name, frameon=False) plt.title(win_name) canvas = fig.canvas dpi = fig.get_dpi() # add a small EPS to avoid precision lost due to matplotlib's truncation # (https://github.com/matplotlib/matplotlib/issues/15363) fig.set_size_inches((width + EPS) / dpi, (height + EPS) / dpi) # remove white edges by set subplot margin plt.subplots_adjust(left=0, right=1, bottom=0, top=1) ax = plt.gca() ax.axis('off') polygons = [] color = [] for i, (bbox, label) in enumerate(zip(bboxes, labels)): if visual_label_threshold < 0 or visual_label_offset <= label <= visual_label_threshold: bbox_int = bbox.astype(np.int32) poly = [[bbox_int[0], bbox_int[1]], [bbox_int[0], bbox_int[3]], [bbox_int[2], bbox_int[3]], [bbox_int[2], bbox_int[1]]] np_poly = np.array(poly).reshape((4, 2)) polygons.append(Polygon(np_poly)) color.append(bbox_color) label_text = class_names[ label] if class_names is not None else f'class {label}' if len(bbox) > 4: label_text += f'|{bbox[-1]:.02f}' ax.text(bbox_int[0], bbox_int[1], f'{label_text}', bbox={ 'facecolor': 'black', 'alpha': 0.8, 'pad': 0.7, 'edgecolor': 'none' }, color=text_color, fontsize=font_size, verticalalignment='top', horizontalalignment='left') if segms is not None: color_mask = mask_colors[labels[i]] mask = segms[i].astype(bool) img[mask] = img[mask] * 0.5 + color_mask * 0.5 plt.imshow(img) p = PatchCollection(polygons, facecolor='none', edgecolors=color, linewidths=thickness) ax.add_collection(p) stream, _ = canvas.print_to_buffer() buffer = np.frombuffer(stream, dtype='uint8') img_rgba = buffer.reshape(height, width, 4) rgb, alpha = np.split(img_rgba, [3], axis=2) img = rgb.astype('uint8') img = mmcv.rgb2bgr(img) if show: # We do not use cv2 for display because in some cases, opencv will # conflict with Qt, it will output a warning: Current thread # is not the object's thread. You can refer to # https://github.com/opencv/opencv-python/issues/46 for details if wait_time == 0: plt.show() else: plt.show(block=False) plt.pause(wait_time) if out_file is not None: mmcv.imwrite(img, out_file) plt.close() return img
def imshow_det_bboxes(img, bboxes, labels, segms=None, class_names=None, score_thr=0, bbox_color='green', text_color='green', mask_color=None, #thickness=2, thickness=0, # clw modify font_scale=0.5, #font_size=13, font_size=6, win_name='', #fig_size=(15, 10), fig_size=(3, 2), show=True, wait_time=0, out_file=None): """Draw bboxes and class labels (with scores) on an image. Args: img (str or ndarray): The image to be displayed. bboxes (ndarray): Bounding boxes (with scores), shaped (n, 4) or (n, 5). labels (ndarray): Labels of bboxes. segms (ndarray or None): Masks, shaped (n,h,w) or None class_names (list[str]): Names of each classes. score_thr (float): Minimum score of bboxes to be shown. Default: 0 bbox_color (str or tuple(int) or :obj:`Color`):Color of bbox lines. The tuple of color should be in BGR order. Default: 'green' text_color (str or tuple(int) or :obj:`Color`):Color of texts. The tuple of color should be in BGR order. Default: 'green' mask_color (None or str or tuple(int) or :obj:`Color`): Color of masks. The tuple of color should be in BGR order. Default: None thickness (int): Thickness of lines. Default: 2 font_scale (float): Font scales of texts. Default: 0.5 font_size (int): Font size of texts. Default: 13 show (bool): Whether to show the image. Default: True win_name (str): The window name. Default: '' fig_size (tuple): Figure size of the pyplot figure. Default: (15, 10) wait_time (float): Value of waitKey param. Default: 0. out_file (str or None): The filename to write the image. Default: None Returns: ndarray: The image with bboxes drawn on it. """ warnings.warn('"font_scale" will be deprecated in v2.9.0,' 'Please use "font_size"') assert bboxes.ndim == 2, \ f' bboxes ndim should be 2, but its ndim is {bboxes.ndim}.' assert labels.ndim == 1, \ f' labels ndim should be 1, but its ndim is {labels.ndim}.' assert bboxes.shape[0] == labels.shape[0], \ 'bboxes.shape[0] and labels.shape[0] should have the same length.' assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5,\ f' bboxes.shape[1] should be 4 or 5, but its {bboxes.shape[1]}.' img = mmcv.imread(img).copy() if score_thr > 0: assert bboxes.shape[1] == 5 scores = bboxes[:, -1] inds = scores > score_thr bboxes = bboxes[inds, :] labels = labels[inds] if segms is not None: segms = segms[inds, ...] mask_colors = [] if labels.shape[0] > 0: if mask_color is None: # random color np.random.seed(42) mask_colors = [ np.random.randint(0, 256, (1, 3), dtype=np.uint8) for _ in range(max(labels) + 1) ] else: # specify color mask_colors = [ np.array(mmcv.color_val(mask_color)[::-1], dtype=np.uint8) ] * ( max(labels) + 1) bbox_color = color_val_matplotlib(bbox_color) text_color = color_val_matplotlib(text_color) img = mmcv.bgr2rgb(img) img = np.ascontiguousarray(img) plt.figure(win_name, figsize=fig_size) plt.title(win_name) plt.axis('off') ax = plt.gca() polygons = [] color = [] for i, (bbox, label) in enumerate(zip(bboxes, labels)): bbox_int = bbox.astype(np.int32) poly = [[bbox_int[0], bbox_int[1]], [bbox_int[0], bbox_int[3]], [bbox_int[2], bbox_int[3]], [bbox_int[2], bbox_int[1]]] np_poly = np.array(poly).reshape((4, 2)) polygons.append(Polygon(np_poly)) color.append(bbox_color) label_text = class_names[ label] if class_names is not None else f'class {label}' if len(bbox) > 4: label_text += f'|{bbox[-1]:.02f}' ax.text( bbox_int[0], bbox_int[1], f'{label_text}', bbox={ 'facecolor': 'black', #'alpha': 0.8, 'alpha': 0, # clw modify #'pad': 0.7, 'pad': 0, # clw modify 'edgecolor': 'none' }, color=text_color, fontsize=font_size, verticalalignment='top', horizontalalignment='left') if segms is not None: color_mask = mask_colors[labels[i]] mask = segms[i].astype(bool) img[mask] = img[mask] * 0.5 + color_mask * 0.5 plt.imshow(img) p = PatchCollection( polygons, facecolor='none', edgecolors=color, linewidths=thickness) ax.add_collection(p) if out_file is not None: dir_name = osp.abspath(osp.dirname(out_file)) mmcv.mkdir_or_exist(dir_name) plt.savefig(out_file) #plt.savefig(out_file, dpi=600) # clw modify: too slow if not show: plt.close() if show: if wait_time == 0: plt.show() else: plt.show(block=False) plt.pause(wait_time) plt.close() return mmcv.rgb2bgr(img)
def imshow_det_bboxes(img, bboxes=None, labels=None, segms=None, class_names=None, score_thr=0, bbox_color='green', text_color='green', mask_color=None, thickness=2, font_size=8, win_name='', show=True, wait_time=0, out_file=None): """Draw bboxes and class labels (with scores) on an image. Args: img (str | ndarray): The image to be displayed. bboxes (ndarray): Bounding boxes (with scores), shaped (n, 4) or (n, 5). labels (ndarray): Labels of bboxes. segms (ndarray | None): Masks, shaped (n,h,w) or None. class_names (list[str]): Names of each classes. score_thr (float): Minimum score of bboxes to be shown. Default: 0. bbox_color (list[tuple] | tuple | str | None): Colors of bbox lines. If a single color is given, it will be applied to all classes. The tuple of color should be in RGB order. Default: 'green'. text_color (list[tuple] | tuple | str | None): Colors of texts. If a single color is given, it will be applied to all classes. The tuple of color should be in RGB order. Default: 'green'. mask_color (list[tuple] | tuple | str | None, optional): Colors of masks. If a single color is given, it will be applied to all classes. The tuple of color should be in RGB order. Default: None. thickness (int): Thickness of lines. Default: 2. font_size (int): Font size of texts. Default: 13. show (bool): Whether to show the image. Default: True. win_name (str): The window name. Default: ''. wait_time (float): Value of waitKey param. Default: 0. out_file (str, optional): The filename to write the image. Default: None. Returns: ndarray: The image with bboxes drawn on it. """ assert bboxes is None or bboxes.ndim == 2, \ f' bboxes ndim should be 2, but its ndim is {bboxes.ndim}.' assert labels.ndim == 1, \ f' labels ndim should be 1, but its ndim is {labels.ndim}.' assert bboxes is None or bboxes.shape[1] == 4 or bboxes.shape[1] == 5, \ f' bboxes.shape[1] should be 4 or 5, but its {bboxes.shape[1]}.' assert bboxes is None or bboxes.shape[0] <= labels.shape[0], \ 'labels.shape[0] should not be less than bboxes.shape[0].' assert segms is None or segms.shape[0] == labels.shape[0], \ 'segms.shape[0] and labels.shape[0] should have the same length.' assert segms is not None or bboxes is not None, \ 'segms and bboxes should not be None at the same time.' img = mmcv.imread(img).astype(np.uint8) if score_thr > 0: assert bboxes is not None and bboxes.shape[1] == 5 scores = bboxes[:, -1] inds = scores > score_thr bboxes = bboxes[inds, :] labels = labels[inds] if segms is not None: segms = segms[inds, ...] img = mmcv.bgr2rgb(img) width, height = img.shape[1], img.shape[0] img = np.ascontiguousarray(img) fig = plt.figure(win_name, frameon=False) plt.title(win_name) canvas = fig.canvas dpi = fig.get_dpi() # add a small EPS to avoid precision lost due to matplotlib's truncation # (https://github.com/matplotlib/matplotlib/issues/15363) fig.set_size_inches((width + EPS) / dpi, (height + EPS) / dpi) # remove white edges by set subplot margin plt.subplots_adjust(left=0, right=1, bottom=0, top=1) ax = plt.gca() ax.axis('off') max_label = int(max(labels) if len(labels) > 0 else 0) text_palette = palette_val(get_palette(text_color, max_label + 1)) text_colors = [text_palette[label] for label in labels] num_bboxes = 0 if bboxes is not None: num_bboxes = bboxes.shape[0] bbox_palette = palette_val(get_palette(bbox_color, max_label + 1)) colors = [bbox_palette[label] for label in labels[:num_bboxes]] draw_bboxes(ax, bboxes, colors, alpha=0.8, thickness=thickness) horizontal_alignment = 'left' positions = bboxes[:, :2].astype(np.int32) + thickness areas = (bboxes[:, 3] - bboxes[:, 1]) * (bboxes[:, 2] - bboxes[:, 0]) scales = _get_adaptive_scales(areas) scores = bboxes[:, 4] if bboxes.shape[1] == 5 else None draw_labels(ax, labels[:num_bboxes], positions, scores=scores, class_names=class_names, color=text_colors, font_size=font_size, scales=scales, horizontal_alignment=horizontal_alignment) if segms is not None: mask_palette = get_palette(mask_color, max_label + 1) colors = [mask_palette[label] for label in labels] colors = np.array(colors, dtype=np.uint8) draw_masks(ax, img, segms, colors, with_edge=True) if num_bboxes < segms.shape[0]: segms = segms[num_bboxes:] horizontal_alignment = 'center' areas = [] positions = [] for mask in segms: _, _, stats, centroids = cv2.connectedComponentsWithStats( mask.astype(np.uint8), connectivity=8) largest_id = np.argmax(stats[1:, -1]) + 1 positions.append(centroids[largest_id]) areas.append(stats[largest_id, -1]) areas = np.stack(areas, axis=0) scales = _get_adaptive_scales(areas) draw_labels(ax, labels[num_bboxes:], positions, class_names=class_names, color=text_colors, font_size=font_size, scales=scales, horizontal_alignment=horizontal_alignment) plt.imshow(img) stream, _ = canvas.print_to_buffer() buffer = np.frombuffer(stream, dtype='uint8') img_rgba = buffer.reshape(height, width, 4) rgb, alpha = np.split(img_rgba, [3], axis=2) img = rgb.astype('uint8') img = mmcv.rgb2bgr(img) if show: # We do not use cv2 for display because in some cases, opencv will # conflict with Qt, it will output a warning: Current thread # is not the object's thread. You can refer to # https://github.com/opencv/opencv-python/issues/46 for details if wait_time == 0: plt.show() else: plt.show(block=False) plt.pause(wait_time) if out_file is not None: mmcv.imwrite(img, out_file) plt.close() return img
def imshow_keypoints_3d( pose_result, img=None, skeleton=None, pose_kpt_color=None, pose_link_color=None, vis_height=400, kpt_score_thr=0.3, num_instances=-1, *, axis_azimuth=70, axis_limit=1.7, axis_dist=10.0, axis_elev=15.0, ): """Draw 3D keypoints and links in 3D coordinates. Args: pose_result (list[dict]): 3D pose results containing: - "keypoints_3d" ([K,4]): 3D keypoints - "title" (str): Optional. A string to specify the title of the visualization of this pose result img (str|np.ndarray): Opptional. The image or image path to show input image and/or 2D pose. Note that the image should be given in BGR channel order. skeleton (list of [idx_i,idx_j]): Skeleton described by a list of links, each is a pair of joint indices. pose_kpt_color (np.ndarray[Nx3]`): Color of N keypoints. If None, do not nddraw keypoints. pose_link_color (np.array[Mx3]): Color of M links. If None, do not draw links. vis_height (int): The image height of the visualization. The width will be N*vis_height depending on the number of visualized items. kpt_score_thr (float): Minimum score of keypoints to be shown. Default: 0.3. num_instances (int): Number of instances to be shown in 3D. If smaller than 0, all the instances in the pose_result will be shown. Otherwise, pad or truncate the pose_result to a length of num_instances. axis_azimuth (float): axis azimuth angle for 3D visualizations. axis_dist (float): axis distance for 3D visualizations. axis_elev (float): axis elevation view angle for 3D visualizations. axis_limit (float): The axis limit to visualize 3d pose. The xyz range will be set as: - x: [x_c - axis_limit/2, x_c + axis_limit/2] - y: [y_c - axis_limit/2, y_c + axis_limit/2] - z: [0, axis_limit] Where x_c, y_c is the mean value of x and y coordinates figsize: (float): figure size in inch. """ show_img = img is not None if num_instances < 0: num_instances = len(pose_result) else: if len(pose_result) > num_instances: pose_result = pose_result[:num_instances] elif len(pose_result) < num_instances: pose_result += [dict()] * (num_instances - len(pose_result)) num_axis = num_instances + 1 if show_img else num_instances plt.ioff() fig = plt.figure(figsize=(vis_height * num_axis * 0.01, vis_height * 0.01)) if show_img: img = mmcv.imread(img, channel_order='bgr') img = mmcv.bgr2rgb(img) img = mmcv.imrescale(img, scale=vis_height / img.shape[0]) ax_img = fig.add_subplot(1, num_axis, 1) ax_img.get_xaxis().set_visible(False) ax_img.get_yaxis().set_visible(False) ax_img.set_axis_off() ax_img.set_title('Input') ax_img.imshow(img, aspect='equal') for idx, res in enumerate(pose_result): dummy = len(res) == 0 kpts = np.zeros((1, 3)) if dummy else res['keypoints_3d'] if kpts.shape[1] == 3: kpts = np.concatenate([kpts, np.ones((kpts.shape[0], 1))], axis=1) valid = kpts[:, 3] >= kpt_score_thr ax_idx = idx + 2 if show_img else idx + 1 ax = fig.add_subplot(1, num_axis, ax_idx, projection='3d') ax.view_init( elev=axis_elev, azim=axis_azimuth, ) x_c = np.mean(kpts[valid, 0]) if sum(valid) > 0 else 0 y_c = np.mean(kpts[valid, 1]) if sum(valid) > 0 else 0 ax.set_xlim3d([x_c - axis_limit / 2, x_c + axis_limit / 2]) ax.set_ylim3d([y_c - axis_limit / 2, y_c + axis_limit / 2]) ax.set_zlim3d([0, axis_limit]) ax.set_aspect('auto') ax.set_xticks([]) ax.set_yticks([]) ax.set_zticks([]) ax.set_xticklabels([]) ax.set_yticklabels([]) ax.set_zticklabels([]) ax.dist = axis_dist if not dummy and pose_kpt_color is not None: pose_kpt_color = np.array(pose_kpt_color) assert len(pose_kpt_color) == len(kpts) x_3d, y_3d, z_3d = np.split(kpts[:, :3], [1, 2], axis=1) # matplotlib uses RGB color in [0, 1] value range _color = pose_kpt_color[..., ::-1] / 255. ax.scatter( x_3d[valid], y_3d[valid], z_3d[valid], marker='o', color=_color[valid], ) if not dummy and skeleton is not None and pose_link_color is not None: pose_link_color = np.array(pose_link_color) assert len(pose_link_color) == len(skeleton) for link, link_color in zip(skeleton, pose_link_color): link_indices = [_i for _i in link] xs_3d = kpts[link_indices, 0] ys_3d = kpts[link_indices, 1] zs_3d = kpts[link_indices, 2] kpt_score = kpts[link_indices, 3] if kpt_score.min() > kpt_score_thr: # matplotlib uses RGB color in [0, 1] value range _color = link_color[::-1] / 255. ax.plot(xs_3d, ys_3d, zs_3d, color=_color, zdir='z') if 'title' in res: ax.set_title(res['title']) # convert figure to numpy array fig.tight_layout() fig.canvas.draw() img_w, img_h = fig.canvas.get_width_height() img_vis = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8).reshape(img_h, img_w, -1) img_vis = mmcv.rgb2bgr(img_vis) plt.close(fig) return img_vis
def create_frame_by_matplotlib(image_dir, nrows=1, fig_size=(300, 300), font_size=15): """Create gif frame image through matplotlib. Args: image_dir (str): Root directory of result images nrows (int): Number of rows displayed, Default: 1 fig_size (tuple): Figure size of the pyplot figure. Default: (300, 300) font_size (int): Font size of texts. Default: 15 Returns: list[ndarray]: image frames """ result_dir_names = os.listdir(image_dir) assert len(result_dir_names) == 2 # Longer length has higher priority result_dir_names.reverse() images_list = [] for dir_names in result_dir_names: images_list.append(mmcv.scandir(osp.join(image_dir, dir_names))) frames = [] for paths in _generate_batch_data(zip(*images_list), nrows): fig, axes = plt.subplots(nrows=nrows, ncols=2) fig.suptitle('Good/bad case selected according ' 'to the COCO mAP of the single image') det_patch = mpatches.Patch(color='salmon', label='prediction') gt_patch = mpatches.Patch(color='royalblue', label='ground truth') # bbox_to_anchor may need to be finetuned plt.legend(handles=[det_patch, gt_patch], bbox_to_anchor=(1, -0.18), loc='lower right', borderaxespad=0.) if nrows == 1: axes = [axes] dpi = fig.get_dpi() # set fig size and margin fig.set_size_inches( (fig_size[0] * 2 + fig_size[0] // 20) / dpi, (fig_size[1] * nrows + fig_size[1] // 3) / dpi, ) fig.tight_layout() # set subplot margin plt.subplots_adjust(hspace=.05, wspace=0.05, left=0.02, right=0.98, bottom=0.02, top=0.98) for i, (path_tuple, ax_tuple) in enumerate(zip(paths, axes)): image_path_left = osp.join( osp.join(image_dir, result_dir_names[0], path_tuple[0])) image_path_right = osp.join( osp.join(image_dir, result_dir_names[1], path_tuple[1])) image_left = mmcv.imread(image_path_left) image_left = mmcv.rgb2bgr(image_left) image_right = mmcv.imread(image_path_right) image_right = mmcv.rgb2bgr(image_right) if i == 0: ax_tuple[0].set_title(result_dir_names[0], fontdict={'size': font_size}) ax_tuple[1].set_title(result_dir_names[1], fontdict={'size': font_size}) ax_tuple[0].imshow(image_left, extent=(0, *fig_size, 0), interpolation='bilinear') ax_tuple[0].axis('off') ax_tuple[1].imshow(image_right, extent=(0, *fig_size, 0), interpolation='bilinear') ax_tuple[1].axis('off') canvas = fig.canvas s, (width, height) = canvas.print_to_buffer() buffer = np.frombuffer(s, dtype='uint8') img_rgba = buffer.reshape(height, width, 4) rgb, alpha = np.split(img_rgba, [3], axis=2) img = rgb.astype('uint8') frames.append(img) return frames