Merge pull request #31 from fastlabel/develop

Merge to main

Author: faycute

.gitignore

@@ -11,4 +11,5 @@ annotation.json
 *.jpg
 *.jpeg
 *.png
-tests/
+tests/
+output/

README.md

@@ -501,10 +501,10 @@ map = client.get_task_id_name_map(project="YOUR_PROJECT_SLUG")
 
 ```python
 annotation_id = client.create_annotation(
-    project="YOUR_PROJECT_SLUG", type="bbox", value="cat", title="Cat", color="#FF0000")
+    project="YOUR_PROJECT_SLUG", type="bbox", value="cat", title="Cat")
 ```
 
-- Create a new annotation with attributes.
+- Create a new annotation with color and attributes.
 
 ```python
 attributes = [
@@ -663,13 +663,39 @@ client.delete_annotation(annotation_id="YOUR_ANNOTATIPN_ID")
 
 ## Converter
 
+Supporting bbox or polygon annotation type.
+
 ### COCO
 
-- Get tasks and convert to [COCO format](https://cocodataset.org/#format-data) (supporting bbox or polygon annotation type).
+- Get tasks and export as a [COCO format](https://cocodataset.org/#format-data) file.
+
+```python
+tasks = client.get_image_tasks(project="YOUR_PROJECT_SLUG")
+client.export_coco(tasks)
+```
+
+- Export with specifying output directory.
+
+```python
+client.export_coco(tasks=tasks, output_dir="YOUR_DIRECTROY")
+```
+
+### YOLO
+
+- Get tasks and export as YOLO format files.
+
+```python
+tasks = client.get_image_tasks(project="YOUR_PROJECT_SLUG")
+client.export_yolo(tasks)
+```
+
+### Pascal VOC
+
+- Get tasks and export as Pascal VOC format files.
 
 ```python
 tasks = client.get_image_tasks(project="YOUR_PROJECT_SLUG")
