Merge pull request #239 from Labelbox/ms/mask-data

msokoloff1 · web-flow · commit b35dc30400fe · 2021-08-11T15:11:44.000-04:00
Add MaskData
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@
 * Geometry.raster now has a consistent interface and improved functionality
 * renamed schema_id to feature_schema_id in the `FeatureSchema` class
 * ImageData can load 2d images or masks
+* Mask objects now use `MaskData` to represent segmentation masks instead of `ImageData`
 
 # Version 3.0.0-rc2 (2021-08-09)
 ## Updates
diff --git a/examples/annotation_types/basics.ipynb b/examples/annotation_types/basics.ipynb
@@ -56,6 +56,7 @@
     "from labelbox.data.annotation_types import (\n",
     "    Label,\n",
     "    ImageData,\n",
+    "    MaskData,\n",
     "    LabelList,\n",
     "    TextData,\n",
     "    VideoData,\n",
@@ -551,12 +552,12 @@
    "source": [
     "\n",
     "mask_annotation = ObjectAnnotation(\n",
-    "        value = Mask(mask =  ImageData(arr = np.zeros((32,32,3), dtype = np.uint8)), color = (255,255,255)),\n",
+    "        value = Mask(mask =  MaskData(arr = np.zeros((32,32,3), dtype = np.uint8)), color = (255,255,255)),\n",
     "            name = \"mask class name\"\n",
     ")\n",
     "\n",
     "mask_annotation = ObjectAnnotation(\n",
-    "        value = Mask(mask =  ImageData(arr = np.zeros((32,32,3), dtype = np.uint8)), color = (255,255,255)),\n",
+    "        value = Mask(mask =  MaskData(arr = np.zeros((32,32,3), dtype = np.uint8)), color = (255,255,255)),\n",
     "            feature_schema_id = \"ckrgcel71000008jtd9mn0szu\"\n",
     ")"
    ]
@@ -577,7 +578,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "raster_data = ImageData(arr = np.zeros((32,32,3), dtype = np.uint8))\n",
+    "raster_data = MaskData(arr = np.zeros((32,32,3), dtype = np.uint8))\n",
     "mask_annotation = ObjectAnnotation(\n",
     "        value = Mask(mask  = raster_data, color = [255,255,255]),\n",
     "            name = \"eyes\"\n",
@@ -604,7 +605,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "raster_data = ImageData(arr = np.zeros((32,32, 3), dtype = np.uint8))\n",
+    "raster_data = MaskData(arr = np.zeros((32,32, 3), dtype = np.uint8))\n",
     "mask_annotation = ObjectAnnotation(\n",
     "        value = Mask(mask  = raster_data, color = (128,255,255)),\n",
     "            name = \"eye\"\n",
@@ -734,7 +735,7 @@
    "outputs": [],
    "source": [
     "mask_annotation = Mask(\n",
-    "    mask = ImageData(arr = np_mask),\n",
+    "    mask = MaskData(arr = np_mask),\n",
     "    color = color \n",
     ")\n",
     "\n",
@@ -790,7 +791,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "mask_data = ImageData(arr = np_seg_mask)\n",
+    "mask_data = MaskData(arr = np_seg_mask)\n",
     "eye_mask = Mask(mask = mask_data, color = eye_color)\n",
     "nose_mask = Mask(mask = mask_data, color = nose_color)\n"
    ]
@@ -954,7 +955,7 @@
     "        ),\n",
     "        ObjectAnnotation(\n",
     "            name = \"deer_eyes\",\n",
-    "            value = Mask(mask = ImageData(arr = np_mask), color = color)\n",
+    "            value = Mask(mask = MaskData(arr = np_mask), color = color)\n",
     "        )\n",
     "    ]\n",
     ")"
@@ -1283,6 +1284,22 @@
    "metadata": {},
    "outputs": [],
    "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "progressive-courage",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "asian-album",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
diff --git a/examples/annotation_types/converters.ipynb b/examples/annotation_types/converters.ipynb
diff --git a/examples/annotation_types/label_containers.ipynb b/examples/annotation_types/label_containers.ipynb
@@ -36,7 +36,8 @@
     "from labelbox import LabelingFrontend\n",
     "from labelbox.data.annotation_types import (\n",
     "    Label, \n",
