feat(label): add MaskCategoriesMixin, MaskCategoryInfo to MaskSubcatalog

tensorbay/label/__init__.py

@@ -34,7 +34,7 @@
     Polyline2DSubcatalog,
 )
 from .label_sentence import LabeledSentence, SentenceSubcatalog, Word
-from .supports import CategoryInfo, KeypointsInfo
+from .supports import CategoryInfo, KeypointsInfo, MaskCategoryInfo
 
 __all__ = [
     "AttributeInfo",
@@ -59,6 +59,7 @@
     "LabeledPolyline2D",
     "LabeledRLE",
     "LabeledSentence",
+    "MaskCategoryInfo",
     "MultiPolygonSubcatalog",
     "MultiPolyline2DSubcatalog",
     "PanopticMask",

tensorbay/label/label_mask.py

@@ -9,10 +9,10 @@
 
 from ..utility import FileMixin, RemoteFileMixin, ReprMixin
 from .basic import AttributeType, SubcatalogBase
-from .supports import AttributesMixin, CategoriesMixin, IsTrackingMixin
+from .supports import AttributesMixin, IsTrackingMixin, MaskCategoriesMixin
 
 
-class SemanticMaskSubcatalog(SubcatalogBase, CategoriesMixin, AttributesMixin):
+class SemanticMaskSubcatalog(SubcatalogBase, MaskCategoriesMixin, AttributesMixin):
     """This class defines the subcatalog for semantic mask type of labels.
 
     Attributes:
@@ -33,7 +33,10 @@ class SemanticMaskSubcatalog(SubcatalogBase, CategoriesMixin, AttributesMixin):
 
         >>> catalog = {
         ...     "SEMANTIC_MASK": {
-        ...         "categories": [{'name': 'cat'}, {'name': 'dog'}],
+        ...         "categories": [
+        ...             {'name': 'cat', "categoryId": 1},
+        ...             {'name': 'dog', "categoryId": 2}
+        ...         ],
         ...         "attributes": [{'name': 'occluded', 'type': 'boolean'}],
         ...     }
         ... }
@@ -46,8 +49,8 @@ class SemanticMaskSubcatalog(SubcatalogBase, CategoriesMixin, AttributesMixin):
         *Initialization Method 2:* Init an empty SemanticMaskSubcatalog and then add the attributes.
 
         >>> semantic_mask_subcatalog = SemanticMaskSubcatalog()
-        >>> semantic_mask_subcatalog.add_category("cat")
-        >>> semantic_mask_subcatalog.add_category("dog")
+        >>> semantic_mask_subcatalog.add_category("cat", 1)
+        >>> semantic_mask_subcatalog.add_category("dog", 2)
         >>> semantic_mask_subcatalog.add_attribute("occluded", type_="boolean")
         >>> semantic_mask_subcatalog
         SemanticMaskSubcatalog(
@@ -100,7 +103,7 @@ class InstanceMaskSubcatalog(SubcatalogBase, IsTrackingMixin, AttributesMixin):
     """
 
 
-class PanopticMaskSubcatalog(SubcatalogBase, CategoriesMixin, AttributesMixin):
+class PanopticMaskSubcatalog(SubcatalogBase, MaskCategoriesMixin, AttributesMixin):
     """This class defines the subcatalog for panoptic mask type of labels.
 
     Attributes:
@@ -121,7 +124,10 @@ class PanopticMaskSubcatalog(SubcatalogBase, CategoriesMixin, AttributesMixin):
 
         >>> catalog = {
         ...     "PANOPTIC_MASK": {
-        ...         "categories": [{'name': 'cat'}, {'name': 'dog'}],
+        ...         "categories": [
+        ...             {'name': 'cat', "categoryId": 1},
+        ...             {'name': 'dog', "categoryId": 2}
+        ...         ],
         ...         "attributes": [{'name': 'occluded', 'type': 'boolean'}],
         ...     }
         ... }
@@ -134,8 +140,8 @@ class PanopticMaskSubcatalog(SubcatalogBase, CategoriesMixin, AttributesMixin):
         *Initialization Method 2:* Init an empty PanopticMaskSubcatalog and then add the attributes.
 
         >>> panoptic_mask_subcatalog = PanopticMaskSubcatalog()
-        >>> panoptic_mask_subcatalog.add_category("cat")
-        >>> panoptic_mask_subcatalog.add_category("dog")
+        >>> panoptic_mask_subcatalog.add_category("cat", 1)
+        >>> panoptic_mask_subcatalog.add_category("dog", 2)
         >>> panoptic_mask_subcatalog.add_attribute("occluded", type_="boolean")
         >>> panoptic_mask_subcatalog
         PanopticMaskSubcatalog(

tensorbay/label/supports.py

@@ -3,10 +3,12 @@
 # Copyright 2021 Graviti. Licensed under MIT License.
 #
 
-"""CatagoryInfo, KeypointsInfo and different SubcatalogMixin classes.
+"""CatagoryInfo, MaskCategoryInfo, KeypointsInfo and different SubcatalogMixin classes.
 
 :class:`CatagoryInfo` defines a category with the name and description of it.
 
+:class:`MaskCategoryInfo` defines a category with the name, categoryId and description of it.
+
 :class:`KeypointsInfo` defines the structure of a set of keypoints.
 
 .. table:: mixin classes for subcatalog
@@ -90,6 +92,35 @@ def dumps(self) -> Dict[str, str]:
         return self._dumps()
 
 
+class MaskCategoryInfo(CategoryInfo):
+    """This class represents the information of a category, including category name and description.
+
+    Arguments:
+        name: The name of the category.
+        categoryId: The id of the category.
+        description: The description of the category.
+
+    Attributes:
+        name: The name of the category.
+        categoryId: The id of the category.
+        description: The description of the category.
+
+    Examples:
+        >>> MaskCategoryInfo(name="example", category_id=1, description="This is an example")
+        MaskCategoryInfo("example")(
+          (categoryId): 1
+        )
+
+    """
+
+    _repr_attrs = ("category_id",)
+    category_id: int = attr(key=camel)
+
+    def __init__(self, name: str, category_id: int, description: str = "") -> None:
+        super().__init__(name, description)
+        self.category_id = category_id
+
+
 class _VisibleType(Enum):
     """All the possible visible types of keypoints labels."""
 
@@ -341,6 +372,60 @@ def add_category(self, name: str, description: str = "") -> None:
         self.categories.append(CategoryInfo(name, description))
 
 
+class MaskCategoriesMixin(AttrsMixin):
+    """A mixin class supporting category information of a subcatalog.
+
+    Attributes:
+        categories: All the possible categories in the corresponding dataset
+            stored in a :class:`~tensorbay.utility.name.NameList`
+            with the category names as keys
+            and the :class:`~tensorbay.label.supports.MaskCategoryInfo` as values.
+        category_delimiter: The delimiter in category values indicating parent-child relationship.
+
+    """
+
+    category_delimiter: str = attr(is_dynamic=True, key=camel)
+    categories: NameList[MaskCategoryInfo] = attr(is_dynamic=True)
+
+    def get_category_to_index(self) -> Dict[str, int]:
+        """Return the dict containing the conversion from category name to category id.
+
+        Returns:
+            A dict containing the conversion from category name to category id.
+
+        """
+        if not hasattr(self, "categories"):
+            return {}
+
+        return {item.name: item.category_id for item in self.categories}
+
+    def get_index_to_category(self) -> Dict[int, str]:
+        """Return the dict containing the conversion from category id to category name.
+
+        Returns:
+            A dict containing the conversion from category id to category name.
+
+        """
+        if not hasattr(self, "categories"):
+            return {}
+
+        return {item.category_id: item.name for item in self.categories}
+
+    def add_category(self, name: str, category_id: int, description: str = "") -> None:
+        """Add a category to the Subcatalog.
+
+        Arguments:
+            name: The name of the category.
+            category_id: The id of the category.
+            description: The description of the category.
+
+        """
+        if not hasattr(self, "categories"):
+            self.categories = NameList()
+
+        self.categories.append(MaskCategoryInfo(name, category_id, description))
+
+
 class AttributesMixin(AttrsMixin):
     """A mixin class supporting attribute information of a subcatalog.