feat: key point label transformer (#642)

encord-team · Sep 12, 2023 · 36be608 · 36be608
1 parent bbc50dc
commit 36be608
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 5 deletions.
diff --git a/src/encord_active/frontend/src/components/explorer/Explorer.tsx b/src/encord_active/frontend/src/components/explorer/Explorer.tsx
@@ -1164,15 +1164,15 @@ const ImageWithPolygons = ({
                   <g key={index}>
                     <circle
                       key={index + "_inner"}
-                      cx={points[0].x}
-                      cy={points[0].y}
+                      cx={points[0].x * width}
+                      cy={points[0].y * height}
                       r="5px"
                       fill={color}
                     />
                     <circle
                       key={index + "_outer"}
-                      cx={points[0].x}
-                      cy={points[0].y}
+                      cx={points[0].x * width}
+                      cy={points[0].y * height}
                       r="7px"
                       fill="none"
                       stroke={color}

diff --git a/src/encord_active/lib/labels/label_transformer.py b/src/encord_active/lib/labels/label_transformer.py
@@ -37,6 +37,13 @@ def __str__(self) -> str:
         return f"BounndingBoxLabel({self.class_}, {self.bounding_box})"
 
 
+class KeypointLabel(Label):
+    point: Point
+
+    def __str__(self) -> str:
+        return f"KeypointLabel({self.class_})"
+
+
 @dataclass
 class PolygonLabel:
     class_: str
@@ -50,7 +57,7 @@ def __str__(self) -> str:
         return f"PolygonLabel({self.class_}, [{n}, 2])"
 
 
-LabelOptions = Union[BoundingBoxLabel, ClassificationLabel, PolygonLabel]
+LabelOptions = Union[BoundingBoxLabel, ClassificationLabel, PolygonLabel, KeypointLabel]
 
 
 class DataLabel(NamedTuple):
@@ -191,6 +198,13 @@ def _add_bounding_box_label(self, label: BoundingBoxLabel, data_unit: dict):
         bbox_obj = make_object_dict(ont_obj, bbox_dict)
         data_unit.setdefault("labels", {}).setdefault("objects", []).append(bbox_obj)
 
+    def _add_keypoint_label(self, label: KeypointLabel, data_unit: dict):
+        shape = Shape.POINT
+
+        ont_obj = self._get_object_by_name(label.class_, shape=shape, create_when_missing=True)
+        bbox_obj = make_object_dict(ont_obj, label.point)
+        data_unit.setdefault("labels", {}).setdefault("objects", []).append(bbox_obj)
+
     def _add_polygon_label(self, label: PolygonLabel, data_unit: dict):
         ont_obj = self._get_object_by_name(label.class_, shape=Shape.POLYGON, create_when_missing=True)
         points = [Point(*r) for r in label.polygon]
@@ -204,6 +218,8 @@ def _add_label(self, label: LabelOptions, label_row: dict, data_unit: dict):
             self._add_bounding_box_label(label, data_unit)
         elif isinstance(label, PolygonLabel):
             self._add_polygon_label(label, data_unit)
+        elif isinstance(label, KeypointLabel):
+            self._add_keypoint_label(label, data_unit)
 
     def add_labels(self, label_paths: List[Path], data_paths: List[Path]):
         if not self.label_transformer: