[feat] export BrushLabels to COCO

label_studio_converter/converter.py

@@ -1,6 +1,7 @@
 import os
 import json
 import io
+import itertools
 import math
 import logging
 import pandas as pd
@@ -16,6 +17,8 @@
 from operator import itemgetter
 from copy import deepcopy
 from PIL import Image
+import numpy as np
+import pycocotools.mask
 
 from label_studio_converter.utils import (
     parse_config, create_tokens_and_tags, download, get_image_size, get_image_size_and_channels, ensure_dir,
@@ -27,6 +30,26 @@
 logger = logging.getLogger(__name__)
 
 
+def binary_mask_to_rle(binary_mask: np.ndarray):
+    counts = []
+    for i, (value, elements) in enumerate(itertools.groupby(binary_mask.ravel(order='F'))):
+        if i == 0 and value == 1:
+            counts.append(0)
+        counts.append(len(list(elements)))
+    return {'counts': counts, 'size': list(binary_mask.shape)}
+
+
+def ls_rle_to_coco_rle(ls_rle, height, width):
+    ls_mask = brush.decode_rle(ls_rle)
+    ls_mask = np.reshape(ls_mask, [height, width, 4])[:, :, 3]
+    ls_mask = np.where(ls_mask > 0, 1, 0)
+    binary_mask = np.asfortranarray(ls_mask)
+    coco_rle = binary_mask_to_rle(binary_mask)
+    result = pycocotools.mask.frPyObjects(coco_rle, *coco_rle.get('size'))
+    result["counts"] = result["counts"].decode()
+    return result
+
+
 class FormatNotSupportedError(NotImplementedError):
     pass
 
@@ -531,7 +554,7 @@ def add_image(images, width, height, image_id, image_path):
             for label in labels:
 
                 category_name = None
-                for key in ['rectanglelabels', 'polygonlabels', 'labels']:
+                for key in ['rectanglelabels', 'polygonlabels', 'brushlabels' 'labels']:
                     if key in label and len(label[key]) > 0:
                         category_name = label[key][0]
                         break
@@ -584,6 +607,16 @@ def add_image(images, width, height, image_id, image_path):
                         'iscrowd': 0,
                         'area': get_polygon_area(x, y)
                     })
+                elif 'brushlabels' in label:
+                    segmentation = ls_rle_to_coco_rle(label["rle"], height, width)
+                    annotations.append({
+                        "image_id": image_id,
+                        "segmentation": segmentation,
+                        "area": int(pycocotools.mask.area(segmentation)),
+                        "bbox": pycocotools.mask.toBbox(segmentation).tolist(),
+                        "iscrowd": 1,
+                        "category_id": category_id,
+                    })
                 else:
                     raise ValueError("Unknown label type")