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
