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..36c6412 100644 --- a/fastlabel/__init__.py +++ b/fastlabel/__init__.py @@ -529,7 +529,8 @@ def create_video_classification_task( """ endpoint = "tasks/video/classification" if not utils.is_video_supported_ext(file_path): - raise FastLabelInvalidException("Supported extensions are mp4.", 422) + raise FastLabelInvalidException( + "Supported extensions are mp4.", 422) file = utils.base64_encode(file_path) payload = {"project": project, "name": name, "file": file} if status: @@ -651,14 +652,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) @@ -709,12 +713,12 @@ def export_labelme(self, tasks: list, output_dir: str = os.path.join("output", " with open(file_path, 'w') as f: json.dump(labelme, f, indent=4, ensure_ascii=False) - # Instance / Semantic Segmetation + def export_instance_segmentation(self, tasks: list, output_dir: str = os.path.join("output", "instance_segmentation"), pallete: List[int] = const.COLOR_PALETTE) -> None: """ Convert tasks to index color instance segmentation (PNG files). - Supports only bbox, polygon and segmentation annotation types. Hollowed points are not supported. + Supports only bbox, polygon and segmentation annotation types. Supports up to 57 instances in default colors palette. Check const.COLOR_PALETTE for more details. tasks is a list of tasks. (Required) @@ -723,12 +727,13 @@ def export_instance_segmentation(self, tasks: list, output_dir: str = os.path.jo """ tasks = converters.to_pixel_coordinates(tasks) for task in tasks: - self.__export_index_color_image(task=task, output_dir=output_dir, pallete=pallete, is_instance_segmentation=True) - + self.__export_index_color_image( + task=task, output_dir=output_dir, pallete=pallete, is_instance_segmentation=True) + def export_semantic_segmentation(self, tasks: list, output_dir: str = os.path.join("output", "semantic_segmentation"), pallete: List[int] = const.COLOR_PALETTE) -> None: """ Convert tasks to index color semantic segmentation (PNG files). - Supports only bbox, polygon and segmentation annotation types. Hollowed points are not supported. + Supports only bbox, polygon and segmentation annotation types. Check const.COLOR_PALETTE for color pallete. tasks is a list of tasks. (Required) @@ -744,7 +749,8 @@ def export_semantic_segmentation(self, tasks: list, output_dir: str = os.path.jo tasks = converters.to_pixel_coordinates(tasks) for task in tasks: - self.__export_index_color_image(task=task, output_dir=output_dir, pallete=pallete, is_instance_segmentation=False, classes=classes) + self.__export_index_color_image( + task=task, output_dir=output_dir, pallete=pallete, is_instance_segmentation=False, classes=classes) def __export_index_color_image(self, task: list, output_dir: str, pallete: List[int], is_instance_segmentation: bool = True, classes: list = []) -> None: image = Image.new("RGB", (task["width"], task["height"]), 0) @@ -753,28 +759,39 @@ def __export_index_color_image(self, task: list, output_dir: str, pallete: List[ index = 1 for annotation in task["annotations"]: - color = index if is_instance_segmentation else classes.index(annotation["value"]) + 1 + color = index if is_instance_segmentation else classes.index( + annotation["value"]) + 1 if annotation["type"] == AnnotationType.segmentation.value: for region in annotation["points"]: count = 0 for points in region: - cv_draw_points = self.__get_cv_draw_points(points) if count == 0: - cv2.fillPoly(image, [cv_draw_points], color, lineType=cv2.LINE_8, shift=0) + cv_draw_points = self.__get_cv_draw_points(points) + cv2.fillPoly( + image, [cv_draw_points], color, lineType=cv2.LINE_8, shift=0) else: - cv2.fillPoly(image, [cv_draw_points], 0, lineType=cv2.LINE_8, shift=0) + # Reverse hollow points for opencv because this points are counter clockwise + cv_draw_points = self.__get_cv_draw_points( + utils.reverse_points(points)) + cv2.fillPoly( + image, [cv_draw_points], 0, lineType=cv2.LINE_8, shift=0) count += 1 elif annotation["type"] == AnnotationType.polygon.value: - cv_draw_points = self.__get_cv_draw_points(annotation["points"]) - cv2.fillPoly(image, [cv_draw_points], color, lineType=cv2.LINE_8, shift=0) + cv_draw_points = self.__get_cv_draw_points( + annotation["points"]) + cv2.fillPoly(image, [cv_draw_points], color, + lineType=cv2.LINE_8, shift=0) elif annotation["type"] == AnnotationType.bbox.value: - cv_draw_points = self.__get_cv_draw_points(annotation["points"]) - cv2.fillPoly(image, [cv_draw_points], color, lineType=cv2.LINE_8, shift=0) + cv_draw_points = self.__get_cv_draw_points( + annotation["points"]) + cv2.fillPoly(image, [cv_draw_points], color, + lineType=cv2.LINE_8, shift=0) else: continue index += 1 - image_path = os.path.join(output_dir, utils.get_basename(task["name"]) + ".png") + image_path = os.path.join( + output_dir, utils.get_basename(task["name"]) + ".png") os.makedirs(os.path.dirname(image_path), exist_ok=True) image = Image.fromarray(image) image = image.convert('P') @@ -823,7 +840,6 @@ def __get_cv_draw_points(self, points: List[int]) -> List[int]: cv_points.append((new_points[i * 2], new_points[i * 2 + 1])) return np.array(cv_points) - # Annotation def find_annotation(self, annotation_id: str) -> dict: diff --git a/fastlabel/converters.py b/fastlabel/converters.py index ed56ae6..47b9409 100644 --- a/fastlabel/converters.py +++ b/fastlabel/converters.py @@ -4,6 +4,7 @@ import copy import geojson import numpy as np +import math from fastlabel.const import AnnotationType # COCO @@ -144,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: @@ -198,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 @@ -207,10 +273,75 @@ def _truncate(n, decimals=0) -> float: def to_pascalvoc(tasks: list) -> list: - coco = to_coco(tasks) - pascalvoc = __coco2pascalvoc(coco) + pascalvoc = [] + for task in tasks: + if task["height"] == 0 or task["width"] == 0: + continue + + pascal_objs = [] + data = [{"annotation": annotation} + for annotation in task["annotations"]] + with ThreadPoolExecutor(max_workers=8) as executor: + results = executor.map(__get_pascalvoc_obj, data) + + for result in results: + if not result: + continue + pascal_objs.append(result) + + voc = { + "annotation": { + "filename": task["name"], + "size": { + "width": task["width"], + "height": task["height"], + "depth": 3, + }, + "segmented": 0, + "object": pascal_objs + } + } + pascalvoc.append(voc) return pascalvoc +def __get_pascalvoc_obj(data: dict) -> dict: + annotation = data["annotation"] + points = annotation["points"] + annotation_type = annotation["type"] + 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 + bbox = __to_bbox(points) + x = bbox[0] + y = bbox[1] + w = bbox[2] + h = bbox[3] + + return { + "name": annotation["value"], + "pose": "Unspecified", + "truncated": __get_pascalvoc_tag_value(annotation, "truncated"), + "occluded": __get_pascalvoc_tag_value(annotation, "occluded"), + "difficult": __get_pascalvoc_tag_value(annotation, "difficult"), + "bndbox": { + "xmin": math.floor(x), + "ymin": math.floor(y), + "xmax": math.floor(x + w), + "ymax": math.floor(y + h), + }, + } + +def __get_pascalvoc_tag_value(annotation: dict, target_tag_name: str) -> int: + attributes = annotation["attributes"] + if not attributes: + return 0 + related_attr = next( + (attribute for attribute in attributes if attribute["type"] == "switch" and attribute["key"] == target_tag_name), None) + return int(related_attr["value"]) if related_attr else 0 + # labelme @@ -388,58 +519,3 @@ def __get_pixel_coordinates(points: List[int or float]) -> List[int]: new_points.append(int(prev_x + int(xdiff / mindiff * (i + 1)))) new_points.append(int(prev_y + int(ydiff / mindiff * (i + 1)))) return new_points - -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 diff --git a/fastlabel/utils.py b/fastlabel/utils.py index 7984222..58d9254 100644 --- a/fastlabel/utils.py +++ b/fastlabel/utils.py @@ -21,3 +21,18 @@ def get_basename(file_path: str) -> str: path/to/file.jpg -> path/to/file """ return os.path.splitext(file_path)[0] + + +def reverse_points(points: list[int]) -> list[int]: + """ + e.g.) + [4, 5, 4, 9, 8, 9, 8, 5, 4, 5] => [4, 5, 8, 5, 8, 9, 4, 9, 4, 5] + """ + reversed_points = [] + for index, _ in enumerate(points): + if index % 2 == 0: + reversed_points.insert( + 0, points[index + 1]) + reversed_points.insert( + 0, points[index]) + return reversed_points diff --git a/setup.py b/setup.py index 1beaa53..f22dd75 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,7 @@ setuptools.setup( name="fastlabel", - version="0.11.1", + version="0.11.4", author="eisuke-ueta", author_email="eisuke.ueta@fastlabel.ai", description="The official Python SDK for FastLabel API, the Data Platform for AI",