refactored input converters

Author: agyurjin

superannotate/common.py

@@ -1,3 +1,4 @@
+import numpy as np
 import functools
 import logging
 from pathlib import Path
@@ -138,3 +139,26 @@ def rgb_to_hex(rgb_tuple):
     """Converts RGB values to HEX values
     """
     return '#%02x%02x%02x' % rgb_tuple
+
+
+def blue_color_generator(n, hex_values=True):
+    """ Blue colors generator for SuperAnnotate blue mask.
+    """
+    hex_colors = []
+    for i in range(n + 1):
+        int_color = i * 15
+        bgr_color = np.array(
+            [
+                int_color & 255, (int_color >> 8) & 255,
+                (int_color >> 16) & 255, 255
+            ],
+            dtype=np.uint8
+        )
+        hex_color = '#' + "{:02x}".format(
+            bgr_color[2]
+        ) + "{:02x}".format(bgr_color[1], ) + "{:02x}".format(bgr_color[0])
+        if hex_values:
+            hex_colors.append(hex_color)
+        else:
+            hex_colors.append(hex_to_rgb(hex_color))
+    return hex_colors[1:]

superannotate/input_converters/converters/coco_converters/coco_converter.py

@@ -14,19 +14,14 @@
 
 
 class CoCoConverter(object):
-    def __init__(
-        self,
-        dataset_name_,
-        export_root_,
-        project_type_,
-        output_dir_,
-        task_=None
-    ):
-        self.project_type = project_type_
-        self.dataset_name = dataset_name_
-        self.export_root = export_root_
-        self.output_dir = output_dir_
-        self.task = task_
+    def __init__(self, args):
+        self.project_type = args.project_type
+        self.dataset_name = args.dataset_name
+        self.export_root = args.export_root
+        self.output_dir = args.output_dir
+        self.task = args.task
+        self.direction = args.direction
+
         self.failed_conversion_cnt = 0
 
     def _create_single_category(self, item):

superannotate/input_converters/converters/coco_converters/coco_strategies.py

@@ -14,14 +14,11 @@
 from .coco_to_sa_vector import coco_keypoint_detection_to_sa_vector, coco_instance_segmentation_to_sa_vector
 
 
-class PanopticConverterStrategy(CoCoConverter):
+class CocoPanopticConverterStrategy(CoCoConverter):
     name = "Panoptic converter"
 
-    def __init__(
-        self, dataset_name, export_root, project_type, output_dir, direction
-    ):
-        super().__init__(dataset_name, export_root, project_type, output_dir)
-        self.direction = direction
+    def __init__(self, args):
+        super().__init__(args)
         self.__setup_conversion_algorithm()
 
     def __setup_conversion_algorithm(self):
@@ -100,17 +97,11 @@ def to_sa_format(self):
         #     loader = self.conversion_algorithm(json_data, self.output_dir)
 
 
-class ObjectDetectionStrategy(CoCoConverter):
+class CocoObjectDetectionStrategy(CoCoConverter):
     name = "ObjectDetection converter"
 
-    def __init__(
-        self, dataset_name, export_root, project_type, output_dir, task,
-        direction
-    ):
-        self.direction = direction
-        super().__init__(
-            dataset_name, export_root, project_type, output_dir, task
-        )
+    def __init__(self, args):
+        super().__init__(args)
         self.__setup_conversion_algorithm()
 
     def __setup_conversion_algorithm(self):
@@ -220,16 +211,11 @@ def to_sa_format(self):
         #     loader = self.conversion_algorithm(json_data, self.output_dir)
 
 
-class KeypointDetectionStrategy(CoCoConverter):
+class CocoKeypointDetectionStrategy(CoCoConverter):
     name = 'Keypoint Detection Converter'
 
