From df55636fab880d625190e5df6834476405ccd40c Mon Sep 17 00:00:00 2001
From: ueta-eisuke <ueta.eisuke1@gmail.com>
Date: Sat, 3 Jul 2021 12:48:55 +0900
Subject: [PATCH] add labelme format

---
 README.md               | 10 ++++++++
 fastlabel/__init__.py   | 16 ++++++++++++
 fastlabel/converters.py | 54 +++++++++++++++++++++++++++++++++++++++++
 setup.py                |  2 +-
 4 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 6096c85..e3e4fd5 100644
--- a/README.md
+++ b/README.md
@@ -19,6 +19,7 @@ _If you are using FastLabel prototype, please install version 0.2.2._
   - [COCO](#coco)
   - [YOLO](#yolo)
   - [Pascal VOC](#pascal-voc)
+  - [labelme](#labelme)
 
 ## Installation
 
@@ -815,6 +816,15 @@ tasks = client.get_image_tasks(project="YOUR_PROJECT_SLUG")
 client.export_pascalvoc(tasks)
 ```
 
+### labelme
+
+- Get tasks and export as labelme format files.
+
+```python
+tasks = client.get_image_tasks(project="YOUR_PROJECT_SLUG")
+client.export_labelme(tasks)
+```
+
 ## API Docs
 
 Check [this](https://api.fastlabel.ai/docs/) for further information.
diff --git a/fastlabel/__init__.py b/fastlabel/__init__.py
index 9d0d310..766ec9e 100644
--- a/fastlabel/__init__.py
+++ b/fastlabel/__init__.py
@@ -509,6 +509,22 @@ def export_pascalvoc(self, tasks: list, output_dir: str = os.path.join("output",
             with open(file_path, 'w', encoding="utf8") as f:
                 f.write(xml)
 
+    def export_labelme(self, tasks: list, output_dir: str = os.path.join("output", "labelme")) -> None:
+        """
+        Convert tasks to labelme format as files.
+
+        tasks is a list of tasks. (Required)
+        output_dir is output directory(default: output/labelme). (Optional)
+        """
+        labelmes = converters.to_labelme(tasks)
+        for labelme in labelmes:
+            file_name = labelme["imagePath"]
+            basename = utils.get_basename(file_name)
+            file_path = os.path.join(output_dir, basename + ".json")
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+            with open(file_path, 'w') as f:
+                json.dump(labelme, f, indent=4, ensure_ascii=False)
+
     # Annotation
 
     def find_annotation(self, annotation_id: str) -> dict:
diff --git a/fastlabel/converters.py b/fastlabel/converters.py
index eb1f9b2..2ee6c0e 100644
--- a/fastlabel/converters.py
+++ b/fastlabel/converters.py
@@ -220,6 +220,60 @@ def to_pascalvoc(tasks: list) -> list:
     return pascalvoc
 
 
+# labelme
+
+
+def to_labelme(tasks: list) -> list:
+    labelmes =[]
+    for task in tasks:
+        shapes = []
+        for annotation in task["annotations"]:
+                shape_type = __to_labelme_shape_type(annotation["type"])
+                if not shape_type:
+                    continue
+                points = annotation["points"]
+                if len(points) == 0:
+                    continue
+
+                shape_points = []
+                if annotation["type"] == "segmentation":
+                    for i in range(int(len(points[0][0]) / 2)):
+                        shape_points.append([points[0][0][i * 2], points[0][0][(i * 2) + 1]]) 
+                else:
+                    for i in range(int(len(points) / 2)):
+                        shape_points.append([points[i * 2], points[(i * 2) + 1]]) 
+
+                shape = {
+                        "label": annotation["value"],
+                        "points": shape_points,
+                        "group_id": None,
+                        "shape_type": shape_type,
+                        "flags": {}
+                }
+                shapes.append(shape)
+        labelmes.append({
+                "version": "4.5.9",
+                "flags": {},
+                "shapes": shapes,
+                "imagePath": task["name"],
+                "imageData": None,
+                "imageHeight": task["height"],
+                "imageWidth": task["width"],
+        })
+    return labelmes
+
+def __to_labelme_shape_type(annotation_type: str) -> str:
+    if annotation_type == "polygon" or annotation_type == "segmentation":
+        return "polygon"
+    if annotation_type == "bbox":
+        return "rectangle"
+    if annotation_type == "keypoint":
+        return "point"
+    if annotation_type == "line":
+        return "line"
+    return None
+
+
 def __coco2pascalvoc(coco: dict) -> list:
     pascalvoc = []
 
diff --git a/setup.py b/setup.py
index bdaefc0..49210ab 100644
--- a/setup.py
+++ b/setup.py
@@ -8,7 +8,7 @@
 
 setuptools.setup(
     name="fastlabel",
-    version="0.9.3",
+    version="0.9.4",
     author="eisuke-ueta",
     author_email="eisuke.ueta@fastlabel.ai",
     description="The official Python SDK for FastLabel API, the Data Platform for AI",