Merge pull request #7 from DagsHub/new/ellipsis

Feature: Ellipse annotations

dagshub_annotation_converter/converters/cvat.py

@@ -1,6 +1,6 @@
 import logging
 from os import PathLike
-from typing import Sequence, List, Dict
+from typing import Sequence, List, Dict, Union
 from zipfile import ZipFile
 
 import lxml.etree
@@ -38,12 +38,12 @@ def load_cvat_from_xml_string(
     return annotations
 
 
-def load_cvat_from_xml_file(xml_file: PathLike) -> Dict[str, Sequence[IRImageAnnotationBase]]:
+def load_cvat_from_xml_file(xml_file: Union[str, PathLike]) -> Dict[str, Sequence[IRImageAnnotationBase]]:
     with open(xml_file, "rb") as f:
         return load_cvat_from_xml_string(f.read())
 
 
-def load_cvat_from_zip(zip_path: PathLike) -> Dict[str, Sequence[IRImageAnnotationBase]]:
+def load_cvat_from_zip(zip_path: Union[str, PathLike]) -> Dict[str, Sequence[IRImageAnnotationBase]]:
     with ZipFile(zip_path) as proj_zip:
         with proj_zip.open("annotations.xml") as f:
             return load_cvat_from_xml_string(f.read())

dagshub_annotation_converter/formats/cvat/__init__.py

@@ -3,6 +3,7 @@
 from lxml.etree import ElementBase
 
 from .box import parse_box
+from .ellipse import parse_ellipse
 from .polygon import parse_polygon
 from .points import parse_points
 from .skeleton import parse_skeleton
@@ -15,4 +16,5 @@
     "polygon": parse_polygon,
     "points": parse_points,
     "skeleton": parse_skeleton,
+    "ellipse": parse_ellipse,
 }

dagshub_annotation_converter/formats/cvat/box.py

@@ -58,6 +58,7 @@ def parse_box(elem: ElementBase, containing_image: ElementBase) -> IRBBoxImageAn
         left=left,
         width=width,
         height=height,
+        rotation=rotation,
         image_width=image_info.width,
         image_height=image_info.height,
         filename=image_info.name,

dagshub_annotation_converter/formats/cvat/ellipse.py

@@ -0,0 +1,28 @@
+from lxml.etree import ElementBase
+
+from dagshub_annotation_converter.formats.cvat.context import parse_image_tag
+from dagshub_annotation_converter.ir.image import IREllipseImageAnnotation, CoordinateStyle
+
+
+def parse_ellipse(elem: ElementBase, containing_image: ElementBase) -> IREllipseImageAnnotation:
+    center_x = float(elem.attrib["cx"])
+    center_y = float(elem.attrib["cy"])
+    radius_x = float(elem.attrib["rx"])
+    radius_y = float(elem.attrib["ry"])
+
+    rotation = float(elem.attrib.get("rotation", 0.0))
+
+    image_info = parse_image_tag(containing_image)
+
+    return IREllipseImageAnnotation(
+        categories={str(elem.attrib["label"]): 1.0},
+        center_x=round(center_x),
+        center_y=round(center_y),
+        radius_x=radius_x,
+        radius_y=radius_y,
+        rotation=rotation,
+        image_width=image_info.width,
+        image_height=image_info.height,
+        filename=image_info.name,
+        coordinate_style=CoordinateStyle.DENORMALIZED,
+    )

dagshub_annotation_converter/formats/label_studio/ellipselabels.py