-    def __init__(
-        self, dataset_name, export_root, project_type, output_dir, direction
-    ):
-        super().__init__(
-            dataset_name, export_root, project_type, output_dir, direction
-        )
-        self.direction = direction
+    def __init__(self, args):
+        super().__init__(args)
         self.__setup_conversion_algorithm()
 
     def __str__(self):

superannotate/input_converters/converters/coco_converters/coco_to_sa_pixel.py

@@ -7,37 +7,12 @@
 from panopticapi.utils import id2rgb
 from tqdm import tqdm
 
-from ....common import hex_to_rgb
-
-# def _rle_to_bitmask(coco_json, annotation):
-#     return coco.annToMask(annotation)
-
-
-# Generates blue colors in range(n)
-def _blue_color_generator(n, hex_values=True):
-    hex_colors = []
-    for i in range(n):
-        int_color = i * 15
-        bgr_color = np.array(
-            [
-                int_color & 255, (int_color >> 8) & 255,
-                (int_color >> 16) & 255, 255
-            ],
-            dtype=np.uint8
-        )
-        hex_color = '#' + "{:02x}".format(
-            bgr_color[2]
-        ) + "{:02x}".format(bgr_color[1], ) + "{:02x}".format(bgr_color[0])
-        if hex_values:
-            hex_colors.append(hex_color)
-        else:
-            hex_colors.append(hex_to_rgb(hex_color))
-    return hex_colors
+from ....common import hex_to_rgb, blue_color_generator
 
 
 def coco_panoptic_segmentation_to_sa_pixel(coco_path, images_path):
     coco_json = json.load(open(coco_path))
-    hex_colors = _blue_color_generator(len(coco_json["categories"]))
+    hex_colors = blue_color_generator(len(coco_json["categories"]))
     annotate_list = coco_json["annotations"]
 
     for annotate in tqdm(annotate_list, "Converting"):
@@ -55,18 +30,18 @@ def coco_panoptic_segmentation_to_sa_pixel(coco_path, images_path):
         H, W, C = img.shape
         img = img.reshape((H * W, C))
         segments = annotate["segments_info"]
-        hex_colors = _blue_color_generator(len(segments) + 1)
+        hex_colors = blue_color_generator(len(segments))
 
         out_json = []
         for i, seg in enumerate(segments):
             img[np.all(img == id2rgb(seg["id"]),
-                       axis=1)] = hex_to_rgb(hex_colors[i + 1])
+                       axis=1)] = hex_to_rgb(hex_colors[i])
             dd = {
                 "classId": seg["category_id"],
                 "probability": 100,
                 "visible": True,
                 "parts": [{
-                    "color": hex_colors[i + 1]
+                    "color": hex_colors[i]
                 }],
                 "attributes": [],
                 "attributeNames": [],
@@ -103,11 +78,15 @@ def coco_instance_segmentation_to_sa_pixel(coco_path, images_path):
         }
 
     sa_json = {}
-    for annot in tqdm(coco_json['annotations'], "Convertring annotations"):
+    hexcolors = blue_color_generator(len(coco_json['annotations']))
+    for i, annot in enumerate(coco_json['annotations']):
         if str(annot['image_id']) not in images_dict:
             continue
