add segmentation logic

Author: shuichi

examples/bbox_to_visual_inspection_ai.py

@@ -0,0 +1,41 @@
+import json
+import os
+
+from visual_inspection_ai import VisualInspectionAi
+
+
+class VisualInspectionAiConverter:
+    def convert_from_bbox(
+        self, file_path: str, export_file_path: str, project_type: str
+    ) -> dict:
+        json = self.open_bbox(file_path)
+        if project_type == "bbox":
+            visual_inspection_ai = VisualInspectionAi.from_bbox(json)
+        else:
+            visual_inspection_ai = VisualInspectionAi.from_segmentation(json)
+
+        os.makedirs(os.path.dirname(export_file_path), exist_ok=True)
+        with open(export_file_path, "w") as file:
+            file.write(visual_inspection_ai.as_jsonl())
+
+    def open_bbox(self, file_path: str) -> any:
+        return json.load(open(file_path, "r"))
+
+
+converter = VisualInspectionAiConverter()
+
+export_file_path = (
+    "./export/visual_inspection_ai/test.jsonl"  # 出力先のファイルパスを指定
+)
+
+# bboxの場合
+# file_path = (
+#     "./import/visual_inspection_ai/bbox/annotations.json"  # 入力元のファイルパスを指定
+# )
+# project_type = "bbox"
+
+# segmentationの場合
+file_path = "./import/visual_inspection_ai/segmentation/annotations.json"  # 入力元のファイルパスを指定
+project_type = "segmentation"
+
+result = converter.convert_from_bbox(file_path, export_file_path, project_type)

examples/to_visual_inspection_ai.py

@@ -9,7 +9,7 @@ def __init__(self, vertex, project_type):
         self.project_type = project_type
 
     @classmethod
-    def fromBbox(cls, width, height, points):
+    def from_bbox(cls, width, height, points):
         vertices = [
             (points[0] / width, points[1] / height),  # 左上
             (points[2] / width, points[1] / height),  # 右上
@@ -18,21 +18,37 @@ def fromBbox(cls, width, height, points):
         ]
         return VisualInspectionAiFigure(vertices, "bbox")
 
+    @classmethod
+    def from_segmentation(cls, width, height, points):
+        # サンプルのjsonには矩形しかなかったので、とりあえずで変換
+        figure_points = points[0][
+            0
+        ]  # FIXME fastlabelでは飛地がある場合に配列に複数の要素がある場合があるが、VisualInspectionAiで飛地をどのように扱うか不明なため、一旦最初の図形だけパース。
+        vertices = [
+            {"x": figure_points[i] / width, "y": figure_points[i + 1] / height}
+            for i in range(0, len(figure_points), 2)
+        ]
+        return VisualInspectionAiFigure(vertices, "segmentation")
+
     # VisualInspectionAiのjsonlを作成する。
     def make_jsonl(self) -> dict:
         return {
-            "image_gcs_uri": "gs://ffod-98sm/20240613/NG/test/BQ1334B__11-7_18-19_0004.png",
+            "image_gcs_uri": "gs://ffod-98sm/20240613/NG/test/BQ1334B__11-7_18-19_0004.png",  # TODO ここに入れ込めそうな値はjsonにはないので、調査
             "vi_annotations": self.__make_vi_annotations_jsonl(),
-            "dataItemResourceLabels": {"goog_vi_ml_use": "test"},
+            "dataItemResourceLabels": {
+                "goog_vi_ml_use": "test"
+            },  # TODO この値が何を指しているのか調査
         }
 
     def __make_vi_annotations_jsonl(self) -> dict:
         return {
             "viBoundingPoly": {"vertex": self.vertex},
-            "annotationSpec": "defect",
-            "annotationSet": "Polygons Regions"
-            if self.project_type == "bbox"
-            else "TODO segmentationの場合の表記調査",
+            "annotationSpec": "defect",  # TODO ここにクラスが入る場合は入れる。
+            "annotationSet": (
+                "Polygons Regions"
+                if self.project_type == "bbox"
+                else "TODO segmentationの場合の表記調査"
+            ),
         }
 
 
@@ -44,23 +60,37 @@ def __init__(self, figures):
     def from_bbox(cls, bbox_json):
         figures = []
         for bbox in bbox_json:
-            figures.extend(VisualInspectionAi.parseBboxJson(bbox))
+            figures.extend(VisualInspectionAi.parse_bbox_json(bbox))
         return cls(figures)
 
     @classmethod
-    def parseBboxJson(cls, bbox_json) -> list[VisualInspectionAiFigure]:
+    def parse_bbox_json(cls, bbox_json) -> list[VisualInspectionAiFigure]:
         result = []
         width = bbox_json["width"]
         height = bbox_json["height"]
         for annotation in bbox_json["annotations"]:
             points = annotation["points"]
-            result.append(VisualInspectionAiFigure.fromBbox(width, height, points))
+            result.append(VisualInspectionAiFigure.from_bbox(width, height, points))
         return result
 
     @classmethod
-    def from_segmentation(cls, segmentation_json, project_type):
-        # TODO segmentationからvertexへの変換処理
-        return cls([])
+    def from_segmentation(cls, segmentation_json):
+        figures = []
+        for bbox in segmentation_json:
+            figures.extend(VisualInspectionAi.parse_segmentation_json(bbox))
+        return cls(figures)
+
+    @classmethod
+    def parse_segmentation_json(cls, json) -> list[VisualInspectionAiFigure]:
+        result = []
+        width = json["width"]
+        height = json["height"]
+        for annotation in json["annotations"]:
+            points = annotation["points"]
+            result.append(
+                VisualInspectionAiFigure.from_segmentation(width, height, points)
+            )
+        return result
 
     def __make_jsonl(self) -> list:
         print(self.figures)