feat: add "RLESubcatalog" & "LabeledRLE" & "RLE" class

PR Closed: #840

Author: zhen.chen

tensorbay/geometry/__init__.py

@@ -7,7 +7,7 @@
 
 from .box import Box2D, Box3D
 from .keypoint import Keypoint2D, Keypoints2D
-from .polygon import MultiPolygon, Polygon
+from .polygon import RLE, MultiPolygon, Polygon
 from .polyline import MultiPolyline2D, Polyline2D
 from .transform import Transform3D
 from .vector import Vector, Vector2D, Vector3D
@@ -21,6 +21,7 @@
     "Polyline2D",
     "MultiPolygon",
     "MultiPolyline2D",
+    "RLE",
     "Transform3D",
     "Vector",
     "Vector2D",

tensorbay/geometry/polygon.py

@@ -12,7 +12,7 @@
 
 from typing import Dict, Iterable, List, Optional, Type, TypeVar
 
-from ..utility import common_loads
+from ..utility import UserMutableSequence, common_loads
 from .point_list import MultiPointList2D, PointList2D
 from .vector import Vector2D
 
@@ -151,3 +151,70 @@ def dumps(self) -> List[List[Dict[str, float]]]:
 
         """
         return self._dumps()
+
+
+class RLE(UserMutableSequence[int]):
+    """This class defines the concept of RLE.
+
+    :class:`RLE` contains an rle format mask.
+
+    Arguments:
+        rle: A rle format mask.
+
+    Examples:
+        >>> RLE([272, 2, 4, 4, 2, 9])
+        RLE [
+          272,
+          2,
+          ...
+        ]
+
+    """
+
+    _data: List[int]
+
+    def __init__(self, rle: Optional[Iterable[int]]):
+        self._data = list(rle) if rle else []
+
+    def _dumps(self) -> List[int]:
+        return self._data
+
+    def _loads(self, contents: List[int]) -> None:
+        self._data = contents
+
+    @classmethod
+    def loads(cls: Type["RLE"], contents: List[int]) -> "RLE":
+        """Loads a :class:RLE` from the given contents.
+
+        Arguments:
+            contents: One rle mask.
+
+        Returns:
+            The loaded :class:`RLE` object.
+
+        Examples:
+            >>> contents = [272, 2, 4, 4, 2, 9]
+            >>> rle = RLE.loads(contents)
+            >>> rle
+            RLE [
+              272,
+              2,
+              ...
+            ]
+
+        """
+        return common_loads(cls, contents)
+
+    def dumps(self) -> List[int]:
+        """Dumps a :class:`RLE` into one rle mask.
+
+        Returns:
+            All the information of the :class:`RLE`.
+
+        Examples:
+            >>> rle = RLE([272, 2, 4, 4, 2, 9])
+            >>> rle.dumps()
+            [272, 2, 4, 4, 2, 9]
+
+        """
+        return self._dumps()

tensorbay/geometry/tests/test_polygon.py

@@ -3,13 +3,16 @@
 # Copyright 2021 Graviti. Licensed under MIT License.
 #
 
-from .. import Box2D, MultiPolygon, Polygon, Vector2D
+from ..box import Box2D
+from ..polygon import RLE, MultiPolygon, Polygon
+from ..vector import Vector2D
 
 _DATA_POLYGON = [{"x": 1.0, "y": 1.0}, {"x": 2.0, "y": 2.0}, {"x": 2.0, "y": 3.0}]
 _DATA_MULTIPOLYGON = [
     [{"x": 1.0, "y": 4.0}, {"x": 2.0, "y": 3.7}, {"x": 7.0, "y": 4.0}],
     [{"x": 5.0, "y": 7.0}, {"x": 6.0, "y": 7.0}, {"x": 9.0, "y": 8.0}],
 ]
+_DATA_RLE = [272, 2, 4, 4, 2, 9]
 
 
 class TestPolygon:
@@ -63,3 +66,16 @@ def test_dumps(self):
             [[[1.0, 4.0], [2.0, 3.7], [7.0, 4.0]], [[5.0, 7.0], [6.0, 7.0], [9.0, 8.0]]]
         )
         assert multipolygon.dumps() == _DATA_MULTIPOLYGON
+
+
+class TestRLE:
+    def test_init(self):
+        assert RLE([272, 2, 4, 4, 2, 9])._data == [272, 2, 4, 4, 2, 9]
+
+    def test_loads(self):
+        rle = RLE.loads(_DATA_RLE)
+        assert rle._data == _DATA_RLE
+
+    def test_dumps(self):
+        rle = RLE(_DATA_RLE)
+        assert rle.dumps() == _DATA_RLE

tensorbay/label/__init__.py

@@ -15,8 +15,10 @@
 from .label_polygon import (
     LabeledMultiPolygon,
     LabeledPolygon,
+    LabeledRLE,
     MultiPolygonSubcatalog,
     PolygonSubcatalog,
+    RLESubcatalog,
 )
 from .label_polyline import (
     LabeledMultiPolyline2D,
@@ -35,23 +37,25 @@
     "CategoryInfo",
     "Classification",
     "ClassificationSubcatalog",
+    "Items",
     "Keypoints2DSubcatalog",
     "KeypointsInfo",
     "Label",
     "LabelType",
     "LabeledBox2D",
     "LabeledBox3D",
     "LabeledKeypoints2D",
+    "LabeledMultiPolygon",
     "LabeledMultiPolyline2D",
     "LabeledPolygon",
     "LabeledPolyline2D",
+    "LabeledRLE",
     "LabeledSentence",
-    "MultiPolyline2DSubcatalog",
-    "LabeledMultiPolygon",
-    "Items",
     "MultiPolygonSubcatalog",
+    "MultiPolyline2DSubcatalog",
     "PolygonSubcatalog",
     "Polyline2DSubcatalog",
+    "RLESubcatalog",
     "SentenceSubcatalog",
     "Subcatalogs",
     "Word",

tensorbay/label/basic.py

@@ -43,6 +43,7 @@ class LabelType(TypeEnum):
     POLYGON = "polygon"
     POLYLINE2D = "polyline2d"
     MULTI_POLYLINE2D = "multi_polyline2d"
+    RLE = "rle"
     KEYPOINTS2D = "keypoints2d"
     MULTI_POLYGON = "multi_polygon"
     SENTENCE = "sentence"

tensorbay/label/catalog.py

@@ -25,6 +25,7 @@
    :class:`.PolygonSubcatalog`           subcatalog for polygon type of label
    :class:`.Polyline2DSubcatalog`        subcatalog for 2D polyline type of label
    :class:`.MultiPolygonSubcatalog`      subcatalog for multiple polygon type of label
+   :class:`.RLESubcatalog`               subcatalog for rle mask type of label
    :class:`.MultiPolyline2DSubcatalog`   subcatalog for 2D multiple polyline type of label
    :class:`.SentenceSubcatalog`          subcatalog for transcripted sentence type of label
    ===================================   ==================================================
@@ -39,7 +40,7 @@
 from .label_box import Box2DSubcatalog, Box3DSubcatalog
 from .label_classification import ClassificationSubcatalog
 from .label_keypoints import Keypoints2DSubcatalog
-from .label_polygon import MultiPolygonSubcatalog, PolygonSubcatalog
+from .label_polygon import MultiPolygonSubcatalog, PolygonSubcatalog, RLESubcatalog
 from .label_polyline import MultiPolyline2DSubcatalog, Polyline2DSubcatalog
 from .label_sentence import SentenceSubcatalog
 
@@ -52,6 +53,7 @@
     MultiPolyline2DSubcatalog,
     Keypoints2DSubcatalog,
     MultiPolygonSubcatalog,
+    RLESubcatalog,
     SentenceSubcatalog,
 ]
 
@@ -104,6 +106,7 @@ class Catalog(ReprMixin, AttrsMixin):
     multi_polyline2d: MultiPolyline2DSubcatalog = _attr()
     keypoints2d: Keypoints2DSubcatalog = _attr()
     multi_polygon: MultiPolygonSubcatalog = _attr()
+    rle: RLESubcatalog = _attr()
     sentence: SentenceSubcatalog = _attr()
 
     def __bool__(self) -> bool:

tensorbay/label/label.py

@@ -13,17 +13,20 @@
 .. table:: label classes
    :widths: auto
 
-   ============================================================= ===================================
-   label classes                                                 explaination
-   ============================================================= ===================================
-   :class:`~tensorbay.label.label_classification.Classification` classification type of label
-   :class:`~tensorbay.label.label_box.LabeledBox2D`              2D bounding box type of label
-   :class:`~tensorbay.label.label_box.LabeledBox3D`              3D bounding box type of label
-   :class:`~tensorbay.label.label_polygon.LabeledPolygon`        polygon type of label
-   :class:`~tensorbay.label.label_polyline.LabeledPolyline2D`    2D polyline type of label
-   :class:`~tensorbay.label.label_keypoints.LabeledKeypoints2D`  2D keypoints type of label
-   :class:`~tensorbay.label.label_sentence.LabeledSentence`      transcripted sentence type of label
-   ============================================================= ===================================
+   ================================ ===================================
+   label classes                    explaination
+   ================================ ===================================
+   :class:`.Classification`         classification type of label
+   :class:`.LabeledBox2D`           2D bounding box type of label
+   :class:`.LabeledBox3D`           3D bounding box type of label
+   :class:`.LabeledPolygon`         polygon type of label
+   :class:`.LabeledMultiPolygon`    polygon lists type of label
+   :class:`.LabeledRLE`             rle mask type of label
+   :class:`.LabeledPolyline2D`      2D polyline type of label
+   :class:`.LabeledMultiPolyline2D` 2D polyline lists type of label
+   :class:`.LabeledKeypoints2D`     2D keypoints type of label
+   :class:`.LabeledSentence`        transcripted sentence type of label
+   ================================ ===================================
 
 """
 
@@ -35,7 +38,7 @@
 from .label_box import LabeledBox2D, LabeledBox3D
 from .label_classification import Classification
 from .label_keypoints import LabeledKeypoints2D
-from .label_polygon import LabeledMultiPolygon, LabeledPolygon
+from .label_polygon import LabeledMultiPolygon, LabeledPolygon, LabeledRLE
 from .label_polyline import LabeledMultiPolyline2D, LabeledPolyline2D
 from .label_sentence import LabeledSentence
 
@@ -74,6 +77,7 @@ class Label(ReprMixin, AttrsMixin):
     polygon: List[LabeledPolygon] = _attr()
     polyline2d: List[LabeledPolyline2D] = _attr()
     multi_polyline2d: List[LabeledMultiPolyline2D] = _attr()
+    rle: List[LabeledRLE] = _attr()
     keypoints2d: List[LabeledKeypoints2D] = _attr()
     multi_polygon: List[LabeledMultiPolygon] = _attr()
     sentence: List[LabeledSentence] = _attr()