-    "    ImageData, \n",
+    "    ImageData,\n",
+    "    MaskData,\n",
     "    Mask, \n",
     "    Point, \n",
     "    Polygon, \n",
@@ -137,7 +138,7 @@
     "        Point(x = 99, y = 181).raster(im_h, im_w, thickness = 3),\n",
     "    ]\n",
     "    mask_arr = np.max([*eye_masks,nose_mask] , axis = 0)\n",
-    "    mask = ImageData(arr = mask_arr)\n",
+    "    mask = MaskData(arr = mask_arr)\n",
     "    return [Label(\n",
     "        data = ImageData(im_bytes = requests.get(image_url).content),\n",
     "        annotations = [\n",
@@ -572,6 +573,14 @@
    "metadata": {},
    "outputs": [],
    "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "foreign-berkeley",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
diff --git a/examples/annotation_types/mal_using_annotation_types.ipynb b/examples/annotation_types/mal_using_annotation_types.ipynb
@@ -67,6 +67,7 @@
     "from labelbox.data.annotation_types import (\n",
     "    LabelList,\n",
     "    ImageData,\n",
+    "    MaskData,\n",
     "    Rectangle,\n",
     "    ObjectAnnotation,\n",
     "    ClassificationAnnotation,\n",
@@ -170,7 +171,7 @@
     "                start = Point(x = box[1], y = box[0]), end = Point(x = box[3], y = box[2])\n",
     "            )\n",
     "        elif name == 'car':\n",
-    "            value = Mask(mask = ImageData.from_2D_arr(arr = seg), color = (1,1,1))\n",
+    "            value = Mask(mask = MaskData.from_2D_arr(arr = seg), color = (1,1,1))\n",
     "        if value is not None:\n",
     "            annotations.append(\n",
     "                ObjectAnnotation(\n",
@@ -322,22 +323,6 @@
     "for status in upload_task.statuses[:5]:\n",
     "    print(status)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "micro-salem",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "prospective-grave",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
diff --git a/examples/model_diagnostics/model_diagnostics_demo.ipynb b/examples/model_diagnostics/model_diagnostics_demo.ipynb
@@ -136,6 +136,7 @@
     "    LabelList, \n",
     "    Label, \n",
     "    ImageData, \n",
+    "    MaskData,\n",
     "    Mask, \n",
     "    Polygon,\n",
     "    Point, \n",
@@ -270,7 +271,7 @@
     "            elif class_info['kind'] == Tool.Type.POINT:\n",
     "                value = Point(x=(box[1] + box[3]) / 2., y = (box[0] + box[2]) / 2.)\n",
     "            elif class_info['kind'] == Tool.Type.SEGMENTATION:\n",
-    "                value = Mask(mask = ImageData.from_2D_arr(seg * class_info['color']), color = (class_info['color'],)* 3)\n",
+    "                value = Mask(mask = MaskData.from_2D_arr(seg * class_info['color']), color = (class_info['color'],)* 3)\n",
     "            else:\n",
     "                raise ValueError(f\"Unsupported kind found. {class_info['kind']}\")\n",
     "            annotations.append(ObjectAnnotation(name = class_info['name'], value = value))\n",
diff --git a/examples/model_diagnostics/model_diagnostics_guide.ipynb b/examples/model_diagnostics/model_diagnostics_guide.ipynb
@@ -138,6 +138,7 @@
     "    LabelList, \n",
     "    Label, \n",
     "    ImageData, \n",
+    "    MaskData,\n",
     "    Mask,\n",
     "    Polygon,\n",
     "    Line,\n",
@@ -271,7 +272,7 @@
     "        value = Rectangle(start = Point(x = xmin, y = ymin), end = Point(x=xmax, y=ymax))\n",
     "        value = Point(x=x, y =y)\n",
     "        value = Line(points = [Point(x = x, y = y) for x,y in instance])\n",
