diff --git a/eta/core/annotations.py b/eta/core/annotations.py index 4f16ea86..b1cae463 100644 --- a/eta/core/annotations.py +++ b/eta/core/annotations.py @@ -145,8 +145,25 @@ class AnnotationConfig(Config): different names in different colors per_keypoints_label_colors: (True) whether to render keypoints with different labels in different colors - keypoints_size: (4) the size to render keypoints + keypoints_size: (6) the size to render keypoints keypoints_alpha: (0.75) the transparency of keypoints + keypoints_skeleton: (None) an optional keypoint skeleton dictionary of + the following form:: + + { + "labels": [ + "left hand" "left shoulder", "right shoulder", + "right hand", "left eye", "right eye", "mouth" + ], + "edges": [[0, 1, 2, 3], [4, 5, 6]], + } + + draw_keypoints_skeletons: (True) whether to render keypoint skeletons, + if available + keypoints_edge_linewidth (2): the linewidth, in pixels, of keypoint + skeleton edges + keypoints_edge_alpha: (0.75) the transparency of keypoint skeleton + edges show_polyline_names: (True) whether to render polyline names, if available show_polyline_labels: (True) whether to render polyline labels, if @@ -372,7 +389,7 @@ def __init__(self, d): d, "show_keypoints_attr_names", default=True ) self.show_keypoints_attr_confidences = self.parse_bool( - d, "show_keypoints_attr_confidences", default=True + d, "show_keypoints_attr_confidences", default=False ) self.per_keypoints_name_colors = self.parse_bool( d, "per_keypoints_name_colors", default=True @@ -380,10 +397,22 @@ def __init__(self, d): self.per_keypoints_label_colors = self.parse_bool( d, "per_keypoints_label_colors", default=True ) - self.keypoints_size = self.parse_number(d, "keypoints_size", default=4) + self.keypoints_size = self.parse_number(d, "keypoints_size", default=6) self.keypoints_alpha = self.parse_number( d, "keypoints_alpha", default=0.75 ) + self.keypoints_skeleton = self.parse_dict( + d, "keypoints_skeleton", default=None + ) + self.draw_keypoints_skeletons = self.parse_bool( + d, "draw_keypoints_skeletons", default=True + ) + self.keypoints_edge_linewidth = self.parse_number( + d, "keypoints_edge_linewidth", default=2 + ) + self.keypoints_edge_alpha = self.parse_number( + d, "keypoints_edge_alpha", default=0.75 + ) # POLYLINES ########################################################### @@ -400,7 +429,7 @@ def __init__(self, d): d, "show_polyline_attr_names", default=True ) self.show_polyline_attr_confidences = self.parse_bool( - d, "show_polyline_attr_confidences", default=True + d, "show_polyline_attr_confidences", default=False ) self.hide_non_filled_polyline_annos = self.parse_bool( d, "hide_non_filled_polyline_annos", default=False @@ -1087,6 +1116,15 @@ def _draw_keypoints(img, keypoints, annotation_config): ) ) gap = annotation_config.linewidth + skeleton = annotation_config.keypoints_skeleton + draw_skeletons = annotation_config.draw_keypoints_skeletons + edge_linewidth = int( + round( + annotation_config.scale_factor + * annotation_config.keypoints_edge_linewidth + ) + ) + edge_alpha = annotation_config.keypoints_edge_alpha colormap = annotation_config.colormap @@ -1128,24 +1166,57 @@ def _draw_keypoints(img, keypoints, annotation_config): color = _parse_hex_color(colormap.get_color(title_hash)) # Render coordinates for image - points = keypoints.coords_in(img=img) + frame_size = np.array(etai.to_frame_size(img=img)) + points = np.array(keypoints.points, dtype=float) # converts None -> NaN + points = np.nan_to_num(points, nan=-1, posinf=-1, neginf=-1) + points = np.rint(frame_size * points).astype(np.int32) + _points = points[(points >= 0).all(axis=1)] # omits nans + + if len(_points) == 0: + return img # # Draw keypoints # overlay = img.copy() + img_anno = img - for x, y in points: + # Draw skeleton + if draw_skeletons and skeleton is not None: + edges = [] + for inds in skeleton["edges"]: + segment = [] + for pt in points[inds]: + if pt[0] < 0 or pt[1] < 0: + if segment: + edges.append(np.array(segment)) + segment = [] + else: + segment.append(pt) + + if segment: + edges.append(np.array(segment)) + + if edges: + overlay = cv2.polylines( + overlay, edges, False, color, thickness=edge_linewidth + ) + img_anno = cv2.addWeighted( + overlay, edge_alpha, img_anno, 1 - edge_alpha, 0 + ) + + # Draw points + for x, y in _points: overlay = cv2.circle(overlay, (x, y), size, color, thickness=-1) - img_anno = cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0) + img_anno = cv2.addWeighted(overlay, alpha, img_anno, 1 - alpha, 0) # # Draw title string # - tcx, tcy = points[int(round(0.5 * len(points)))] + tcx, tcy = _points[0] tcx += size tcy += size ttlx, ttly, tw, th = _get_panel_coords( diff --git a/eta/core/keypoints.py b/eta/core/keypoints.py index b96e9efc..2b8f723f 100644 --- a/eta/core/keypoints.py +++ b/eta/core/keypoints.py @@ -37,9 +37,9 @@ class Keypoints(etal.Labels): name: (optional) the name for the keypoints, e.g., ``ground_truth`` or the name of the model that produced it label: (optional) keypoints label - confidence: (optional) a confidence for the keypoints, in ``[0, 1]`` index: (optional) an index assigned to the keypoints points: a list of ``(x, y)`` keypoints in ``[0, 1] x [0, 1]`` + confidence: (optional) a list of per-point confidences in ``[0, 1]`` attrs: (optional) an :class:`eta.core.data.AttributeContainer` of attributes for the keypoints tags: (optional) a list of tag strings @@ -48,9 +48,9 @@ class Keypoints(etal.Labels): name (None): a name for the keypoints, e.g., ``ground_truth`` or the name of the model that produced it label (None): a label for the keypoints - confidence (None): a confidence for the keypoints, in ``[0, 1]`` index (None): an integer index assigned to the keypoints points (None): a list of ``(x, y)`` keypoints in ``[0, 1] x [0, 1]`` + confidence (None): a list of per-point confidences in ``[0, 1]`` attrs (None): an :class:`eta.core.data.AttributeContainer` of attributes for the keypoints tags (None): a list of tag strings @@ -60,18 +60,18 @@ def __init__( self, name=None, label=None, - confidence=None, index=None, points=None, + confidence=None, attrs=None, tags=None, ): self.type = etau.get_class_name(self) self.name = name self.label = label - self.confidence = confidence self.index = index self.points = points or [] + self.confidence = confidence self.attrs = attrs or etad.AttributeContainer() self.tags = tags or [] @@ -92,7 +92,7 @@ def has_label(self): @property def has_confidence(self): - """Whether the keypoints has a ``confidence``.""" + """Whether the keypoints have ``confidence``.""" return self.confidence is not None @property @@ -313,9 +313,9 @@ def from_dict(cls, d): return cls( name=name, label=label, - confidence=confidence, index=index, points=points, + confidence=confidence, attrs=attrs, tags=tags, ) @@ -362,16 +362,6 @@ def clear_indexes(self): for keypoints in self: keypoints.clear_index() - def sort_by_confidence(self, reverse=False): - """Sorts the :class:`Keypoints` instances by confidence. - - Keypoints whose confidence is ``None`` are always put last. - - Args: - reverse (False): whether to sort in descending order - """ - self.sort_by("confidence", reverse=reverse) - def sort_by_index(self, reverse=False): """Sorts the :class:`Keypoints` instances by index.