@@ -0,0 +1,56 @@
+from typing import List, Sequence
+
+from dagshub_annotation_converter.formats.label_studio.base import ImageAnnotationResultABC
+from dagshub_annotation_converter.ir.image import IRImageAnnotationBase, IREllipseImageAnnotation, CoordinateStyle
+from dagshub_annotation_converter.util.pydantic_util import ParentModel
+
+
+class EllipseLabelsAnnotationsValue(ParentModel):
+    x: float
+    y: float
+    radiusX: float
+    radiusY: float
+    rotation: float = 0.0
+    ellipselabels: List[str]
+
+
+class EllipseLabelsAnnotation(ImageAnnotationResultABC):
+    value: EllipseLabelsAnnotationsValue
+    type: str = "ellipselabels"
+
+    def to_ir_annotation(self) -> Sequence[IRImageAnnotationBase]:
+        res = IREllipseImageAnnotation(
+            categories={self.value.ellipselabels[0]: 1.0},
+            coordinate_style=CoordinateStyle.NORMALIZED,
+            center_x=self.value.x / 100,
+            center_y=self.value.y / 100,
+            radius_x=self.value.radiusX / 100,
+            radius_y=self.value.radiusY / 100,
+            rotation=self.value.rotation,
+            image_width=self.original_width,
+            image_height=self.original_height,
+        )
+        res.imported_id = self.id
+        return [res]
+
+    @staticmethod
+    def from_ir_annotation(ir_annotation: IRImageAnnotationBase) -> Sequence["ImageAnnotationResultABC"]:
+        assert isinstance(ir_annotation, IREllipseImageAnnotation)
+
+        ir_annotation = ir_annotation.normalized()
+        category = ir_annotation.ensure_has_one_category()
+
+        return [
+            EllipseLabelsAnnotation(
+                original_width=ir_annotation.image_width,
+                original_height=ir_annotation.image_height,
+                value=EllipseLabelsAnnotationsValue(
+                    x=ir_annotation.center_x * 100,
+                    y=ir_annotation.center_y * 100,
+                    radiusX=ir_annotation.radius_x * 100,
+                    radiusY=ir_annotation.radius_y * 100,
+                    rotation=ir_annotation.rotation,
+                    ellipselabels=[category],
+                ),
+            )
+        ]

dagshub_annotation_converter/formats/label_studio/task.py

@@ -8,6 +8,7 @@
 from pydantic import SerializeAsAny, Field, BeforeValidator
 
 from dagshub_annotation_converter.formats.label_studio.base import AnnotationResultABC, ImageAnnotationResultABC
+from dagshub_annotation_converter.formats.label_studio.ellipselabels import EllipseLabelsAnnotation
 from dagshub_annotation_converter.formats.label_studio.keypointlabels import KeyPointLabelsAnnotation
 from dagshub_annotation_converter.formats.label_studio.polygonlabels import PolygonLabelsAnnotation
 from dagshub_annotation_converter.formats.label_studio.rectanglelabels import RectangleLabelsAnnotation
@@ -18,19 +19,22 @@
     CoordinateStyle,
     IRPosePoint,
     IRSegmentationImageAnnotation,
+    IREllipseImageAnnotation,
 )
 from dagshub_annotation_converter.util.pydantic_util import ParentModel
 
 task_lookup: Dict[str, Type[AnnotationResultABC]] = {
     "polygonlabels": PolygonLabelsAnnotation,
     "rectanglelabels": RectangleLabelsAnnotation,
     "keypointlabels": KeyPointLabelsAnnotation,
+    "ellipselabels": EllipseLabelsAnnotation,
 }
 
 ir_annotation_lookup: Dict[Type[IRImageAnnotationBase], Type[ImageAnnotationResultABC]] = {
     IRPoseImageAnnotation: KeyPointLabelsAnnotation,
     IRBBoxImageAnnotation: RectangleLabelsAnnotation,
     IRSegmentationImageAnnotation: PolygonLabelsAnnotation,
+    IREllipseImageAnnotation: EllipseLabelsAnnotation,
 }
 
 logger = logging.getLogger(__name__)

dagshub_annotation_converter/ir/image/__init__.py

@@ -3,6 +3,7 @@
 from .annotations.pose import IRPoseImageAnnotation, IRPosePoint
 from .annotations.segmentation import IRSegmentationImageAnnotation, IRSegmentationPoint
 from .annotations.bbox import IRBBoxImageAnnotation
+from .annotations.ellipse import IREllipseImageAnnotation
 
 __all__ = [
     "CoordinateStyle",
@@ -12,4 +13,5 @@
     "IRSegmentationImageAnnotation",
     "IRSegmentationPoint",
     "IRBBoxImageAnnotation",
+    "IREllipseImageAnnotation",
 ]

dagshub_annotation_converter/ir/image/annotations/ellipse.py

@@ -0,0 +1,22 @@
+from dagshub_annotation_converter.ir.image import IRImageAnnotationBase
+
+
+class IREllipseImageAnnotation(IRImageAnnotationBase):
+    center_x: float
+    center_y: float
+    radius_x: float
+    radius_y: float
+    rotation: float = 0.0
+    """Rotation in degrees (pivot point - top-left)"""
+
+    def _normalize(self):
+        self.center_x = self.center_x / self.image_width
+        self.center_y = self.center_y / self.image_height
+        self.radius_x = self.radius_x / self.image_width
+        self.radius_y = self.radius_y / self.image_height
+
+    def _denormalize(self):
+        self.center_x = self.center_x * self.image_width
+        self.center_y = self.center_y * self.image_height
+        self.radius_x = self.radius_x * self.image_width
+        self.radius_y = self.radius_y * self.image_height

