Merge pull request #178 from Labelbox/ms/mea-seg-metrics

mult-class seg iou

Author: msokoloff1

labelbox/data/metrics/iou.py

@@ -19,7 +19,7 @@
 
 def mask_miou(predictions: List[NDMask], labels: List[NDMask]) -> float:
     """
-    Creates prediction and label binary mask for all features with the same feature scheama id.
+    Creates prediction and label binary mask for all features with the same feature schema id.
 
     Args:
         predictions: List of masks objects
@@ -28,10 +28,18 @@ def mask_miou(predictions: List[NDMask], labels: List[NDMask]) -> float:
         float indicating iou score
     """
 
+    colors_pred = {tuple(pred.mask['colorRGB']) for pred in predictions}
+    colors_label = {tuple(label.mask['colorRGB']) for label in labels}
+    error_msg = "segmentation {} should all have the same color. Found {}"
+    if len(colors_pred) > 1:
+        raise ValueError(error_msg.format("predictions", colors_pred))
+    elif len(colors_label) > 1:
+        raise ValueError(error_msg.format("labels", colors_label))
+
     pred_mask = _instance_urls_to_binary_mask(
-        [pred.mask['instanceURI'] for pred in predictions])
+        [pred.mask['instanceURI'] for pred in predictions], colors_pred.pop())
     label_mask = _instance_urls_to_binary_mask(
-        [label.mask['instanceURI'] for label in labels])
+        [label.mask['instanceURI'] for label in labels], colors_label.pop())
     assert label_mask.shape == pred_mask.shape
     return _mask_iou(label_mask, pred_mask)
 
@@ -282,7 +290,13 @@ def _mask_iou(mask1: np.ndarray, mask2: np.ndarray) -> float:
     return np.sum(mask1 & mask2) / np.sum(mask1 | mask2)
 
 
-def _instance_urls_to_binary_mask(urls: List[str]) -> np.ndarray:
+def _remove_opacity_channel(masks: List[np.ndarray]) -> List[np.ndarray]:
+    return [mask[:, :, :3] if mask.shape[-1] == 4 else mask for mask in masks]
+
+
+def _instance_urls_to_binary_mask(urls: List[str],
+                                  color: Tuple[int, int, int]) -> np.ndarray:
     """Downloads segmentation masks and turns the image into a binary mask."""
-    masks = [url_to_numpy(url) for url in urls]
-    return np.sum(masks, axis=(0, 3)) > 0
+    masks = _remove_opacity_channel([url_to_numpy(url) for url in urls])
+    return np.sum([np.all(mask == color, axis=-1) for mask in masks],
+                  axis=0) > 0

labelbox/data/metrics/preprocess.py

@@ -1,3 +1,4 @@
+import functools
 from typing import List, Dict, Any, Union
 from collections import defaultdict
 import numpy as np  # type: ignore
@@ -35,7 +36,8 @@ def label_to_ndannotation(label: Dict[str, Any],
     if tool in SEGMENTATION_TOOLS:
         label['mask'] = {
             'instanceURI': label['instanceURI'],
-            'colorRGB': (0, 0, 0)
+            # Matches the color in the seg masks in the exports
+            'colorRGB': (255, 255, 255)
         }
     for unused_key in unused_keys:
         label.pop(unused_key, None)
@@ -66,7 +68,8 @@ def create_schema_lookup(rows: List[NDBase]) -> Dict[str, List[Any]]:
     return data
 
 
-@retry.Retry(deadline=10.)
+@retry.Retry(deadline=15.)
+@functools.lru_cache(maxsize=256)
 def url_to_numpy(mask_url: str) -> np.ndarray:
     """ Downloads an image and converts to a numpy array """
     return np.array(Image.open(BytesIO(requests.get(mask_url).content)))

tests/data/metrics/conftest.py

@@ -3,6 +3,7 @@
 import pytest
 import numpy as np
 from PIL import Image
+import base64
 
 
 class NameSpace(SimpleNamespace):
@@ -225,21 +226,22 @@ def unmatched_label():
                      expected=0.25)
 
 
-def create_mask_url(indices, h, w):
+def create_mask_url(indices, h, w, value):
     mask = np.zeros((h, w, 3), dtype=np.uint8)
     for idx in indices:
-        mask[idx] = 1
-    return mask.tobytes()
+        mask[idx] = value
+    return base64.b64encode(mask.tobytes()).decode('utf-8')
 
 
 @pytest.fixture
 def mask_pair():
-    #* Use your own signed urls so that you can resign the data
-    #* This is just to make the demo work
     return NameSpace(labels=[{
-        'featureId': '1234567890111213141516171',
-        'schemaId': 'ckppid25v0000aeyjmxfwlc7t',
-        'instanceURI': create_mask_url([(0, 0, 0), (0, 1, 0)], 32, 32)
+        'featureId':
+            '1234567890111213141516171',
+        'schemaId':
+            'ckppid25v0000aeyjmxfwlc7t',
+        'instanceURI':
+            create_mask_url([(0, 0), (0, 1)], 32, 32, (255, 255, 255))
     }],
                      predictions=[{
                          'uuid': '76e0dcea-fe46-43e5-95f5-a5e3f378520a',
@@ -249,7 +251,7 @@ def mask_pair():
                          },
                          'mask': {
                              'instanceURI':
-                                 create_mask_url([(0, 0, 0)], 32, 32),
+                                 create_mask_url([(0, 0)], 32, 32, (1, 1, 1)),
                              'colorRGB': (1, 1, 1)
                          }
                      }],

tests/data/metrics/test_iou.py

@@ -2,6 +2,7 @@
 from unittest.mock import patch
 import math
 import numpy as np
+import base64
 
 from labelbox.data.metrics.iou import datarow_miou
 
@@ -18,9 +19,9 @@ def test_overlapping(polygon_pair, box_pair, mask_pair):
     check_iou(polygon_pair)
     check_iou(box_pair)
     with patch('labelbox.data.metrics.iou.url_to_numpy',
-               side_effect=lambda x: np.frombuffer(x.encode('utf-8'),
-                                                   dtype=np.uint8).reshape(
-                                                       (32, 32, 3))):
+               side_effect=lambda x: np.frombuffer(
+                   base64.b64decode(x.encode('utf-8')), dtype=np.uint8).reshape(
+                       (32, 32, 3))):
         check_iou(mask_pair)