diff --git a/README.md b/README.md index 0c46f2b..3ce13ff 100644 --- a/README.md +++ b/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 diff --git a/fastlabel/__init__.py b/fastlabel/__init__.py index 17cbb57..3d2a8d3 100644 --- a/fastlabel/__init__.py +++ b/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) diff --git a/fastlabel/converters.py b/fastlabel/converters.py index 38131d7..47b9409 100644 --- a/fastlabel/converters.py +++ b/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 diff --git a/setup.py b/setup.py index 26781b3..43ebc31 100644 --- a/setup.py +++ b/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",