add option for export yolo with intended annotations

README.md

@@ -889,7 +889,18 @@ Get tasks and export as YOLO format files.
 
 ```python
 tasks = client.get_image_tasks(project="YOUR_PROJECT_SLUG")
-client.export_yolo(tasks)
+client.export_yolo(tasks, output_dir="YOUR_DIRECTROY")
+```
+
+Get tasks and export as YOLO format files with classes.txt
+You can use fixed classes.txt and arrange order of each annotaiton file's order
+
+```python
+project_slug = "YOUR_PROJECT_SLUG"
+tasks = client.get_image_tasks(project=project_slug)
+annotations = client.get_annotations(project=project_slug)
+classes = list(map(lambda annotation: annotation["value"], annotations))
+client.export_yolo(tasks=tasks, classes=classes, output_dir="YOUR_DIRECTROY")
 ```
 
 ### Pascal VOC

fastlabel/__init__.py

@@ -651,14 +651,17 @@ def export_coco(self, tasks: list, output_dir: str = os.path.join("output", "coc
         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:
+    def export_yolo(self, tasks: list, classes: list = [], output_dir: str = os.path.join("output", "yolo")) -> None:
         """
         Convert tasks to YOLO format and export as files.
+        If you pass classes, classes.txt will be generated based on it .
+        If not , classes.txt will be generated based on passed tasks .(Annotations never used in your project will not be exported.)
 
         tasks is a list of tasks. (Required)
+        classes is a list of annotation values.  e.g. ['dog','bird'] (Optional)
         output_dir is output directory(default: output/yolo). (Optional)
         """
-        annos, categories = converters.to_yolo(tasks)
+        annos, categories = converters.to_yolo(tasks, classes)
         for anno in annos:
             file_name = anno["filename"]
             basename = utils.get_basename(file_name)

fastlabel/converters.py

@@ -145,10 +145,12 @@ def __calc_area(annotation_type: str, points: list) -> float:
 # YOLO
 
 
-def to_yolo(tasks: list) -> tuple:
-    coco = to_coco(tasks)
-    yolo = __coco2yolo(coco)
-    return yolo
+def to_yolo(tasks: list, classes: list) -> tuple:
+    if len(classes) == 0:
+        coco = to_coco(tasks)
+        return __coco2yolo(coco)
+    else:
+        return __to_yolo(tasks, classes)
 
 
 def __coco2yolo(coco: dict) -> tuple:
@@ -199,6 +201,69 @@ def __coco2yolo(coco: dict) -> tuple:
     return annos, categories
 
 
+def __to_yolo(tasks: list, classes: list) -> tuple:
+    annos = []
+    for task in tasks:
+        if task["height"] == 0 or task["width"] == 0:
+            continue
+        objs = []
+        data = [{"annotation": annotation, "task": task, "classes": classes}
+                for annotation in task["annotations"]]
+        with ThreadPoolExecutor(max_workers=8) as executor:
+            results = executor.map(__get_yolo_annotation, data)
+            for result in results:
+                if not result:
+                    continue
+                objs.append(" ".join(result))
+            anno = {
+                "filename": task["name"],
+                "object": objs
+            }
+            annos.append(anno)
+
+    categories = map(lambda val: {"name": val}, classes)
+
+    return annos, categories
+
+
+def __get_yolo_annotation(data: dict) -> dict:
+    annotation = data["annotation"]
+    points = annotation["points"]
+    annotation_type = annotation["type"]
+    value = annotation["value"]
+    classes = list(data["classes"])
+    task = data["task"]
+    if annotation_type != AnnotationType.bbox.value and annotation_type != AnnotationType.polygon.value:
+        return None
+    if not points or len(points) == 0:
+        return None
+    if annotation_type == AnnotationType.bbox.value and (int(points[0]) == int(points[2]) or int(points[1]) == int(points[3])):
+        return None
+    if not annotation["value"] in classes:
+        return None
+
+    dw = 1. / task["width"]
+    dh = 1. / task["height"]
+
+    bbox = __to_bbox(points)
+    xmin = bbox[0]
+    ymin = bbox[1]
+    xmax = bbox[0] + bbox[2]
+    ymax = bbox[1] + 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))
+    category_index = str(classes.index(value))
+    return [category_index, x, y, w, h]
+
+
 def _truncate(n, decimals=0) -> float:
     multiplier = 10 ** decimals
     return int(n * multiplier) / multiplier

setup.py

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