feat: LSDV-4831: Export BrushLabels to COCO

label_studio_converter/brush.py

@@ -38,6 +38,18 @@
 
 logger = logging.getLogger(__name__)
 
+try:
+    import pycocotools.mask
+except ImportError:
+    pycocotools_imported = False
+    logger.warning(
+        'pycocotools library import failed! You need to setup this library manually:\n'
+        + 'check this https://stackoverflow.com/questions/60189943/how-to-install-pycocotools-through-conda'
+    )
+else:
+    pycocotools_imported = True
+    logger.info('pycocotools library imported successfully!')
+
 
 ### Brush Export ###
 
@@ -188,6 +200,36 @@ def convert_task_dir(items, out_dir, out_format='numpy'):
         convert_task(item, out_dir, out_format)
 
 
+def binary_mask_to_rle(binary_mask: np.ndarray):
+    """from binary image mask to uncompressed coco rle"""
+    counts = []
+    for i, (value, elements) in enumerate(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):
+    """from LS rle to compressed coco rle"""
+    ls_mask = 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
+
+
+def get_cocomask_area(segmentation):
+    return int(pycocotools.mask.area(segmentation))
+
+
+def get_cocomask_bounding_box(segmentation):
+    return pycocotools.mask.toBbox(segmentation).tolist()
+
+
 # convert_task_dir('/ls/test/completions', '/ls/test/completions/output', 'numpy')
 
 

label_studio_converter/converter.py

@@ -1,5 +1,6 @@
 import os
 import io
+import itertools
 import math
 import logging
 import ujson as json
@@ -30,6 +31,8 @@
     get_annotator,
     get_json_root_type,
     prettify_result,
+    get_cocomask_area,
+    get_cocomask_bounding_box,
 )
 from label_studio_converter import brush
 from label_studio_converter.audio import convert_to_asr_json_manifest
@@ -314,8 +317,23 @@ def _get_supported_formats(self):
             or 'PolygonLabels' in output_tag_types
             and 'Labels' in output_tag_types
         ):
-            all_formats.remove(Format.COCO.name)
             all_formats.remove(Format.YOLO.name)
+        if not (
+            'Image' in input_tag_types
+            and (
+                'RectangleLabels' in output_tag_types
+                or 'PolygonLabels' in output_tag_types
+                or ('BrushLabels' in output_tag_types and brush.pycocotools_imported)
+                or ('brushlabels' in output_tag_types and brush.pycocotools_imported)
+            )
+            or 'Rectangle' in output_tag_types
+            and 'Labels' in output_tag_types
+            or 'PolygonLabels' in output_tag_types
+            and 'Labels' in output_tag_types
+            or ('Brush' in output_tag_types and brush.pycocotools_imported)
+            and 'Labels' in output_tag_types
+        ):
+            all_formats.remove(Format.COCO.name)
         if not (
             'Image' in input_tag_types
             and (
@@ -596,7 +614,12 @@ 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
@@ -659,6 +682,18 @@ def add_image(images, width, height, image_id, image_path):
                             'area': get_polygon_area(x, y),
                         }
                     )
+                elif 'brushlabels' in label and brush.pycocotools_imported:
+                    segmentation = brush.ls_rle_to_coco_rle(label["rle"], height, width)
+                    annotations.append(
+                        {
+                            "image_id": image_id,
+                            "segmentation": segmentation,
+                            "area": brush.get_cocomask_area(segmentation),
+                            "bbox": brush.get_cocomask_bounding_box(segmentation),
+                            "iscrowd": 1,
+                            "category_id": category_id,
+                        }
+                    )
                 else:
                     raise ValueError("Unknown label type")
 

label_studio_converter/utils.py

@@ -332,6 +332,14 @@ def get_polygon_bounding_box(x, y):
     return [x1, y1, x2 - x1, y2 - y1]
 
 
+def get_cocomask_area(segmentation):
+    return int(pycocotools.mask.area(segmentation))
+
+
+def get_cocomask_bounding_box(segmentation):
+    return pycocotools.mask.toBbox(segmentation).tolist()
+
+
 def get_annotator(item, default=None, int_id=False):
     """Get annotator id or email from annotation"""
     annotator = item['completed_by']

requirements.txt

@@ -5,4 +5,3 @@ nltk==3.6.7
 label-studio-tools>=0.0.2
 ujson
 ijson~=3.2.0.post0
-

setup.py

@@ -29,6 +29,11 @@
         'Operating System :: OS Independent',
     ],
     install_requires=requirements,
+    extras_require={
+        "coco": [
+            "pycocotools==2.0.6",
+        ],
+    },
     python_requires='>=3.6',
     entry_points={
         'console_scripts': [

tests/test_brush.py

@@ -1,11 +1,18 @@
 """
 Test for the brush.py module
 """
+import pytest
 import unittest
 import urllib3
 import json
+import numpy as np
 
-from label_studio_converter.brush import encode_rle, image2annotation
+from label_studio_converter.brush import (
+    encode_rle,
+    image2annotation,
+    binary_mask_to_rle,
+    ls_rle_to_coco_rle,
+)
 
 
 def test_image2annotation():
@@ -71,3 +78,31 @@ def test_rle_encoding():
         56,
         32,
     ]
+
+
+def test_binary_mask_to_rle():
+    # Test with a simple binary mask
+    binary_mask = np.array([[0, 1, 1, 0], [1, 1, 1, 0], [0, 0, 0, 1]])
+    rle = binary_mask_to_rle(binary_mask)
+    assert rle == {'counts': [1, 1, 1, 2, 1, 2, 3, 1], 'size': [3, 4]}
+
+    # Test with a binary mask that is all zeros
+    binary_mask = np.zeros((3, 4))
+    rle = binary_mask_to_rle(binary_mask)
+    assert rle == {'counts': [12], 'size': [3, 4]}
+
+    # Test with a binary mask that is all ones
+    binary_mask = np.ones((4, 5))
+    rle = binary_mask_to_rle(binary_mask)
+    assert rle == {'counts': [0, 20], 'size': [4, 5]}
+
+
+def test_ls_rle_to_coco_rle():
+    # Test with a simple LS RLE
+    pytest.importorskip("pycocotools")
+    ls_rle = encode_rle([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24])
+    height = 2
+    width = 3
+    coco_rle = ls_rle_to_coco_rle(ls_rle, height, width)
+    assert coco_rle == {'counts': '06', 'size': [2, 3]}
+