From 6c2c68b3a8f7c66e4858fb342d3d6bf82935c9d7 Mon Sep 17 00:00:00 2001 From: ueta-eisuke Date: Fri, 30 Jul 2021 20:20:33 +0900 Subject: [PATCH] fix segmentation bug --- fastlabel/__init__.py | 32 +++++++++++++++---------- fastlabel/converters.py | 53 ++++++++++++++++++++++++++++++++++++++++- requirements.txt | 3 ++- setup.py | 2 +- 4 files changed, 75 insertions(+), 15 deletions(-) diff --git a/fastlabel/__init__.py b/fastlabel/__init__.py index bfc46f3..7f71f3f 100644 --- a/fastlabel/__init__.py +++ b/fastlabel/__init__.py @@ -3,7 +3,9 @@ import json from typing import List from logging import getLogger -from PIL import Image, ImageDraw +from PIL import Image +import cv2 +import numpy as np import xmltodict @@ -653,9 +655,8 @@ def export_semantic_segmentation(self, tasks: list, output_dir: str = os.path.jo 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) - image = image.convert('P') - image.putpalette(pallete) - draw = ImageDraw.Draw(image) + image = np.array(image) + image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) index = 1 for annotation in task["annotations"]: @@ -663,25 +664,28 @@ def __export_index_color_image(self, task: list, output_dir: str, pallete: List[ if annotation["type"] == AnnotationType.segmentation.value: for region in annotation["points"]: for points in region: - pillow_draw_points = self.__get_pillow_draw_points(points) - draw.polygon(pillow_draw_points, fill=color) + cv_draw_points = self.__get_cv_draw_points(points) + cv2.fillPoly(image, [cv_draw_points], color, lineType=cv2.LINE_8, shift=0) # hollowd points are not supported break elif annotation["type"] == AnnotationType.polygon.value: - pillow_draw_points = self.__get_pillow_draw_points(annotation["points"]) - draw.polygon(pillow_draw_points, fill=color) + 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: - pillow_draw_points = self.__get_pillow_draw_points(annotation["points"]) - draw.polygon(pillow_draw_points, fill=color) + 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") os.makedirs(os.path.dirname(image_path), exist_ok=True) + image = Image.fromarray(image) + image = image.convert('P') + image.putpalette(pallete) image.save(image_path) - def __get_pillow_draw_points(self, points: List[int]) -> List[int]: + def __get_cv_draw_points(self, points: List[int]) -> List[int]: """ Convert points to pillow draw points. Diagonal points are not supported. """ @@ -717,7 +721,11 @@ def __get_pillow_draw_points(self, points: List[int]) -> List[int]: for i in range(int(len(points) / 2)): new_points.append(x_points[i * 2]) new_points.append(y_points[i * 2 + 1]) - return new_points + + cv_points = [] + for i in range(int(len(new_points) / 2)): + cv_points.append((new_points[i * 2], new_points[i * 2 + 1])) + return np.array(cv_points) # Annotation diff --git a/fastlabel/converters.py b/fastlabel/converters.py index 51fd0dc..ec07d6d 100644 --- a/fastlabel/converters.py +++ b/fastlabel/converters.py @@ -301,14 +301,65 @@ def to_pixel_coordinates(tasks: list) -> list: ] else: continue + + # Remove duplicate points + for task in tasks: + for anno in task["annotations"]: + if annotation["type"] == AnnotationType.segmentation.value: + new_regions = [] + for region in anno["points"]: + new_region = [] + for points in region: + new_points = __remove_duplicated_coordinates(points) + new_region.append(new_points) + new_regions.append(new_region) + anno["points"] = new_regions + elif annotation["type"] == AnnotationType.polygon.value: + new_points = __remove_duplicated_coordinates(annotation["points"]) + annotation["points"] = new_points return tasks +def __remove_duplicated_coordinates(points: List[int]) -> List[int]: + """ + Remove duplicated coordinates. + """ + if len(points) == 0: + return [] + + new_points = [] + for i in range(int(len(points) / 2)): + if i == 0: + new_points.append(points[i*2]) + new_points.append(points[i*2 + 1]) + + if new_points[-2] == points[i*2] and new_points[-1] == points[i*2 + 1]: + continue + + if len(new_points) <= 2: + new_points.append(points[i*2]) + new_points.append(points[i*2 + 1]) + else: + if new_points[-4] == new_points[-2] and new_points[-2] == points[i*2]: + new_points.pop() + new_points.pop() + new_points.append(points[i*2]) + new_points.append(points[i*2 + 1]) + elif new_points[-3] == new_points[-1] and new_points[-1] == points[i*2 + 1]: + new_points.pop() + new_points.pop() + new_points.append(points[i*2]) + new_points.append(points[i*2 + 1]) + else: + new_points.append(points[i*2]) + new_points.append(points[i*2 + 1]) + return new_points + def __get_pixel_coordinates(points: List[int or float]) -> List[int]: """ Remove diagonal coordinates and return pixel outline coordinates. """ if len(points) == 0: - return + return [] new_points = [] new_points.append(int(points[0])) diff --git a/requirements.txt b/requirements.txt index f74a742..0f2e114 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ requests==2.25.1 numpy==1.20.2 geojson==2.5.0 xmltodict==0.12.0 -Pillow==8.3.1 \ No newline at end of file +Pillow==8.3.1 +opencv-python==4.5.3.56 \ No newline at end of file diff --git a/setup.py b/setup.py index a0a4b13..854c2e4 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,7 @@ setuptools.setup( name="fastlabel", - version="0.9.8", + version="0.9.9", author="eisuke-ueta", author_email="eisuke.ueta@fastlabel.ai", description="The official Python SDK for FastLabel API, the Data Platform for AI",