coco converter new json

Author: agyurjin

superannotate/input_converters/converters/coco_converters/coco_api.py

@@ -1,4 +1,5 @@
 import numpy as np
+import cv2
 
 
 def encode(bitmask):
@@ -94,3 +95,47 @@ def _frString(rle_string):
         counts.append(count)
 
     return counts
+
+
+def _area(bitmask):
+    return np.sum(bitmask)
+
+
+def _toBbox(bitmask):
+    contours, _ = cv2.findContours(
+        bitmask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+    )
+    segments = []
+    for contour in contours:
+        contour = contour.flatten().tolist()
+        segments += contour
+
+    xmin = min(segments[::2])
+    xmax = max(segments[::2])
+    ymin = min(segments[1::2])
+    ymax = max(segments[1::2])
+
+    return [xmin, ymin, xmax - xmin, ymax - ymin]
+
+
+def _merge(list_of_bitmask):
+    shape = list_of_bitmask[0].shape
+    final_bitmask = np.zeros(shape, dtype=np.uint8)
+    for bitmask in list_of_bitmask:
+        final_bitmask |= bitmask
+
+    return final_bitmask
+
+
+def _polytoMask(polygons, height, width):
+    masks = []
+    for polygon in polygons:
+        polygon = np.round(polygon)
+        bitmask = np.zeros((height, width)).astype(np.uint8)
+        pts = np.array(
+            [polygon[2 * i:2 * (i + 1)] for i in range(len(polygon) // 2)],
+            dtype=np.int32
+        )
+        cv2.fillPoly(bitmask, [pts], 1)
+        masks.append(bitmask)
+    return masks

superannotate/input_converters/converters/coco_converters/coco_strategies.py

@@ -95,10 +95,14 @@ def _prepare_single_image_commons_pixel(self, id_, json_path):
         )
         rm_len = len('___pixel.json')
 
-        sa_ann_json = json.load(open(json_path))
+        sa_json = json.load(open(json_path))
+        # sa_ann_json = json.load(open(json_path))
+        sa_ann_json = sa_json['instances']
+
         sa_bluemask_path = str(json_path)[:-rm_len] + '___save.png'
 
-        image_info = self.__make_image_info(json_path, id_, self.project_type)
+        # image_info = self.__make_image_info(json_path, id_, self.project_type)
+        image_info = self.__make_image_info(id_, sa_json['metadata'])
 
         sa_bluemask_rgb = np.asarray(
             Image.open(sa_bluemask_path).convert('RGB'), dtype=np.uint32
@@ -117,20 +121,21 @@ def _prepare_single_image_commons_pixel(self, id_, json_path):
 
         return res
 
-    def __make_image_info(self, json_path, id_, source_type):
-        if source_type == 'Pixel':
-            rm_len = len('___pixel.json')
-        elif source_type == 'Vector':
-            rm_len = len('___objects.json')
+    # def __make_image_info(self, json_path, id_, source_type):
+    # if source_type == 'Pixel':
+    #     rm_len = len('___pixel.json')
+    # elif source_type == 'Vector':
+    #     rm_len = len('___objects.json')
 
-        image_path = str(json_path)[:-rm_len]
+    # image_path = str(json_path)[:-rm_len]
 
-        img_width, img_height = Image.open(image_path).size
+    # img_width, img_height = Image.open(image_path).size
+    def __make_image_info(self, id_, sa_meta_json):
         image_info = {
             'id': id_,
-            'file_name': Path(image_path).name,
-            'height': img_height,
-            'width': img_width,
+            'file_name': sa_meta_json['name'],
+            'height': sa_meta_json['height'],
+            'width': sa_meta_json['width'],
             'license': 1
         }
 
@@ -139,8 +144,12 @@ def __make_image_info(self, json_path, id_, source_type):
     def _prepare_single_image_commons_vector(self, id_, json_path):
         ImgCommons = namedtuple('ImgCommons', ['image_info', 'sa_ann_json'])
 
-        image_info = self.__make_image_info(json_path, id_, self.project_type)
-        sa_ann_json = json.load(open(json_path))
+        # sa_ann_json = json.load(open(json_path))
+        sa_json = json.load(open(json_path))
+        sa_ann_json = sa_json['instances']
+
+        # image_info = self.__make_image_info(json_path, id_, self.project_type)
+        image_info = self.__make_image_info(id_, sa_json['metadata'])
 
         res = ImgCommons(image_info, sa_ann_json)
 
@@ -164,7 +173,6 @@ def _create_sa_classes(self, json_path):
             hexcolor = "#%02x%02x%02x" % tuple(color)
             classes_dict = {
                 'name': data["name"],
-                # 'id': data["id"],
                 'color': hexcolor,
                 'attribute_groups': []
             }
@@ -204,7 +212,7 @@ def sa_to_output_format(self):
         id_generator = self._make_id_generator()
 
         jsons_gen = self.export_root.glob('*pixel.json')
-        jsons = [path for path in jsons_gen]
+        jsons = list(jsons_gen)
 
         for id_, json_ in tqdm(enumerate(jsons, 1)):
             res = self._sa_to_coco_single(id_, json_, id_generator, cat_id_map)

superannotate/input_converters/converters/coco_converters/coco_to_sa_pixel.py

@@ -8,11 +8,11 @@
 
 from .coco_api import (_maskfrRLE, decode)
 
-from ..sa_json_helper import (_create_pixel_instance, _create_empty_sa_json)
+from ..sa_json_helper import (_create_pixel_instance, _create_sa_json)
 
 from ....common import blue_color_generator, hex_to_rgb, id2rgb, write_to_json
 
-logger = logging.getLogger("superannotate-python-sdk")
+logger = logging.getLogger("superannot-python-sdk")
 
 
 def annot_to_bitmask(annot):
@@ -27,14 +27,20 @@ def annot_to_bitmask(annot):
 def coco_panoptic_segmentation_to_sa_pixel(coco_path, output_dir):
     coco_json = json.load(open(coco_path))
     hex_colors = blue_color_generator(len(coco_json["categories"]))
-    annotate_list = coco_json["annotations"]
 
     cat_id_to_cat = {}
     for cat in coco_json['categories']:
         cat_id_to_cat[cat['id']] = cat['name']
 
-    for annotate in tqdm(annotate_list, "Converting annotations"):
-        annot_name = Path(annotate["file_name"]).stem
+    img_id_to_shape = {}
+    for img in coco_json['images']:
+        img_id_to_shape[str(img['id'])] = {
+            'height': img['height'],
+            'width': img['width']
+        }
+
+    for annot in tqdm(coco_json["annotations"], "Converting annotations"):
+        annot_name = Path(annot["file_name"]).stem
         img_cv = cv2.imread(str(output_dir / ("%s.png" % annot_name)))
         if img_cv is None:
             logger.warning(
@@ -47,7 +53,7 @@ def coco_panoptic_segmentation_to_sa_pixel(coco_path, output_dir):
         img = cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)
         H, W, C = img.shape
         img = img.reshape((H * W, C))
-        segments = annotate["segments_info"]
+        segments = annot["segments_info"]
         hex_colors = blue_color_generator(len(segments))
 
         sa_instances = []
@@ -64,7 +70,13 @@ def coco_panoptic_segmentation_to_sa_pixel(coco_path, output_dir):
         cv2.imwrite(str(output_dir / ("%s.jpg___save.png" % annot_name)), img)
 
         file_name = "%s.jpg___pixel.json" % annot_name
-        write_to_json(output_dir / file_name, sa_instances)
+        sa_metadata = {
+            'name': annot_name,
+            'width': img_id_to_shape[str(annot['image_id'])]['width'],
+            'height': img_id_to_shape[str(annot['image_id'])]['height']
+        }
+        json_template = _create_sa_json(sa_instances, metadata)
+        write_to_json(output_dir / file_name, json_template)
         (output_dir / ("%s.png" % annot_name)).unlink()
 
 
@@ -119,14 +131,13 @@ def coco_instance_segmentation_to_sa_pixel(coco_path, output_dir):
             )
             sa_instances.append(sa_obj)
 
-        write_to_json(output_dir / file_name, sa_instances)
-        # json_template['instance'] = sa_instances
-        # json_template['metadata'] = {
-        #     'name': images_dict['file_name'],
-        #     'width': images_dict['shape'][1],
-        #     'height': images_dict['shape'][0]
-        # }
-        # write_to_json(output_dir / file_name, json_template)
+        sa_metadata = {
+            'name': images_dict[id_]['file_name'],
+            'width': images_dict[id_]['shape'][1],
+            'height': images_dict[id_]['shape'][0]
+        }
+        json_template = _create_sa_json(sa_instances, sa_metadata)
+        write_to_json(output_dir / file_name, json_template)
         cv2.imwrite(
             str(output_dir / ('%s___save.png' % annotations['file_name'])), mask
         )

superannotate/input_converters/converters/coco_converters/coco_to_sa_vector.py

@@ -11,7 +11,7 @@
 
 from .coco_api import (_maskfrRLE, decode)
 
-from ..sa_json_helper import (_create_vector_instance, _create_empty_sa_json)
+from ..sa_json_helper import (_create_vector_instance, _create_sa_json)
 
 from ....common import id2rgb, write_to_json
 
@@ -36,7 +36,6 @@ def annot_to_polygon(annot):
 
 def save_sa_jsons(coco_json, img_id_to_annot, output_dir):
     for img in tqdm(coco_json['images'], "Writing annotations to disk"):
-        # json_template = _create_empty_sa_json()
         if 'file_name' in img:
             image_path = Path(img['file_name']).name
         else:
@@ -47,15 +46,14 @@ def save_sa_jsons(coco_json, img_id_to_annot, output_dir):
         else:
             sa_instances = img_id_to_annot[str(img['id'])]
         file_name = "%s___objects.json" % image_path
-        write_to_json(output_dir / file_name, sa_instances)
-
-        # json_template['instance'] = sa_instances
-        # json_template['metadata'] = {
-        #     'name': image_path,
-        #     'width': img['width'],
-        #     'height': img['height']
-        # }
-        # write_to_json(output_dir / file_name, json_template)
+
+        sa_metadata = {
+            'name': image_path,
+            'width': img['width'],
+            'height': img['height']
+        }
+        json_template = _create_sa_json(sa_instances, sa_metadata)
+        write_to_json(output_dir / file_name, json_template)
 
 
 def coco_instance_segmentation_to_sa_vector(coco_path, output_dir):

superannotate/input_converters/converters/coco_converters/sa_pixel_to_coco.py

@@ -1,7 +1,8 @@
 import cv2 as cv
 import numpy as np
 
-from ....pycocotools_sa import mask as cocomask
+from .coco_api import (encode, _toBbox, _area)
+# from ....pycocotools_sa import mask as cocomask
 
 
 def __instance_object_commons_per_instance(
@@ -22,10 +23,13 @@ def __instance_object_commons_per_instance(
     contours, _ = cv.findContours(
         databytes, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_NONE
     )
-    coco_instance_mask = cocomask.encode(np.asfortranarray(instance_bitmask))
 
-    bbox = cocomask.toBbox(coco_instance_mask).tolist()
-    area = int(cocomask.area(coco_instance_mask))
+    # coco_instance_mask = cocomask.encode(np.asfortranarray(instance_bitmask))
+    # coco_instance_mask = encode(instance_bitmask.astype(np.uint8))
+    # bbox = cocomask.toBbox(coco_instance_mask).tolist()
+    bbox = list(_toBbox(instance_bitmask.astype(np.uint8)))
+    # area = int(cocomask.area(coco_instance_mask))
+    area = int(_area(instance_bitmask.astype(np.uint8)))
     return (bbox, area, contours, category_id, anno_id)
 
 
@@ -92,11 +96,15 @@ def sa_pixel_to_coco_panoptic_segmentation(
         instance_bitmask = np.isin(flat_mask, parts)
         segment_id = next(id_generator)
         ann_mask[instance_bitmask] = segment_id
-        coco_instance_mask = cocomask.encode(
-            np.asfortranarray(instance_bitmask)
-        )
-        bbox = cocomask.toBbox(coco_instance_mask).tolist()
-        area = int(cocomask.area(coco_instance_mask))
+        # coco_instance_mask = cocomask.encode(
+        #     np.asfortranarray(instance_bitmask)
+        # )
+        # bbox = cocomask.toBbox(coco_instance_mask).tolist()
+        # area = int(cocomask.area(coco_instance_mask))
+
+        # coco_instance_mask = encode(instance_bitmask.astype(np.uint8))
+        bbox = list(_toBbox(instance_bitmask.astype(np.uint8)))
+        area = int(_area(instance_bitmask.astype(np.uint8)))
 
         segment_info = {
             'id': segment_id,

superannotate/input_converters/converters/coco_converters/sa_vector_to_coco.py

@@ -1,6 +1,8 @@
 import json
 import logging
 
+from .coco_api import (encode, _toBbox, _merge, _area, _polytoMask)
+
 from ....pycocotools_sa import mask as cocomask
 
 logger = logging.getLogger("superannotate-python-sdk")
@@ -80,12 +82,19 @@ def sa_vector_to_coco_instance_segmentation(
         for cat_id, polygons in polygon_group.items():
             anno_id = next(id_generator)
             try:
-                masks = cocomask.frPyObjects(
+                # masks = cocomask.frPyObjects(
+                #     polygons, image_info['height'], image_info['width']
+                # )
+                # mask = cocomask.merge(masks)
+                # area = int(cocomask.area(mask))
+                # bbox = cocomask.toBbox(mask).tolist()
+                masks = _polytoMask(
                     polygons, image_info['height'], image_info['width']
                 )
-                mask = cocomask.merge(masks)
-                area = int(cocomask.area(mask))
-                bbox = cocomask.toBbox(mask).tolist()
+                mask = _merge(masks)
+                area = int(_area(mask))
+                bbox = list(_toBbox(mask))
+
                 annotation = make_annotation(
                     cat_id, image_info['id'], bbox, polygons, area, anno_id
                 )