-    "        value = Mask(mask = ImageData.from_2D_arr(seg * grayscale_color), color = (grayscale_color,)* 3)\n",
+    "        value = Mask(mask = MaskData.from_2D_arr(seg * grayscale_color), color = (grayscale_color,)* 3)\n",
     "        \n",
     "        annotations.append(ObjectAnnotation(name =class_name, value = value))\n",
     "            \n",
diff --git a/labelbox/data/annotation_types/__init__.py b/labelbox/data/annotation_types/__init__.py
@@ -19,6 +19,7 @@
 from .classification import Text
 
 from .data import ImageData
+from .data import MaskData
 from .data import TextData
 from .data import VideoData
 
diff --git a/labelbox/data/annotation_types/collection.py b/labelbox/data/annotation_types/collection.py
@@ -77,7 +77,7 @@ def add_to_dataset(self,
     def add_url_to_masks(self, signer, max_concurrency=20) -> "LabelList":
         """
         Creates signed urls for all masks in the LabelList.
-        Multiple masks can reference the same ImageData mask so this makes sure we only upload that url once.
+        Multiple masks objects can reference the same MaskData so this makes sure we only upload that url once.
         Only uploads url if one doesn't already exist.
 
         Args:
@@ -235,7 +235,7 @@ def add_url_to_masks(self, signer: Callable[[bytes],
                                                 str]) -> "LabelGenerator":
         """
         Creates signed urls for all masks in the LabelGenerator.
-        Multiple masks can reference the same ImageData mask so this makes sure we only upload that url once.
+        Multiple masks can reference the same MaskData so this makes sure we only upload that url once.
         Only uploads url if one doesn't already exist.
 
         Args:
diff --git a/labelbox/data/annotation_types/data/__init__.py b/labelbox/data/annotation_types/data/__init__.py
@@ -1,3 +1,4 @@
 from .raster import ImageData
+from .raster import MaskData
 from .text import TextData
 from .video import VideoData
diff --git a/labelbox/data/annotation_types/data/raster.py b/labelbox/data/annotation_types/data/raster.py
@@ -2,6 +2,7 @@
 from io import BytesIO
 
 import numpy as np
+from pydantic import BaseModel
 import requests
 from google.api_core import retry
 from typing_extensions import Literal
@@ -12,7 +13,7 @@
 from ..types import TypedArray
 
 
-class ImageData(BaseData):
+class RasterData(BaseModel):
     """
     Represents an image or segmentation mask.
     """
@@ -52,8 +53,10 @@ def np_to_bytes(self, arr: np.ndarray) -> bytes:
             png encoded bytes
         """
         if len(arr.shape) != 3:
-            raise ValueError("unsupported image format. Must be 3D ([H,W,C])."
-                             "Use ImageData.from_2D_arr to construct from 2D")
+            raise ValueError(
+                "unsupported image format. Must be 3D ([H,W,C])."
+                f"Use {self.__class__.__name__}.from_2D_arr to construct from 2D"
+            )
         if arr.dtype != np.uint8:
             raise TypeError(f"image data type must be uint8. Found {arr.dtype}")
 
@@ -143,12 +146,12 @@ def validate_args(cls, values):
             elif len(arr.shape) != 3:
                 raise ValueError(
                     "unsupported image format. Must be 3D ([H,W,C])."
-                    "Use ImageData.from_2D_arr to construct from 2D")
+                    f"Use {cls.__name__}.from_2D_arr to construct from 2D")
         return values
 
     def __repr__(self) -> str:
         symbol_or_none = lambda data: '...' if data is not None else None
-        return  f"ImageData(im_bytes={symbol_or_none(self.im_bytes)}," \
+        return  f"{self.__class__.__name__}(im_bytes={symbol_or_none(self.im_bytes)}," \
                 f"file_path={self.file_path}," \
                 f"url={self.url}," \
                 f"arr={symbol_or_none(self.arr)})"
@@ -158,3 +161,11 @@ class Config:
         copy_on_model_validation = False
         # Required for discriminating between data types
         extra = 'forbid'
+
+
+class MaskData(RasterData):
+    ...
+
+
+class ImageData(RasterData, BaseData):
+    ...
diff --git a/labelbox/data/annotation_types/geometry/mask.py b/labelbox/data/annotation_types/geometry/mask.py
@@ -6,13 +6,13 @@
 from shapely.geometry import MultiPolygon, shape
 import cv2
 
-from ..data.raster import ImageData
+from ..data import MaskData
 from .geometry import Geometry
 
 
 class Mask(Geometry):
     # Raster data can be shared across multiple masks... or not
-    mask: ImageData
+    mask: MaskData
     # RGB or Grayscale
     color: Tuple[int, int, int]
 
diff --git a/labelbox/data/annotation_types/label.py b/labelbox/data/annotation_types/label.py
@@ -63,7 +63,7 @@ def add_url_to_data(self, signer) -> "Label":
     def add_url_to_masks(self, signer) -> "Label":
         """
         Creates signed urls for all masks in the Label.
-        Multiple masks can reference the same ImageData mask so this makes sure we only upload that url once.
+        Multiple masks can reference the same MaskData mask so this makes sure we only upload that url once.
         Only uploads url if one doesn't already exist.
 
         Args:
diff --git a/labelbox/data/serialization/labelbox_v1/objects.py b/labelbox/data/serialization/labelbox_v1/objects.py
@@ -4,7 +4,7 @@
 
 from ...annotation_types.annotation import (ClassificationAnnotation,
                                             ObjectAnnotation)
-from ...annotation_types.data import ImageData
+from ...annotation_types.data import MaskData
 from ...annotation_types.geometry import Line, Mask, Point, Polygon, Rectangle
 from ...annotation_types.ner import TextEntity
 from ...annotation_types.types import Cuid
@@ -122,7 +122,7 @@ class LBV1Mask(LBV1ObjectBase):
     instanceURI: str
 
     def to_common(self) -> Mask:
-        return Mask(mask=ImageData(url=self.instanceURI), color=(255, 255, 255))
+        return Mask(mask=MaskData(url=self.instanceURI), color=(255, 255, 255))
 
     @classmethod
     def from_common(cls, mask: Mask,
diff --git a/labelbox/data/serialization/ndjson/objects.py b/labelbox/data/serialization/ndjson/objects.py
@@ -2,7 +2,7 @@
 
 from pydantic import BaseModel
 
-from ...annotation_types.data import ImageData, TextData
+from ...annotation_types.data import ImageData, TextData, MaskData
 from ...annotation_types.ner import TextEntity
 from ...annotation_types.types import Cuid
 from ...annotation_types.geometry import Rectangle, Polygon, Line, Point, Mask
@@ -122,7 +122,7 @@ class NDMask(NDBaseObject):
     mask: _Mask
 
     def to_common(self) -> Mask:
-        return Mask(mask=ImageData(url=self.mask.instanceURI),
+        return Mask(mask=MaskData(url=self.mask.instanceURI),
                     color=self.mask.colorRGB)
 
     @classmethod
diff --git a/tests/data/annotation_types/geometry/test_mask.py b/tests/data/annotation_types/geometry/test_mask.py
@@ -4,7 +4,7 @@
 import numpy as np
 import cv2
 
-from labelbox.data.annotation_types import ImageData, Point, Rectangle, Mask
+from labelbox.data.annotation_types import Point, Rectangle, Mask, MaskData
 
 
 def test_mask():
@@ -14,7 +14,7 @@ def test_mask():
     mask_data = np.zeros((32, 32, 3), dtype=np.uint8)
     mask_data = cv2.rectangle(mask_data, (0, 0), (10, 10), (255, 255, 255), -1)
     mask_data = cv2.rectangle(mask_data, (20, 20), (30, 30), (0, 255, 255), -1)
-    mask_data = ImageData(arr=mask_data)
+    mask_data = MaskData(arr=mask_data)
 
     mask1 = Mask(mask=mask_data, color=(255, 255, 255))
 
diff --git a/tests/data/annotation_types/test_collection.py b/tests/data/annotation_types/test_collection.py
