#!/usr/bin/env python3
#
# Copyright 2021 Graviti. Licensed under MIT License.
#
"""The implementation of the TensorBay 2D keypoints label."""
from typing import Any, Dict, Iterable, List, Optional, Type, TypeVar, Union
from tensorbay.geometry import Keypoints2D
from tensorbay.label.basic import SubcatalogBase, _LabelBase
from tensorbay.label.supports import (
AttributesMixin,
CategoriesMixin,
IsTrackingMixin,
KeypointsInfo,
)
from tensorbay.utility import ReprType, attr, attr_base, common_loads
[docs]class Keypoints2DSubcatalog(SubcatalogBase, IsTrackingMixin, CategoriesMixin, AttributesMixin):
"""This class defines the subcatalog for 2D keypoints type of labels.
Arguments:
is_tracking: A boolean value indicates whether the corresponding
subcatalog contains tracking information.
Attributes:
description: The description of the entire 2D keypoints subcatalog.
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.CategoryInfo` as values.
category_delimiter: The delimiter in category values indicating parent-child relationship.
attributes: All the possible attributes in the corresponding dataset
stored in a :class:`~tensorbay.utility.name.NameList`
with the attribute names as keys
and the :class:`~tensorbay.label.attribute.AttributeInfo` as values.
is_tracking: Whether the Subcatalog contains tracking information.
Examples:
*Initialization Method 1:* Init from ``Keypoints2DSubcatalog.loads()`` method.
>>> catalog = {
... "KEYPOINTS2D": {
... "isTracking": True,
... "categories": [{"name": "0"}, {"name": "1"}],
... "attributes": [{"name": "gender", "enum": ["male", "female"]}],
... "keypoints": [
... {
... "number": 2,
... "names": ["L_shoulder", "R_Shoulder"],
... "skeleton": [(0, 1)],
... }
... ],
... }
... }
>>> Keypoints2DSubcatalog.loads(catalog["KEYPOINTS2D"])
Keypoints2DSubcatalog(
(is_tracking): True,
(keypoints): [...],
(categories): NameList [...],
(attributes): NameList [...]
)
*Initialization Method 2:* Init an empty Keypoints2DSubcatalog and then add the attributes.
>>> from tensorbay.label import CategoryInfo, AttributeInfo, KeypointsInfo
>>> from tensorbay.utility import NameList
>>> categories = NameList()
>>> categories.append(CategoryInfo("a"))
>>> attributes = NameList()
>>> attributes.append(AttributeInfo("gender", enum=["female", "male"]))
>>> keypoints2d_subcatalog = Keypoints2DSubcatalog()
>>> keypoints2d_subcatalog.is_tracking = True
>>> keypoints2d_subcatalog.categories = categories
>>> keypoints2d_subcatalog.attributes = attributes
>>> keypoints2d_subcatalog.add_keypoints(
... 2,
... names=["L_shoulder", "R_Shoulder"],
... skeleton=[(0,1)],
... visible="BINARY",
... parent_categories="shoulder",
... description="12345",
... )
>>> keypoints2d_subcatalog
Keypoints2DSubcatalog(
(is_tracking): True,
(keypoints): [...],
(categories): NameList [...],
(attributes): NameList [...]
)
"""
_keypoints: List[KeypointsInfo] = attr(key="keypoints")
def __init__(self, is_tracking: bool = False) -> None:
SubcatalogBase.__init__(self)
IsTrackingMixin.__init__(self, is_tracking)
self._keypoints: List[KeypointsInfo] = []
@property
def keypoints(self) -> List[KeypointsInfo]:
"""Return the KeypointsInfo of the Subcatalog.
Returns:
A list of :class:`~tensorbay.label.supports.KeypointsInfo`.
Examples:
>>> keypoints2d_subcatalog = Keypoints2DSubcatalog()
>>> keypoints2d_subcatalog.add_keypoints(2)
>>> keypoints2d_subcatalog.keypoints
[KeypointsInfo(
(number): 2
)]
"""
return self._keypoints
[docs] def add_keypoints(
self,
number: int,
*,
names: Optional[Iterable[str]] = None,
skeleton: Optional[Iterable[Iterable[int]]] = None,
visible: Optional[str] = None,
parent_categories: Union[None, str, Iterable[str]] = None,
description: str = "",
) -> None:
"""Add a type of keypoints to the subcatalog.
Arguments:
number: The number of keypoints.
names: All the names of keypoints.
skeleton: The skeleton of the keypoints
indicating which keypoint should connect with another.
visible: The visible type of the keypoints, can only be 'BINARY' or 'TERNARY'.
It determines the range of the
:attr:`Keypoint2D.v<tensorbay.geometry.keypoint.Keypoint2D.v>`.
parent_categories: The parent categories of the keypoints.
description: The description of keypoints.
Examples:
>>> keypoints2d_subcatalog = Keypoints2DSubcatalog()
>>> keypoints2d_subcatalog.add_keypoints(
... 2,
... names=["L_shoulder", "R_Shoulder"],
... skeleton=[(0,1)],
... visible="BINARY",
... parent_categories="shoulder",
... description="12345",
... )
>>> keypoints2d_subcatalog.keypoints
[KeypointsInfo(
(number): 2,
(names): [...],
(skeleton): [...],
(visible): 'BINARY',
(parent_categories): [...]
)]
"""
self._keypoints.append(
KeypointsInfo(
number=number,
names=names,
skeleton=skeleton,
visible=visible,
parent_categories=parent_categories,
description=description,
)
)
[docs] def dumps(self) -> Dict[str, Any]:
"""Dumps all the information of the keypoints into a dict.
Returns:
A dict containing all the information of this Keypoints2DSubcatalog.
Examples:
>>> # keypoints2d_subcatalog is the instance initialized above.
>>> keypoints2d_subcatalog.dumps()
{
'isTracking': True,
'categories': [{'name': 'a'}],
'attributes': [{'name': 'gender', 'enum': ['female', 'male']}],
'keypoints': [
{
'number': 2,
'names': ['L_shoulder', 'R_Shoulder'],
'skeleton': [(0, 1)],
}
]
}
"""
return self._dumps()
[docs]class LabeledKeypoints2D(_LabelBase, Keypoints2D):
"""This class defines the concept of 2D keypoints label.
:class:`LabeledKeypoints2D` is the 2D keypoints type of label,
which is often used for CV tasks such as human body pose estimation.
Arguments:
keypoints: A list of 2D keypoint.
category: The category of the label.
attributes: The attributes of the label.
instance: The instance id of the label.
Attributes:
category: The category of the label.
attributes: The attributes of the label.
instance: The instance id of the label.
Examples:
>>> LabeledKeypoints2D(
... [(1, 2), (2, 3)],
... category="example",
... attributes={"key": "value"},
... instance="123",
... )
LabeledKeypoints2D [
Keypoint2D(1, 2),
Keypoint2D(2, 3)
](
(category): 'example',
(attributes): {...},
(instance): '123'
)
"""
_T = TypeVar("_T", bound="LabeledKeypoints2D")
_repr_type = ReprType.SEQUENCE
_repr_attrs = _LabelBase._repr_attrs
_attrs_base: Keypoints2D = attr_base(key="keypoints2d")
def __init__(
self,
keypoints: Optional[Iterable[Iterable[float]]] = None,
*,
category: Optional[str] = None,
attributes: Optional[Dict[str, Any]] = None,
instance: Optional[str] = None,
) -> None:
Keypoints2D.__init__(self, keypoints) # type: ignore[arg-type]
_LabelBase.__init__(self, category, attributes, instance)
[docs] @classmethod
def loads(cls: Type[_T], contents: Dict[str, Any]) -> _T: # type: ignore[override]
"""Loads a LabeledKeypoints2D from a dict containing the information of the label.
Arguments:
contents: A dict containing the information of the 2D keypoints label.
Returns:
The loaded :class:`LabeledKeypoints2D` object.
Examples:
>>> contents = {
... "keypoints2d": [
... {"x": 1, "y": 1, "v": 2},
... {"x": 2, "y": 2, "v": 2},
... ],
... "category": "example",
... "attributes": {"key": "value"},
... "instance": "12345",
... }
>>> LabeledKeypoints2D.loads(contents)
LabeledKeypoints2D [
Keypoint2D(1, 1, 2),
Keypoint2D(2, 2, 2)
](
(category): 'example',
(attributes): {...},
(instance): '12345'
)
"""
return common_loads(cls, contents)
[docs] def dumps(self) -> Dict[str, Any]: # type: ignore[override]
"""Dumps the current 2D keypoints label into a dict.
Returns:
A dict containing all the information of the 2D keypoints label.
Examples:
>>> labeledkeypoints2d = LabeledKeypoints2D(
... [(1, 1, 2), (2, 2, 2)],
... category="example",
... attributes={"key": "value"},
... instance="123",
... )
>>> labeledkeypoints2d.dumps()
{
'category': 'example',
'attributes': {'key': 'value'},
'instance': '123',
'keypoints2d': [{'x': 1, 'y': 1, 'v': 2}, {'x': 2, 'y': 2, 'v': 2}],
}
"""
return self._dumps()