class DatasetItem: id = attrib(converter=lambda x: str(x).replace('\\', '/'), type=str, validator=not_empty) annotations = attrib(factory=list, validator=default_if_none(list)) subset = attrib(converter=lambda v: v or DEFAULT_SUBSET_NAME, default=None) path = attrib(factory=list, validator=default_if_none(list)) image = attrib(type=Image, default=None) @image.validator def _image_validator(self, attribute, image): if callable(image) or isinstance(image, np.ndarray): image = Image(data=image) elif isinstance(image, str): image = Image(path=image) assert image is None or isinstance(image, Image) self.image = image attributes = attrib(factory=dict, validator=default_if_none(dict)) @property def has_image(self): return self.image is not None def wrap(item, **kwargs): return attr.evolve(item, **kwargs)
class Category: # Names for specific points, e.g. eye, hose, mouth etc. # These labels are not required to be in LabelCategories labels: List[str] = field(factory=list, validator=default_if_none(list)) # Pairs of connected point indices joints: Set[Tuple[int, int]] = field(factory=set, validator=default_if_none(set))
class DatasetItem: id: str = field(converter=lambda x: str(x).replace('\\', '/'), validator=not_empty) annotations: List[Annotation] = field(factory=list, validator=default_if_none(list)) subset: str = field(converter=lambda v: v or DEFAULT_SUBSET_NAME, default=None) # TODO: introduce "media" field with type info. Replace image and pcd. image: Optional[Image] = field(default=None) # TODO: introduce pcd type like Image point_cloud: Optional[str] = field( converter=lambda x: str(x).replace('\\', '/') if x else None, default=None) related_images: List[Image] = field(default=None) def __attrs_post_init__(self): if (self.has_image and self.has_point_cloud): raise ValueError("Can't set both image and point cloud info") if self.related_images and not self.has_point_cloud: raise ValueError("Related images require point cloud") def _image_converter(image): if callable(image) or isinstance(image, np.ndarray): image = Image(data=image) elif isinstance(image, str): image = Image(path=image) assert image is None or isinstance(image, Image), type(image) return image image.converter = _image_converter def _related_image_converter(images): return list(map(__class__._image_converter, images or [])) related_images.converter = _related_image_converter @point_cloud.validator def _point_cloud_validator(self, attribute, pcd): assert pcd is None or isinstance(pcd, str), type(pcd) attributes: Dict[str, Any] = field(factory=dict, validator=default_if_none(dict)) @property def has_image(self): return self.image is not None @property def has_point_cloud(self): return self.point_cloud is not None def wrap(item, **kwargs): return attr.evolve(item, **kwargs)
class Annotation: id = attrib(default=0, validator=default_if_none(int)) attributes = attrib(factory=dict, validator=default_if_none(dict)) group = attrib(default=0, validator=default_if_none(int)) def __attrs_post_init__(self): assert isinstance(self.type, AnnotationType) @property def type(self) -> AnnotationType: return self._type # must be set in subclasses def wrap(self, **kwargs): return attr.evolve(self, **kwargs)
class _Shape(Annotation): # Flattened list of point coordinates points: List[float] = field(converter=lambda x: \ np.around(x, COORDINATE_ROUNDING_DIGITS).tolist()) label: Optional[int] = field(converter=attr.converters.optional(int), default=None, kw_only=True) z_order: int = field(default=0, validator=default_if_none(int), kw_only=True) def get_area(self): raise NotImplementedError() def get_bbox(self) -> Tuple[float, float, float, float]: "Returns [x, y, w, h]" points = self.points if not points: return None xs = [p for p in points[0::2]] ys = [p for p in points[1::2]] x0 = min(xs) x1 = max(xs) y0 = min(ys) y1 = max(ys) return [x0, y0, x1 - x0, y1 - y0]
class _Shape(Annotation): points = attrib( converter=lambda x: [round(p, _COORDINATE_ROUNDING_DIGITS) for p in x]) label = attrib(converter=attr.converters.optional(int), default=None, kw_only=True) z_order = attrib(default=0, validator=default_if_none(int), kw_only=True) def get_area(self): raise NotImplementedError() def get_bbox(self): points = self.points if not points: return None print("----p----", p) print(" print(" - ---points - --- ",points)", points) xs = [p for p in points[0::2]] ys = [p for p in points[1::2]] x1 = [p for p in points[3]] y1 = [p for p in points[4]] z1 = [p for p in points[5]] print("___x---->y--->z", x, "--->", y, "--->", z) x0 = min(xs) x1 = max(xs) y0 = min(ys) y1 = max(ys) return [x0, y0, x1 - x0, y1 - y0]
class LabelCategories(Categories): @attrs(repr_ns='LabelCategories') class Category: name = attrib(converter=str, validator=not_empty) parent = attrib(default='', validator=default_if_none(str)) attributes = attrib(factory=set, validator=default_if_none(set)) items = attrib(factory=list, validator=default_if_none(list)) _indices = attrib(factory=dict, init=False, eq=False) @classmethod def from_iterable(cls, iterable): """Generation of LabelCategories from iterable object Args: iterable ([type]): This iterable object can be: 1)simple str - will generate one Category with str as name 2)list of str - will interpreted as list of Category names 3)list of positional argumetns - will generate Categories with this arguments Returns: LabelCategories: LabelCategories object """ temp_categories = cls() if isinstance(iterable, str): iterable = [[iterable]] for category in iterable: if isinstance(category, str): category = [category] temp_categories.add(*category) return temp_categories def __attrs_post_init__(self): self._reindex() def _reindex(self): indices = {} for index, item in enumerate(self.items): assert item.name not in self._indices indices[item.name] = index self._indices = indices def add(self, name: str, parent: str = None, attributes: dict = None): assert name not in self._indices, name index = len(self.items) self.items.append(self.Category(name, parent, attributes)) self._indices[name] = index return index def find(self, name: str): index = self._indices.get(name) if index is not None: return index, self.items[index] return index, None
class PointsCategories(Categories): @attrs(repr_ns="PointsCategories") class Category: labels = attrib(factory=list, validator=default_if_none(list)) joints = attrib(factory=set, validator=default_if_none(set)) items = attrib(factory=dict, validator=default_if_none(dict)) @classmethod def from_iterable(cls, iterable): """Generation of PointsCategories from iterable object Args: iterable ([type]): This iterable object can be: 1) list of positional argumetns - will generate Categories with these arguments Returns: PointsCategories: PointsCategories object """ temp_categories = cls() for category in iterable: temp_categories.add(*category) return temp_categories def add(self, label_id, labels=None, joints=None): if joints is None: joints = [] joints = set(map(tuple, joints)) self.items[label_id] = self.Category(labels, joints)
class MaskCategories(Categories): @classmethod def make_default(cls, size=256): from datumaro.util.mask_tools import generate_colormap return cls(generate_colormap(size)) colormap = attrib(factory=dict, validator=default_if_none(dict)) _inverse_colormap = attrib(default=None, validator=attr.validators.optional(dict)) @property def inverse_colormap(self): from datumaro.util.mask_tools import invert_colormap if self._inverse_colormap is None: if self.colormap is not None: self._inverse_colormap = invert_colormap(self.colormap) return self._inverse_colormap def __eq__(self, other): if not super().__eq__(other): return False if not isinstance(other, __class__): return False for label_id, my_color in self.colormap.items(): other_color = other.colormap.get(label_id) if not np.array_equal(my_color, other_color): return False return True
class PointsCategories(Categories): """ Describes (key-)point metainfo such as point names and joints. """ @attrs(slots=True, order=False) class Category: # Names for specific points, e.g. eye, hose, mouth etc. # These labels are not required to be in LabelCategories labels: List[str] = field(factory=list, validator=default_if_none(list)) # Pairs of connected point indices joints: Set[Tuple[int, int]] = field(factory=set, validator=default_if_none(set)) items: Dict[int, Category] = field(factory=dict, validator=default_if_none(dict)) @classmethod def from_iterable( cls, iterable: Union[Tuple[int, List[str]], Tuple[int, List[str], Set[Tuple[int, int]]], ] ) -> PointsCategories: """ Create PointsCategories from an iterable. Args: - iterable - An Iterable with the following elements: - a label id - a list of positional arguments for Categories Returns: PointsCategories: PointsCategories object """ temp_categories = cls() for args in iterable: temp_categories.add(*args) return temp_categories def add(self, label_id: int, labels: Optional[Iterable[str]] = None, joints: Iterable[Tuple[int, int]] = None): if joints is None: joints = [] joints = set(map(tuple, joints)) self.items[label_id] = self.Category(labels, joints) def __contains__(self, idx: int) -> bool: return idx in self.items def __getitem__(self, idx: int) -> Category: return self.items[idx] def __len__(self) -> int: return len(self.items)
class Categories: """ A base class for annotation metainfo. It is supposed to include dataset-wide metainfo like available labels, label colors, label attributes etc. """ # Describes the list of possible annotation-type specific attributes # in a dataset. attributes: Set[str] = field(factory=set, validator=default_if_none(set), eq=False)
class Annotation: """ A base annotation class. Derived classes must define the '_type' class variable with a value from the AnnotationType enum. """ # Describes an identifier of the annotation # Is not required to be unique within DatasetItem annotations or dataset id: int = field(default=0, validator=default_if_none(int)) # Arbitrary annotation-specific attributes. Typically, includes # metainfo and properties that are not covered by other fields. # If possible, try to limit value types of values by the simple # builtin types (int, float, bool, str) to increase compatibility with # different formats. # There are some established names for common attributes like: # - "occluded" (bool) # - "visible" (bool) # Possible dataset attributes can be described in Categories.attributes. attributes: Dict[str, Any] = field(factory=dict, validator=default_if_none(dict)) # Annotations can be grouped, which means they describe parts of a # single object. The value of 0 means there is no group. group: int = field(default=NO_GROUP, validator=default_if_none(int)) @property def type(self) -> AnnotationType: return self._type # must be set in subclasses def as_dict(self) -> Dict[str, Any]: "Returns a dictionary { field_name: value }" return asdict(self) def wrap(self, **kwargs): "Returns a modified copy of the object" return attr.evolve(self, **kwargs)
class MaskCategories(Categories): """ Describes a color map for segmentation masks. """ @classmethod def generate(cls, size: int = 255, include_background: bool = True) \ -> MaskCategories: """ Generates MaskCategories with the specified size. If include_background is True, the result will include the item "0: (0, 0, 0)", which is typically used as a background color. """ from datumaro.util.mask_tools import generate_colormap return cls( generate_colormap(size, include_background=include_background)) colormap: Colormap = field(factory=dict, validator=default_if_none(dict)) _inverse_colormap: Optional[Dict[RgbColor, int]] = field( default=None, validator=attr.validators.optional(dict)) @property def inverse_colormap(self) -> Dict[RgbColor, int]: from datumaro.util.mask_tools import invert_colormap if self._inverse_colormap is None: if self.colormap is not None: self._inverse_colormap = invert_colormap(self.colormap) return self._inverse_colormap def __contains__(self, idx: int) -> bool: return idx in self.colormap def __getitem__(self, idx: int) -> RgbColor: return self.colormap[idx] def __len__(self) -> int: return len(self.colormap) def __eq__(self, other): if not super().__eq__(other): return False if not isinstance(other, __class__): return False for label_id, my_color in self.colormap.items(): other_color = other.colormap.get(label_id) if not np.array_equal(my_color, other_color): return False return True
class Mask(Annotation): _type = AnnotationType.mask _image = attrib() label = attrib(converter=attr.converters.optional(int), default=None, kw_only=True) z_order = attrib(default=0, validator=default_if_none(int), kw_only=True) def __attrs_post_init__(self): if isinstance(self._image, np.ndarray): self._image = self._image.astype(bool) @property def image(self): if callable(self._image): return self._image() return self._image def as_class_mask(self, label_id=None): if label_id is None: label_id = self.label from datumaro.util.mask_tools import make_index_mask return make_index_mask(self.image, label_id) def as_instance_mask(self, instance_id): from datumaro.util.mask_tools import make_index_mask return make_index_mask(self.image, instance_id) def get_area(self): return np.count_nonzero(self.image) def get_bbox(self): from datumaro.util.mask_tools import find_mask_bbox return find_mask_bbox(self.image) def paint(self, colormap): from datumaro.util.mask_tools import paint_mask return paint_mask(self.as_class_mask(), colormap) def __eq__(self, other): if not super().__eq__(other): return False if not isinstance(other, __class__): return False return \ (self.label == other.label) and \ (self.z_order == other.z_order) and \ (np.array_equal(self.image, other.image))
class PointsCategories(Categories): Category = namedtuple('Category', ['labels', 'joints']) items = attrib(factory=dict, validator=default_if_none(dict)) @classmethod def from_iterable(cls, iterable): """Generation of PointsCategories from iterable object Args: iterable ([type]): This iterable object can be: 1)simple int - will generate one Category with int as label 2)list of int - will interpreted as list of Category labels 3)list of positional argumetns - will generate Categories with this arguments Returns: PointsCategories: PointsCategories object """ temp_categories = cls() if isinstance(iterable, int): iterable = [[iterable]] for category in iterable: if isinstance(category, int): category = [category] temp_categories.add(*category) return temp_categories def add(self, label_id, labels=None, joints=None): if labels is None: labels = [] if joints is None: joints = [] joints = set(map(tuple, joints)) self.items[label_id] = self.Category(labels, joints)
class Categories: attributes = attrib(factory=set, validator=default_if_none(set), eq=False)
class Category: labels = attrib(factory=list, validator=default_if_none(list)) joints = attrib(factory=set, validator=default_if_none(set))
class LabelCategories(Categories): Category = namedtuple('Category', ['name', 'parent', 'attributes']) items = attrib(factory=list, validator=default_if_none(list)) _indices = attrib(factory=dict, init=False, eq=False) @classmethod def from_iterable(cls, iterable): """Generation of LabelCategories from iterable object Args: iterable ([type]): This iterable object can be: 1)simple str - will generate one Category with str as name 2)list of str - will interpreted as list of Category names 3)list of positional argumetns - will generate Categories with this arguments Returns: LabelCategories: LabelCategories object """ temp_categories = cls() if isinstance(iterable, str): iterable = [[iterable]] for category in iterable: if isinstance(category, str): category = [category] temp_categories.add(*category) return temp_categories def __attrs_post_init__(self): self._reindex() def _reindex(self): indices = {} for index, item in enumerate(self.items): assert item.name not in self._indices indices[item.name] = index self._indices = indices def add(self, name, parent=None, attributes=None): assert name not in self._indices, name if attributes is None: attributes = set() else: if not isinstance(attributes, set): attributes = set(attributes) for attr in attributes: assert isinstance(attr, str) if parent is None: parent = '' index = len(self.items) self.items.append(self.Category(name, parent, attributes)) self._indices[name] = index return index def find(self, name): index = self._indices.get(name) if index is not None: return index, self.items[index] return index, None
class Categories: attributes = attrib(factory=set, validator=default_if_none(set), kw_only=True)
class Category: name: str = field(converter=str, validator=not_empty) parent: str = field(default='', validator=default_if_none(str)) attributes: Set[str] = field(factory=set, validator=default_if_none(set))
class Mask(Annotation): """ Represents a 2d single-instance binary segmentation mask. """ _type = AnnotationType.mask _image = field() label: Optional[int] = field(converter=attr.converters.optional(int), default=None, kw_only=True) z_order: int = field(default=0, validator=default_if_none(int), kw_only=True) def __attrs_post_init__(self): if isinstance(self._image, np.ndarray): self._image = self._image.astype(bool) @property def image(self) -> BinaryMaskImage: image = self._image if callable(image): image = image() return image def as_class_mask(self, label_id: Optional[int] = None) -> IndexMaskImage: """ Produces a class index mask. Mask label id can be changed. """ if label_id is None: label_id = self.label from datumaro.util.mask_tools import make_index_mask return make_index_mask(self.image, label_id) def as_instance_mask(self, instance_id: int) -> IndexMaskImage: """ Produces a instance index mask. """ from datumaro.util.mask_tools import make_index_mask return make_index_mask(self.image, instance_id) def get_area(self) -> int: return np.count_nonzero(self.image) def get_bbox(self) -> Tuple[int, int, int, int]: """ Computes the bounding box of the mask. Returns: [x, y, w, h] """ from datumaro.util.mask_tools import find_mask_bbox return find_mask_bbox(self.image) def paint(self, colormap: Colormap) -> np.ndarray: """ Applies a colormap to the mask and produces the resulting image. """ from datumaro.util.mask_tools import paint_mask return paint_mask(self.as_class_mask(), colormap) def __eq__(self, other): if not super().__eq__(other): return False if not isinstance(other, __class__): return False return \ (self.label == other.label) and \ (self.z_order == other.z_order) and \ (np.array_equal(self.image, other.image))
class Category: name = attrib(converter=str, validator=not_empty) parent = attrib(default='', validator=default_if_none(str)) attributes = attrib(factory=set, validator=default_if_none(set))
class LabelCategories(Categories): @attrs(slots=True, order=False) class Category: name: str = field(converter=str, validator=not_empty) parent: str = field(default='', validator=default_if_none(str)) attributes: Set[str] = field(factory=set, validator=default_if_none(set)) items: List[str] = field(factory=list, validator=default_if_none(list)) _indices: Dict[str, int] = field(factory=dict, init=False, eq=False) @classmethod def from_iterable( cls, iterable: Iterable[Union[str, Tuple[str], Tuple[str, str], Tuple[str, str, List[str]], ]] ) -> LabelCategories: """ Creates a LabelCategories from iterable. Args: iterable: This iterable object can be: - a list of str - will be interpreted as list of Category names - a list of positional arguments - will generate Categories with these arguments Returns: a LabelCategories object """ temp_categories = cls() for category in iterable: if isinstance(category, str): category = [category] temp_categories.add(*category) return temp_categories def __attrs_post_init__(self): self._reindex() def _reindex(self): indices = {} for index, item in enumerate(self.items): assert item.name not in self._indices indices[item.name] = index self._indices = indices def add(self, name: str, parent: Optional[str] = None, attributes: Optional[Set[str]] = None) -> int: assert name assert name not in self._indices, name index = len(self.items) self.items.append(self.Category(name, parent, attributes)) self._indices[name] = index return index def find(self, name: str) -> Tuple[Optional[int], Optional[Category]]: index = self._indices.get(name) if index is not None: return index, self.items[index] return index, None def __getitem__(self, idx: int) -> Category: return self.items[idx] def __contains__(self, value: Union[int, str]) -> bool: if isinstance(value, str): return self.find(value)[1] is not None else: return 0 <= value and value < len(self.items) def __len__(self) -> int: return len(self.items) def __iter__(self) -> Iterator[Category]: return iter(self.items)