tests/cvat/test_parsers.py

@@ -7,15 +7,15 @@
     parse_box,
     parse_polygon,
     parse_points,
-    parse_skeleton,
+    parse_skeleton, parse_ellipse,
 )
 from dagshub_annotation_converter.ir.image import (
     CoordinateStyle,
     IRBBoxImageAnnotation,
     IRSegmentationImageAnnotation,
     IRSegmentationPoint,
     IRPosePoint,
-    IRPoseImageAnnotation,
+    IRPoseImageAnnotation, IREllipseImageAnnotation,
 )
 
 
@@ -159,3 +159,29 @@ def test_skeleton():
     )
 
     assert expected == actual
+
+
+def test_ellipse():
+    data = """
+    <ellipse label="circle1" source="manual" occluded="0" cx="392.23" cy="322.84" rx="205.06" ry="202.84" z_order="0">
+    </ellipse> 
+    """
+
+    image, annotation = to_xml(data)
+
+    actual = parse_ellipse(annotation, image)
+
+    expected = IREllipseImageAnnotation(
+        filename="000.png",
+        categories={"circle1": 1.0},
+        image_width=1920,
+        image_height=1200,
+        coordinate_style=CoordinateStyle.DENORMALIZED,
+        center_x=392,
+        center_y=323,
+        radius_x=205.06,
+        radius_y=202.84,
+        rotation=0.0,
+    )
+
+    assert actual == expected

tests/label_studio/test_ellipse.py

@@ -0,0 +1,85 @@
+from typing import Dict
+
+import pytest
+
+from dagshub_annotation_converter.formats.label_studio.ellipselabels import EllipseLabelsAnnotation
+from dagshub_annotation_converter.formats.label_studio.task import LabelStudioTask, parse_ls_task
+from dagshub_annotation_converter.ir.image import IREllipseImageAnnotation, CoordinateStyle
+from tests.label_studio.common import generate_annotation, generate_task
+
+
+@pytest.fixture
+def ellipse_annotation() -> Dict:
+    annotation = {
+        "x": 10,
+        "y": 20,
+        "radiusX": 30,
+        "radiusY": 40,
+        "ellipselabels": ["dog"],
+    }
+
+    return generate_annotation(annotation, "ellipselabels", "deadbeef")
+
+
+@pytest.fixture
+def ellipse_task(ellipse_annotation) -> str:
+    return generate_task([ellipse_annotation])
+
+
+@pytest.fixture
+def parsed_ellipse_task(ellipse_task) -> LabelStudioTask:
+    return parse_ls_task(ellipse_task)
+
+
+def test_ellipse_parsing(parsed_ellipse_task):
+    actual = parsed_ellipse_task
+    assert len(actual.annotations) == 1
+    assert len(actual.annotations[0].result) == 1
+
+    ann = actual.annotations[0].result[0]
+    assert isinstance(ann, EllipseLabelsAnnotation)
+    assert ann.value.x == 10
+    assert ann.value.y == 20
+    assert ann.value.radiusX == 30
+    assert ann.value.radiusY == 40
+
+
+def test_ellipse_ir(parsed_ellipse_task):
+    actual = parsed_ellipse_task.annotations[0].result[0].to_ir_annotation()
+
+    assert len(actual) == 1
+    ann = actual[0]
+    assert isinstance(ann, IREllipseImageAnnotation)
+
+    assert ann.center_x == 0.1
+    assert ann.center_y == 0.2
+    assert ann.radius_x == 0.3
+    assert ann.radius_y == 0.4
+
+
+def test_ir_ellipse_addition():
+    task = LabelStudioTask()
+    ellipse = IREllipseImageAnnotation(
+        categories={"dog": 1.0},
+        center_x=0.1,
+        center_y=0.2,
+        radius_x=0.3,
+        radius_y=0.4,
+        rotation=60.0,
+        image_width=100,
+        image_height=100,
+        coordinate_style=CoordinateStyle.NORMALIZED,
+    )
+
+    task.add_ir_annotation(ellipse)
+
+    assert len(task.annotations) == 1
+    assert len(task.annotations[0].result) == 1
+
+    ann = task.annotations[0].result[0]
+    assert isinstance(ann, EllipseLabelsAnnotation)
+    assert ann.value.x == 10
+    assert ann.value.y == 20
+    assert ann.value.radiusX == 30
+    assert ann.value.radiusY == 40
+    assert ann.value.rotation == 60.0