This commit is contained in:
charlesbvll 2021-04-29 15:28:29 +02:00
parent d13b480f06
commit f52703b795
2 changed files with 30 additions and 30 deletions

View File

@ -198,7 +198,7 @@ def show_activities(args, image_t, output_path, annotations, dic_out):
fig_width=10,
dpi_factor=1.0) as ax:
keypoint_painter.keypoints(
ax, keypoint_sets, size=image_t.size, colors=colors, raise_hand=r_h)
ax, keypoint_sets, size=image_t.size, colors=colors)
draw_orientation(ax, uv_centers, sizes,
angles, colors, mode='front')

View File

@ -74,27 +74,7 @@ def load_image(path, scale=1.0):
return image
class KeypointPainter:
def __init__(self, *,
skeleton=None,
xy_scale=1.0, y_scale=1.0, highlight=None, highlight_invisible=False,
show_box=True, linewidth=2, markersize=3,
color_connections=False,
solid_threshold=0.5):
self.skeleton = skeleton or COCO_PERSON_SKELETON
self.xy_scale = xy_scale
self.y_scale = y_scale
self.highlight = highlight
self.highlight_invisible = highlight_invisible
self.show_box = show_box
self.linewidth = linewidth
self.markersize = markersize
self.color_connections = color_connections
self.solid_threshold = solid_threshold
self.dashed_threshold = 0.1 # Patch to still allow force complete pose (set to zero to resume original)
def _highlighted_arm(self, x, y, connection, color, lwidth, raise_hand, size=None):
def highlighted_arm(x, y, connection, color, lwidth, raise_hand, size=None):
c = color
linewidth = lwidth
@ -119,7 +99,27 @@ class KeypointPainter:
return c, linewidth
def _draw_skeleton(self, ax, x, y, v, i, *, size=None, color=None, activities=None, dic_out=None):
class KeypointPainter:
def __init__(self, *,
skeleton=None,
xy_scale=1.0, y_scale=1.0, highlight=None, highlight_invisible=False,
show_box=True, linewidth=2, markersize=3,
color_connections=False,
solid_threshold=0.5):
self.skeleton = skeleton or COCO_PERSON_SKELETON
self.xy_scale = xy_scale
self.y_scale = y_scale
self.highlight = highlight
self.highlight_invisible = highlight_invisible
self.show_box = show_box
self.linewidth = linewidth
self.markersize = markersize
self.color_connections = color_connections
self.solid_threshold = solid_threshold
self.dashed_threshold = 0.1 # Patch to still allow force complete pose (set to zero to resume original)
def _draw_skeleton(self, ax, x, y, v, *, i=0, size=None, color=None, activities=None, dic_out=None):
if not np.any(v > 0):
return
@ -130,7 +130,7 @@ class KeypointPainter:
if activities:
if 'raise_hand' in activities:
c, linewidth = self._highlighted_arm(x, y, connection, c, linewidth,
c, linewidth = highlighted_arm(x, y, connection, c, linewidth,
dic_out['raising_hand'][:][i], size=size)
if self.color_connections:
@ -231,7 +231,7 @@ class KeypointPainter:
if isinstance(color, (int, np.integer)):
color = matplotlib.cm.get_cmap('tab20')((color % 20 + 0.05) / 20)
self._draw_skeleton(ax, x, y, v, i, size=size, color=color, activities=activities, dic_out=dic_out)
self._draw_skeleton(ax, x, y, v, i=i, size=size, color=color, activities=activities, dic_out=dic_out)
score = scores[i] if scores is not None else None
if score is not None: