Merge pull request #9 from aoxolotl/seg/fromMask

Create segmentations from binary mask

Author: ninthreezy

datatorch/api/entity/sources/image/segmentations.py

@@ -1,9 +1,15 @@
 import numpy as np
 
-from typing import List
+from typing import List, Optional
+from imantics import Polygons, Mask
+import shapely.ops
+from shapely import geometry
 
 from .typings import Segment
 from ..source import Source
+from ....scripts.utils.simplify import simplify_points
+from ....entity.annotation import Annotation
+from datatorch.api import ApiClient
 
 
 __all__ = "Segmentations"
@@ -14,8 +20,66 @@ class Segmentations(Source):
     type: str = "PaperSegmentations"
     path_data: List[Segment]
 
-    def from_mask(self, mask: np.array):
-        pass
+    def __init__(self, path_data: Optional[List[Segment]] = None):
+        super().__init__()
+        self.path_data = path_data or []
+
+    def from_mask(self, mask: np.array, simplify: int = 0):
+        # convert mask to polygons
+        polygons = Mask(mask).polygons().points
+        # TODO: handle shifted polygons in cropped case
+        polygons = [polygon.tolist() for polygon in polygons]
+
+        # polygons -> path_data
+        # simplify
+        self.path_data = polygons
+        if simplify:
+            self.path_data = [
+                simplify_points(polygon, tolerance=simplify, highestQuality=False)
+                for polygon in polygons
+            ]
+
+        # filter polygons
+        self.path_data = list(filter(lambda x: len(x) > 2, self.path_data))
+
+    def combine_segmentations(self, annotation):
+        if len(self.path_data) == 0:
+            raise ValueError("No path data to combine")
+
+        self.annotation_id = annotation.id
+        existing_segmentation = next(
+            x
+            for x in annotation.get("sources")
+            if x.get("type") == "PaperSegmentations"
+        )
+        poly_1 = [geometry.Polygon(points) for points in self.path_data]
+        poly_2 = [geometry.Polygon(points) for points in existing_segmentation]
+
+        multi = shapely.ops.unary_union(poly_1 + poly_2)
+
+        if isinstance(multi, geometry.Polygon):
+            self.path_data.append(list(multi.exterior.coords[:-1]))
+
+        if isinstance(multi, geometry.MultiPolygon):
+            for polygon in multi:
+                self.path_data.append(list(polygon.exterior.coords[:-1]))
+
+        print(
+            f"Updated segmentation for annotation {annotation.id}",
+            flush=True,
+        )
+
+    def create_new_annotation(self, label_id: str, file_id: str):
+        print("Creating new annotation")
+        new_annotation = Annotation()
+        new_annotation.label_id = label_id
+        new_annotation.file_id = file_id
+        new_annotation.create(ApiClient())
+        annotation_id = new_annotation.id
+
+        self.annotation_id = annotation_id
+        self.create(ApiClient())
+        print("Segmentation created")
 
     def to_mask(self) -> np.array:
         pass

setup.py

@@ -12,7 +12,7 @@
     "gql==3.4.0",
     "websockets==10.4",
     "websocket-client",
-    "requests==2.27.1",
+    "requests==2.31.0",
     "typing_extensions>=4.1.0",
     "psutil~=5.9.4",
     "aiodocker~=0.19.0",
@@ -21,6 +21,8 @@
     "aiostream~=0.4.0",
     "markupsafe==2.0.1",
     "requests_toolbelt==0.10.1",
+    "imantics==0.1.12",
+    "shapely==2.0.1",
     "tqdm~=4.65.0",
     "urllib3==1.26.15",
 ]
@@ -29,7 +31,7 @@
 
 setup(
     name="datatorch",
-    version="0.4.7.1",
+    version="0.4.7.2",
     description="A CLI and library for interacting with DataTorch.",
     author="DataTorch",
     author_email="support@datatorch.io",