-        color = np.random.choice(range(256), size=3)
-        hexcolor = "#%02x%02x%02x" % tuple(color)
+
+        hexcolor = hexcolors[images_dict[str(annot['image_id'])]['segments_num']
+                            ]
+        color = hex_to_rgb(hexcolor)
+        images_dict[str(annot['image_id'])]['segments_num'] += 1
 
         images_dict[str(annot['image_id']
                        )]['mask'][maskUtils.decode(annot['segmentation']) == 1

superannotate/input_converters/converters/converters.py

@@ -1,20 +1,10 @@
 '''
-This may look over-engineered at this point however the idea is the following:
-1. We eventually might want to convert to and from other formats than COCO
-2. For each of these formats (COCO included) there should be different strategies
-   for conversion. In COCO's case there are 5
-   1.1 Panoptic
-   1.2 Object Detection
-   1.3 Stuff Detection
-   1.4 Keypoint Detection
-   1.5 Image Captioning
-3. We will have a general Converter object will not care about the format or the
-   conversion strategy. It has to methods:
-   3.1 convert from sa format to desired format` convert_from_sa()
-   3.2 convert from some format to sa format` convert_to_sa()
-4. This object will receive the strategy from outside and will convert according to
+
+This object will receive the strategy from outside and will convert according to
    said strategy.
+
 '''
+
 import glob
 import os
 import json
@@ -23,22 +13,17 @@
 
 from tqdm import tqdm
 
-from .coco_converters.coco_strategies import ObjectDetectionStrategy, KeypointDetectionStrategy, PanopticConverterStrategy
+from .coco_converters.coco_strategies import CocoObjectDetectionStrategy, CocoKeypointDetectionStrategy, CocoPanopticConverterStrategy
 from .voc_converters.voc_strategies import VocObjectDetectionStrategy
 from .labelbox_converters.labelbox_strategies import LabelBoxObjectDetectionStrategy
 from .dataloop_converters.dataloop_strategies import DataLoopObjectDetectionStrategy
 from .supervisely_converters.supervisely_strategies import SuperviselyObjectDetectionStrategy
 
 
 class Converter(object):
-    def __init__(
-        self, project_type, task, dataset_name, export_root, output_dir, method
-    ):
-        self.export_root = export_root
-        self.output_dir = output_dir
-        self._select_strategy(
-            project_type, task, dataset_name, export_root, output_dir, method
-        )
+    def __init__(self, args):
+        self.output_dir = args.output_dir
+        self._select_strategy(args)
 
     def convert_from_sa(self):
         self.strategy.sa_to_output_format()
@@ -51,52 +36,29 @@ def convert_to_sa(self, platform):
     def __set_strategy(self, c_strategy):
         self.strategy = c_strategy
 
-    def _select_strategy(
-        self, project_type, task, dataset_name, export_root, output_dir, method
-    ):
-        if method.dataset_format == "COCO":
-            if task == 'instance_segmentation' or task == 'object_detection':
-                c_strategy = ObjectDetectionStrategy(
-                    dataset_name, export_root, project_type, output_dir, task,
-                    method.direction
-                )
-            if task == 'keypoint_detection':
-                c_strategy = KeypointDetectionStrategy(
-                    dataset_name, export_root, project_type, output_dir,
-                    method.direction
-                )
-            if task == 'panoptic_segmentation':
-                c_strategy = PanopticConverterStrategy(
-                    dataset_name, export_root, project_type, output_dir,
-                    method.direction
-                )
-        elif method.dataset_format == "VOC":
-            if task == 'instance_segmentation' or task == 'object_detection':
-                c_strategy = VocObjectDetectionStrategy(
-                    dataset_name, export_root, project_type, output_dir, task,
-                    method.direction
-                )
-        elif method.dataset_format == "LabelBox":
-            if task == "object_detection" or task == 'instance_segmentation' or task == 'vector_annotation':
-                c_strategy = LabelBoxObjectDetectionStrategy(
-                    dataset_name, export_root, project_type, output_dir, task,
-                    method.direction
-                )
-        elif method.dataset_format == "DataLoop":
-            if task == 'object_detection' or task == 'instance_segmentation' or task == 'vector_annotation':
-                c_strategy = DataLoopObjectDetectionStrategy(
-                    dataset_name, export_root, project_type, output_dir, task,
-                    method.direction
-                )
-        elif method.dataset_format == "Supervisely":
-            if task == 'vector_annotation':
-                c_strategy = SuperviselyObjectDetectionStrategy(
-                    dataset_name, export_root, project_type, output_dir, task,
-                    method.direction
-                )
-
+    def _select_strategy(self, args):
+        if args.dataset_format == "COCO":
+            if args.task == 'instance_segmentation' or args.task == 'object_detection':
+                c_strategy = CocoObjectDetectionStrategy(args)
+            if args.task == 'keypoint_detection':
+                c_strategy = CocoKeypointDetectionStrategy(args)
+            if args.task == 'panoptic_segmentation':
+                c_strategy = CocoPanopticConverterStrategy(args)
+        elif args.dataset_format == "VOC":
+            if args.task == 'instance_segmentation' or args.task == 'object_detection':
+                c_strategy = VocObjectDetectionStrategy(args)
+        elif args.dataset_format == "LabelBox":
+            if args.task == "object_detection" or args.task == 'instance_segmentation' or args.task == 'vector_annotation':
+                c_strategy = LabelBoxObjectDetectionStrategy(args)
+        elif args.dataset_format == "DataLoop":
+            if args.task == 'object_detection' or args.task == 'instance_segmentation' or args.task == 'vector_annotation':
+                c_strategy = DataLoopObjectDetectionStrategy(args)
+        elif args.dataset_format == "Supervisely":
+            if args.task == 'vector_annotation':
+                c_strategy = SuperviselyObjectDetectionStrategy(args)
         else:
             pass
+
         self.__set_strategy(c_strategy)
 
     def _merge_jsons(self, input_dir):

superannotate/input_converters/converters/dataloop_converters/dataloop_converter.py

@@ -1,18 +1,17 @@
 class DataLoopConverter(object):
-    def __init__(
-        self, dataset_name_, exprot_root_, project_type_, output_dir_, task_
-    ):
-        self.dataset_name = dataset_name_
-        self.project_type = project_type_
-        self.task = task_
-        self.output_dir = output_dir_
-        self.export_root = exprot_root_
+    def __init__(self, args):
+        self.dataset_name = args.dataset_name
+        self.project_type = args.project_type
+        self.task = args.task
+        self.output_dir = args.output_dir
+        self.export_root = args.export_root
+        self.direction = args.direction
 
     def set_output_dir(self, output_dir_):
         self.output_dir = output_dir_
 
-    def set_export_root(self, exprot_root_):
-        self.exprot_root = exprot_root_
+    def set_export_root(self, export_root_):
+        self.export_root = export_root_
 
     def set_dataset_name(self, dname):
         self.dataset_name = dname

superannotate/input_converters/converters/dataloop_converters/dataloop_strategies.py

@@ -8,15 +8,8 @@
 class DataLoopObjectDetectionStrategy(DataLoopConverter):
     name = "Object Detection converter"
 
-    def __init__(
-        self, dataset_name, export_root, project_type, output_dir, task,
-        direction
-    ):
-        self.direction = direction
-        super().__init__(
-            dataset_name, export_root, project_type, output_dir, task
-        )
-
+    def __init__(self, args):
+        super().__init__(args)
         self.__setup_conversion_algorithm()
 
     def __setup_conversion_algorithm(self):

superannotate/input_converters/converters/labelbox_converters/labelbox_converter.py

@@ -1,18 +1,17 @@
 class LabelBoxConverter(object):
-    def __init__(
-        self, dataset_name_, exprot_root_, project_type_, output_dir_, task_
-    ):
-        self.dataset_name = dataset_name_
-        self.project_type = project_type_
-        self.task = task_
-        self.output_dir = output_dir_
-        self.export_root = exprot_root_
+    def __init__(self, args):
+        self.dataset_name = args.dataset_name
+        self.project_type = args.project_type
+        self.task = args.task
+        self.output_dir = args.output_dir
+        self.export_root = args.export_root
+        self.direction = args.direction
 
     def set_output_dir(self, output_dir_):
         self.output_dir = output_dir_
 
-    def set_export_root(self, exprot_root_):
-        self.exprot_root = exprot_root_
+    def set_export_root(self, export_root_):
+        self.export_root = export_root_
 
     def set_dataset_name(self, dname):
         self.dataset_name = dname