Merge pull request #55 from fastlabel/hotfix/fix-segmentation-bug

fix segmentation bug

Author: eisuke-ueta

fastlabel/__init__.py

@@ -3,7 +3,9 @@
 import json
 from typing import List
 from logging import getLogger
-from PIL import Image, ImageDraw
+from PIL import Image
+import cv2
+import numpy as np
 
 import xmltodict
 
@@ -653,35 +655,37 @@ def export_semantic_segmentation(self, tasks: list, output_dir: str = os.path.jo
 
     def __export_index_color_image(self, task: list, output_dir: str, pallete: List[int], is_instance_segmentation: bool = True, classes: list = []) -> None:
         image = Image.new("RGB", (task["width"], task["height"]), 0)
-        image = image.convert('P')
-        image.putpalette(pallete)
-        draw = ImageDraw.Draw(image)
+        image = np.array(image)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
 
         index = 1
         for annotation in task["annotations"]:
             color = index if is_instance_segmentation else classes.index(annotation["value"]) + 1
             if annotation["type"] == AnnotationType.segmentation.value:
                 for region in annotation["points"]:
                     for points in region:
-                        pillow_draw_points = self.__get_pillow_draw_points(points)
-                        draw.polygon(pillow_draw_points, fill=color)
+                        cv_draw_points = self.__get_cv_draw_points(points)
+                        cv2.fillPoly(image, [cv_draw_points], color, lineType=cv2.LINE_8, shift=0)
                         # hollowd points are not supported
                         break
             elif annotation["type"] == AnnotationType.polygon.value:
-                pillow_draw_points = self.__get_pillow_draw_points(annotation["points"])
-                draw.polygon(pillow_draw_points, fill=color)
+                cv_draw_points = self.__get_cv_draw_points(annotation["points"])
+                cv2.fillPoly(image, [cv_draw_points], color, lineType=cv2.LINE_8, shift=0)
             elif annotation["type"] == AnnotationType.bbox.value:
-                pillow_draw_points = self.__get_pillow_draw_points(annotation["points"])
-                draw.polygon(pillow_draw_points, fill=color)
+                cv_draw_points = self.__get_cv_draw_points(annotation["points"])
+                cv2.fillPoly(image, [cv_draw_points], color, lineType=cv2.LINE_8, shift=0)
             else:
                 continue
             index += 1
 
         image_path = os.path.join(output_dir, utils.get_basename(task["name"]) + ".png")
         os.makedirs(os.path.dirname(image_path), exist_ok=True)
+        image = Image.fromarray(image)
+        image = image.convert('P')
+        image.putpalette(pallete)
         image.save(image_path)
 
-    def __get_pillow_draw_points(self, points: List[int]) -> List[int]:
+    def __get_cv_draw_points(self, points: List[int]) -> List[int]:
         """
         Convert points to pillow draw points. Diagonal points are not supported.
         """
@@ -717,7 +721,11 @@ def __get_pillow_draw_points(self, points: List[int]) -> List[int]:
         for i in range(int(len(points) / 2)):
             new_points.append(x_points[i * 2])
             new_points.append(y_points[i * 2 + 1])
-        return new_points
+
+        cv_points = []
+        for i in range(int(len(new_points) / 2)):
+            cv_points.append((new_points[i * 2], new_points[i * 2 + 1]))
+        return np.array(cv_points)
 
 
     # Annotation

fastlabel/converters.py

@@ -301,14 +301,65 @@ def to_pixel_coordinates(tasks: list) -> list:
                 ]
             else:
                 continue
+
+    # Remove duplicate points
+    for task in tasks:
+        for anno in task["annotations"]:
+            if annotation["type"] == AnnotationType.segmentation.value:
+                new_regions = []
+                for region in anno["points"]:
+                    new_region = []
+                    for points in region:
+                        new_points = __remove_duplicated_coordinates(points)
+                        new_region.append(new_points)
+                    new_regions.append(new_region)
+                anno["points"] = new_regions
+            elif annotation["type"] == AnnotationType.polygon.value:
+                new_points = __remove_duplicated_coordinates(annotation["points"])
+                annotation["points"] = new_points
     return tasks
 
+def __remove_duplicated_coordinates(points: List[int]) -> List[int]:
+    """
+    Remove duplicated coordinates.
+    """
+    if len(points) == 0:
+        return []
+
+    new_points = []
+    for i in range(int(len(points) / 2)):
+        if i == 0:
+            new_points.append(points[i*2])
+            new_points.append(points[i*2 + 1])
+
+        if new_points[-2] == points[i*2] and new_points[-1] == points[i*2 + 1]:
+            continue
+
+        if len(new_points) <= 2:
+            new_points.append(points[i*2])
+            new_points.append(points[i*2 + 1])
+        else:
+            if new_points[-4] == new_points[-2] and new_points[-2] == points[i*2]:
+                new_points.pop()
+                new_points.pop()
+                new_points.append(points[i*2])
+                new_points.append(points[i*2 + 1])
+            elif new_points[-3] == new_points[-1] and new_points[-1] == points[i*2 + 1]:
+                new_points.pop()
+                new_points.pop()
+                new_points.append(points[i*2])
+                new_points.append(points[i*2 + 1])
+            else:
+                new_points.append(points[i*2])
+                new_points.append(points[i*2 + 1])
+    return new_points
+
 def __get_pixel_coordinates(points: List[int or float]) -> List[int]:
     """
     Remove diagonal coordinates and return pixel outline coordinates.
     """
     if len(points) == 0:
-        return
+        return []
 
     new_points = []
     new_points.append(int(points[0]))

requirements.txt

@@ -2,4 +2,5 @@ requests==2.25.1
 numpy==1.20.2
 geojson==2.5.0
 xmltodict==0.12.0
-Pillow==8.3.1
+Pillow==8.3.1
+opencv-python==4.5.3.56

setup.py

@@ -8,7 +8,7 @@
 
 setuptools.setup(
     name="fastlabel",
-    version="0.9.8",
+    version="0.9.9",
     author="eisuke-ueta",
     author_email="[email protected]",
     description="The official Python SDK for FastLabel API, the Data Platform for AI",