Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Enhance] Inferencer supports visualization black background #2458

Merged
merged 1 commit into from
Jun 14, 2023
Merged

[Enhance] Inferencer supports visualization black background #2458

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 35 additions & 35 deletions configs/_base_/datasets/coco_openpose.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,77 +12,77 @@
),
keypoint_info={
0:
dict(name='nose', id=0, color=[255, 0, 85], type='upper', swap=''),
dict(name='nose', id=0, color=[255, 0, 0], type='upper', swap=''),
1:
dict(name='neck', id=1, color=[255, 0, 0], type='upper', swap=''),
dict(name='neck', id=1, color=[255, 85, 0], type='upper', swap=''),
2:
dict(
name='right_shoulder',
id=2,
color=[255, 85, 0],
color=[255, 170, 0],
type='upper',
swap='left_shoulder'),
3:
dict(
name='right_elbow',
id=3,
color=[255, 170, 0],
color=[255, 255, 0],
type='upper',
swap='left_elbow'),
4:
dict(
name='right_wrist',
id=4,
color=[255, 255, 0],
color=[170, 255, 0],
type='upper',
swap='left_wrist'),
5:
dict(
name='left_shoulder',
id=5,
color=[170, 255, 0],
color=[85, 255, 0],
type='upper',
swap='right_shoulder'),
6:
dict(
name='left_elbow',
id=6,
color=[85, 255, 0],
color=[0, 255, 0],
type='upper',
swap='right_elbow'),
7:
dict(
name='left_wrist',
id=7,
color=[0, 255, 0],
color=[0, 255, 85],
type='upper',
swap='right_wrist'),
8:
dict(
name='right_hip',
id=8,
color=[255, 0, 170],
color=[0, 255, 170],
type='lower',
swap='left_hip'),
9:
dict(
name='right_knee',
id=9,
color=[255, 0, 255],
color=[0, 255, 255],
type='lower',
swap='left_knee'),
10:
dict(
name='right_ankle',
id=10,
color=[170, 0, 255],
color=[0, 170, 255],
type='lower',
swap='left_ankle'),
11:
dict(
name='left_hip',
id=11,
color=[85, 255, 0],
color=[0, 85, 255],
type='lower',
swap='right_hip'),
12:
Expand All @@ -96,59 +96,59 @@
dict(
name='left_ankle',
id=13,
color=[0, 85, 255],
color=[85, 0, 255],
type='lower',
swap='right_ankle'),
14:
dict(
name='right_eye',
id=14,
color=[0, 255, 170],
color=[170, 0, 255],
type='upper',
swap='left_eye'),
15:
dict(
name='left_eye',
id=15,
color=[0, 255, 255],
color=[255, 0, 255],
type='upper',
swap='right_eye'),
16:
dict(
name='right_ear',
id=16,
color=[0, 170, 255],
color=[255, 0, 170],
type='upper',
swap='left_ear'),
17:
dict(
name='left_ear',
id=17,
color=[0, 170, 255],
color=[255, 0, 85],
type='upper',
swap='right_ear'),
},
skeleton_info={
0: dict(link=('neck', 'right_shoulder'), id=0, color=[255, 0, 85]),
1: dict(link=('neck', 'left_shoulder'), id=1, color=[255, 0, 0]),
2:
dict(link=('right_shoulder', 'right_elbow'), id=2, color=[255, 85, 0]),
0: dict(link=('neck', 'right_shoulder'), id=0, color=[255, 0, 0]),
1: dict(link=('neck', 'left_shoulder'), id=1, color=[255, 85, 0]),
2: dict(
link=('right_shoulder', 'right_elbow'), id=2, color=[255, 170, 0]),
3:
dict(link=('right_elbow', 'right_wrist'), id=3, color=[255, 170, 0]),
dict(link=('right_elbow', 'right_wrist'), id=3, color=[255, 255, 0]),
4:
dict(link=('left_shoulder', 'left_elbow'), id=4, color=[255, 255, 0]),
5: dict(link=('left_elbow', 'left_wrist'), id=5, color=[170, 255, 0]),
6: dict(link=('neck', 'right_hip'), id=6, color=[85, 255, 0]),
7: dict(link=('right_hip', 'right_knee'), id=7, color=[0, 255, 0]),
8: dict(link=('right_knee', 'right_ankle'), id=8, color=[0, 255, 85]),
9: dict(link=('neck', 'left_hip'), id=9, color=[0, 255, 170]),
10: dict(link=('left_hip', 'left_knee'), id=10, color=[0, 255, 225]),
11: dict(link=('left_knee', 'left_ankle'), id=11, color=[0, 170, 255]),
12: dict(link=('neck', 'nose'), id=12, color=[0, 85, 255]),
13: dict(link=('nose', 'right_eye'), id=13, color=[0, 0, 255]),
14: dict(link=('right_eye', 'right_ear'), id=14, color=[255, 0, 170]),
15: dict(link=('nose', 'left_eye'), id=15, color=[170, 0, 255]),
16: dict(link=('left_eye', 'left_ear'), id=16, color=[255, 0, 255]),
dict(link=('left_shoulder', 'left_elbow'), id=4, color=[170, 255, 0]),
5: dict(link=('left_elbow', 'left_wrist'), id=5, color=[85, 255, 0]),
6: dict(link=('neck', 'right_hip'), id=6, color=[0, 255, 0]),
7: dict(link=('right_hip', 'right_knee'), id=7, color=[0, 255, 85]),
8: dict(link=('right_knee', 'right_ankle'), id=8, color=[0, 255, 170]),
9: dict(link=('neck', 'left_hip'), id=9, color=[0, 255, 225]),
10: dict(link=('left_hip', 'left_knee'), id=10, color=[0, 170, 255]),
11: dict(link=('left_knee', 'left_ankle'), id=11, color=[0, 85, 255]),
12: dict(link=('neck', 'nose'), id=12, color=[0, 0, 255]),
13: dict(link=('nose', 'right_eye'), id=13, color=[255, 0, 170]),
14: dict(link=('right_eye', 'right_ear'), id=14, color=[170, 0, 255]),
15: dict(link=('nose', 'left_eye'), id=15, color=[255, 0, 255]),
16: dict(link=('left_eye', 'left_ear'), id=16, color=[255, 0, 170]),
},
joint_weights=[1.] * 18,
sigmas=[
Expand Down
4 changes: 4 additions & 0 deletions demo/inferencer_demo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,10 @@ def parse_args():
type=str,
choices=['mmpose', 'openpose'],
help='Skeleton style selection')
parser.add_argument(
'--black-background',
action='store_true',
help='Plot predictions on a black image')
parser.add_argument(
'--vis-out-dir',
type=str,
Expand Down
15 changes: 7 additions & 8 deletions mmpose/apis/inferencers/base_mmpose_inferencer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,14 +39,8 @@ class BaseMMPoseInferencer(BaseInferencer):
preprocess_kwargs: set = {'bbox_thr', 'nms_thr', 'bboxes'}
forward_kwargs: set = set()
visualize_kwargs: set = {
'return_vis',
'show',
'wait_time',
'draw_bbox',
'radius',
'thickness',
'kpt_thr',
'vis_out_dir',
'return_vis', 'show', 'wait_time', 'draw_bbox', 'radius', 'thickness',
'kpt_thr', 'vis_out_dir', 'black_background'
}
postprocess_kwargs: set = {'pred_out_dir'}

Expand Down Expand Up @@ -271,6 +265,7 @@ def visualize(self,
kpt_thr: float = 0.3,
vis_out_dir: str = '',
window_name: str = '',
black_background: bool = False,
**kwargs) -> List[np.ndarray]:
"""Visualize predictions.

Expand All @@ -291,6 +286,8 @@ def visualize(self,
results w/o predictions. If left as empty, no file will
be saved. Defaults to ''.
window_name (str, optional): Title of display window.
black_background (bool, optional): Whether to plot keypoints on a
black image instead of the input image. Defaults to False.

Returns:
List[np.ndarray]: Visualization results.
Expand All @@ -315,6 +312,8 @@ def visualize(self,
else:
raise ValueError('Unsupported input type: '
f'{type(single_input)}')
if black_background:
img = img * 0

img_name = os.path.basename(pred.metainfo['img_path'])
window_name = window_name if window_name else img_name
Expand Down
3 changes: 2 additions & 1 deletion mmpose/apis/inferencers/mmpose_inferencer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,8 @@ class MMPoseInferencer(BaseMMPoseInferencer):
forward_kwargs: set = {'rebase_keypoint_height'}
visualize_kwargs: set = {
'return_vis', 'show', 'wait_time', 'draw_bbox', 'radius', 'thickness',
'kpt_thr', 'vis_out_dir', 'skeleton_style', 'draw_heatmap'
'kpt_thr', 'vis_out_dir', 'skeleton_style', 'draw_heatmap',
'black_background'
}
postprocess_kwargs: set = {'pred_out_dir'}

Expand Down
2 changes: 1 addition & 1 deletion mmpose/apis/inferencers/pose2d_inferencer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ class Pose2DInferencer(BaseMMPoseInferencer):
'vis_out_dir',
'skeleton_style',
'draw_heatmap',
'black_background',
}
postprocess_kwargs: set = {'pred_out_dir'}

Expand Down Expand Up @@ -152,7 +153,6 @@ def update_model_visualizer_settings(self,
if skeleton_style == 'openpose':
self.visualizer.set_dataset_meta(self.model.dataset_meta,
skeleton_style)
self.visualizer.backend = 'matplotlib'

def preprocess_single(self,
input: InputType,
Expand Down
10 changes: 5 additions & 5 deletions mmpose/visualization/local_visualizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ class PoseLocalVisualizer(OpencvBackendVisualizer):
radius (int, float): The radius of keypoints. Defaults to 4
show_keypoint_weight (bool): Whether to adjust the transparency
of keypoints according to their score. Defaults to ``False``
alpha (int, float): The transparency of bboxes. Defaults to ``0.8``
alpha (int, float): The transparency of bboxes. Defaults to ``1.0``

Examples:
>>> import numpy as np
Expand Down Expand Up @@ -116,7 +116,7 @@ def __init__(self,
radius: Union[int, float] = 3,
show_keypoint_weight: bool = False,
backend: str = 'opencv',
alpha: float = 0.8):
alpha: float = 1.0):
super().__init__(
name=name,
image=image,
Expand Down Expand Up @@ -345,13 +345,13 @@ def _draw_instances_kpts(self,
mX = np.mean(X)
mY = np.mean(Y)
length = ((Y[0] - Y[1])**2 + (X[0] - X[1])**2)**0.5
transparency = 0.6
angle = math.degrees(
math.atan2(Y[0] - Y[1], X[0] - X[1]))
stickwidth = 2
polygons = cv2.ellipse2Poly(
(int(mX), int(mY)),
(int(length / 2), int(stickwidth)), int(angle),
0, 360, 1)
(int(length / 2), int(self.line_width)),
int(angle), 0, 360, 1)

self.draw_polygons(
polygons,
Expand Down
34 changes: 27 additions & 7 deletions mmpose/visualization/opencv_backend_visualizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ class OpencvBackendVisualizer(Visualizer):
Defaults to empty dict.
backend (str): Backend used to draw elements on the image and display
the image. Defaults to 'matplotlib'.
alpha (int, float): The transparency of bboxes. Defaults to ``1.0``
"""

def __init__(self,
Expand Down Expand Up @@ -87,6 +88,7 @@ def draw_circles(self,
radius: Union[np.ndarray, torch.Tensor],
face_colors: Union[str, tuple, List[str],
List[tuple]] = 'none',
alpha: float = 1.0,
**kwargs) -> 'Visualizer':
"""Draw single or multiple circles.

Expand Down Expand Up @@ -123,13 +125,22 @@ def draw_circles(self,
center=center,
radius=radius,
face_colors=face_colors,
alpha=alpha,
**kwargs)
elif self.backend == 'opencv':
if isinstance(face_colors, str):
face_colors = mmcv.color_val(face_colors)
self._image = cv2.circle(self._image,
(int(center[0]), int(center[1])),
int(radius), face_colors, -1)

if alpha == 1.0:
self._image = cv2.circle(self._image,
(int(center[0]), int(center[1])),
int(radius), face_colors, -1)
else:
img = cv2.circle(self._image.copy(),
(int(center[0]), int(center[1])), int(radius),
face_colors, -1)
self._image = cv2.addWeighted(self._image, 1 - alpha, img,
alpha, 0)
else:
raise ValueError(f'got unsupported backend {self.backend}')

Expand Down Expand Up @@ -362,6 +373,7 @@ def draw_polygons(self,
List[Union[np.ndarray, torch.Tensor]]],
edge_colors: Union[str, tuple, List[str],
List[tuple]] = 'g',
alpha: float = 1.0,
**kwargs) -> 'Visualizer':
"""Draw single or multiple bboxes.

Expand Down Expand Up @@ -394,12 +406,20 @@ def draw_polygons(self,
"""
if self.backend == 'matplotlib':
super().draw_polygons(
polygons=polygons, edge_colors=edge_colors, **kwargs)
polygons=polygons,
edge_colors=edge_colors,
alpha=alpha,
**kwargs)

elif self.backend == 'opencv':

self._image = cv2.fillConvexPoly(self._image, polygons,
edge_colors)
if alpha == 1.0:
self._image = cv2.fillConvexPoly(self._image, polygons,
edge_colors)
else:
img = cv2.fillConvexPoly(self._image.copy(), polygons,
edge_colors)
self._image = cv2.addWeighted(self._image, 1 - alpha, img,
alpha, 0)
else:
raise ValueError(f'got unsupported backend {self.backend}')

Expand Down