-pprint(client.to_coco(tasks))
+client.export_pascalvoc(tasks)
 ```
 
 ## API Docs

fastlabel/__init__.py

@@ -1,7 +1,10 @@
 import os
 import glob
+import json
 from logging import getLogger
 
+import xmltodict
+
 from .exceptions import FastLabelInvalidException
 from .api import Api
 from fastlabel import converters, utils
@@ -246,7 +249,7 @@ def get_task_id_name_map(
         self, project: str,
         offset: int = None,
         limit: int = 1000,
-    ) -> list:
+    ) -> dict:
         """
         Returns a map of task ids and names.
         e.g.) {
@@ -447,12 +450,58 @@ def delete_task(self, task_id: str) -> None:
 
     # Task Convert
 
-    def to_coco(self, tasks: list) -> dict:
+    def export_coco(self, tasks: list, output_dir: str = os.path.join("output", "coco")) -> None:
+        """
+        Convert tasks to COCO format and export as a file.
+
+        tasks is a list of tasks. (Required)
+        output_dir is output directory(default: output/coco). (Optional)
+        """
+        coco = converters.to_coco(tasks)
+        os.makedirs(output_dir, exist_ok=True)
+        file_path = os.path.join(output_dir, "annotations.json")
+        with open(file_path, 'w') as f:
+            json.dump(coco, f, indent=4, ensure_ascii=False)
+
+    def export_yolo(self, tasks: list, output_dir: str = os.path.join("output", "yolo")) -> None:
+        """
+        Convert tasks to YOLO format and export as files.
+
+        tasks is a list of tasks. (Required)
+        output_dir is output directory(default: output/yolo). (Optional)
         """
-        Convert tasks to COCO format.
+        annos, categories = converters.to_yolo(tasks)
+        for anno in annos:
+            file_name = anno["filename"]
+            basename = utils.get_basename(file_name)
+            file_path = os.path.join(
+                output_dir, "annotations", basename + ".txt")
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+            with open(file_path, 'w', encoding="utf8") as f:
+                for obj in anno["object"]:
+                    f.write(obj)
+                    f.write("\n")
+        with open(os.path.join(output_dir, "classes.txt"), 'w', encoding="utf8") as f:
+            for category in categories:
+                f.write(category["name"])
+                f.write("\n")
+
+    def export_pascalvoc(self, tasks: list, output_dir: str = os.path.join("output", "pascalvoc")) -> None:
         """
+        Convert tasks to Pascal VOC format as files.
 
-        return converters.to_coco(tasks)
+        tasks is a list of tasks. (Required)
+        output_dir is output directory(default: output/pascalvoc). (Optional)
+        """
+        pascalvoc = converters.to_pascalvoc(tasks)
+        for voc in pascalvoc:
+            file_name = voc["annotation"]["filename"]
+            basename = utils.get_basename(file_name)
+            file_path = os.path.join(output_dir, basename + ".xml")
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+            xml = xmltodict.unparse(voc, pretty=True, full_document=False)
+            with open(file_path, 'w', encoding="utf8") as f:
+                f.write(xml)
 
     # Annotation
 
@@ -507,7 +556,7 @@ def create_annotation(
         type: str,
         value: str,
         title: str,
-        color: str,
+        color: str = None,
         attributes: list = []
     ) -> str:
         """
@@ -517,7 +566,7 @@ def create_annotation(
         type can be 'bbox', 'polygon', 'keypoint', 'classification', 'line', 'segmentation'. (Required)
         value is an unique identifier of annotation in your project. (Required)
         title is a display name of value. (Required)
-        color is hex color code like #ffffff. (Required)
+        color is hex color code like #ffffff. (Optional)
         attributes is a list of attribute. (Optional)
         """
         endpoint = "annotations"
@@ -526,8 +575,9 @@ def create_annotation(
             "type": type,
             "value": value,
             "title": title,
-            "color": color
         }
+        if color:
+            payload["color"] = color
         if attributes:
             payload["attributes"] = attributes
         return self.api.post_request(endpoint, payload=payload)

fastlabel/converters.py

@@ -14,10 +14,10 @@ class AnnotationType(Enum):
     segmentation = "segmentation"
 
 
+# COCO
+
+
 def to_coco(tasks: list) -> dict:
-    """
-    Convert tasks to COCO format.
-    """
     # Get categories
     categories = __get_categories(tasks)
 
@@ -147,3 +147,130 @@ def __calc_area(annotation_type: str, points: list) -> float:
         area = 0.5 * np.abs(np.dot(x, np.roll(y, 1)) -
                             np.dot(y, np.roll(x, 1)))
     return area
+
+
+# YOLO
+
+
+def to_yolo(tasks: list) -> tuple:
+    coco = to_coco(tasks)
+    yolo = __coco2yolo(coco)
+    return yolo
+
+
+def __coco2yolo(coco: dict) -> tuple:
+    categories = coco["categories"]
+
+    annos = []
+    for image in coco["images"]:
+        dw = 1. / image["width"]
+        dh = 1. / image["height"]
+
+        # Get objects
+        objs = []
+        for annotation in coco["annotations"]:
+            if image["id"] != annotation["image_id"]:
+                continue
+
+            category_index = "0"
+            for index, category in enumerate(categories):
+                if category["id"] == annotation["category_id"]:
+                    category_index = str(index)
+                    break
+
+            xmin = annotation["bbox"][0]
+            ymin = annotation["bbox"][1]
+            xmax = annotation["bbox"][0] + annotation["bbox"][2]
+            ymax = annotation["bbox"][1] + annotation["bbox"][3]
+
+            x = (xmin + xmax) / 2
+            y = (ymin + ymax) / 2
+            w = xmax - xmin
+            h = ymax - ymin
+
+            x = str(_truncate(x * dw, 7))
+            y = str(_truncate(y * dh, 7))
+            w = str(_truncate(w * dw, 7))
+            h = str(_truncate(h * dh, 7))
+
+            obj = [category_index, x, y, w, h]
+            objs.append(" ".join(obj))
+
+        # get annotation
+        anno = {
+            "filename": image["file_name"],
+            "object": objs
+        }
+        annos.append(anno)
+
+    return annos, categories
+
+
+def _truncate(n, decimals=0) -> float:
+    multiplier = 10 ** decimals
+    return int(n * multiplier) / multiplier
+
+
+# Pascal VOC
+
+
+def to_pascalvoc(tasks: list) -> list:
+    coco = to_coco(tasks)
+    pascalvoc = __coco2pascalvoc(coco)
+    return pascalvoc
+
+
+def __coco2pascalvoc(coco: dict) -> list:
+    pascalvoc = []
+
+    for image in coco["images"]:
+
+        # Get objects
+        objs = []
+        for annotation in coco["annotations"]:
+            if image["id"] != annotation["image_id"]:
+                continue
+            category = _get_category_by_id(
+                coco["categories"], annotation["category_id"])
+
+            x = annotation["bbox"][0]
+            y = annotation["bbox"][1]
+            w = annotation["bbox"][2]
+            h = annotation["bbox"][3]
+
+            obj = {
+                "name": category["name"],
+                "pose": "Unspecified",
+                "truncated": 0,
+                "difficult": 0,
+                "bndbox": {
+                        "xmin": x,
+                        "ymin": y,
+                        "xmax": x + w,
+                        "ymax": y + h,
+                },
+            }
+            objs.append(obj)
+
+        # get annotation
+        voc = {
+            "annotation": {
+                "filename": image["file_name"],
+                "size": {
+                    "width": image["width"],
+                    "height": image["height"],
+                    "depth": 3,
+                },
+                "segmented": 0,
+                "object": objs
+            }
+        }
+        pascalvoc.append(voc)
+
+    return pascalvoc
+
+
+def _get_category_by_id(categories: list, id_: str) -> str:
+    category = [
+        category for category in categories if category["id"] == id_][0]
+    return category

fastlabel/utils.py

@@ -1,3 +1,4 @@
+import os
 import base64
 
 
@@ -12,3 +13,11 @@ def is_image_supported_ext(file_path: str) -> bool:
 
 def is_video_supported_ext(file_path: str) -> bool:
     return file_path.lower().endswith(('.mp4'))
+
+
+def get_basename(file_path: str) -> str:
+    """
+    e.g.) file.jpg -> file
+          path/to/file.jpg -> path/to/file
+    """
+    return os.path.splitext(file_path)[0]

requirements.txt

@@ -1,3 +1,4 @@
 requests==2.25.1
 numpy==1.20.2
-geojson==2.5.0
+geojson==2.5.0
+xmltodict==0.12.0

setup.py

@@ -8,7 +8,7 @@
 
 setuptools.setup(
     name="fastlabel",
-    version="0.7.1",
+    version="0.8.0",
     author="eisuke-ueta",
     author_email="eisuke.ueta@fastlabel.ai",
     description="The official Python SDK for FastLabel API, the Data Platform for AI",