diff --git a/interactive_ai/libs/iai_core_py/iai_core/entities/label.py b/interactive_ai/libs/iai_core_py/iai_core/entities/label.py index d36afa548f..c73ee81835 100644 --- a/interactive_ai/libs/iai_core_py/iai_core/entities/label.py +++ b/interactive_ai/libs/iai_core_py/iai_core/entities/label.py @@ -99,6 +99,8 @@ class Label(PersistentEntity): will be assigned upon saving. If the argument is None, it will be set to ID() :param is_anomalous: boolean that indicates whether the label is the Anomalous label. Always set to False for non-anomaly projects. + :param is_background: boolean that indicates whether the label is the background label. + This label is exclusive to semantic segmentation projects. False for other projects. """ def __init__( # noqa: PLR0913 @@ -111,6 +113,7 @@ def __init__( # noqa: PLR0913 creation_date: datetime.datetime | None = None, is_empty: bool = False, is_anomalous: bool = False, + is_background: bool = False, ephemeral: bool = True, ): color = Color.random() if color is None else color @@ -123,6 +126,7 @@ def __init__( # noqa: PLR0913 self._is_empty = is_empty self._creation_date = creation_date self.is_anomalous = is_anomalous + self.is_background = is_background PersistentEntity.__init__(self, id_=id_, ephemeral=ephemeral) @property @@ -138,8 +142,8 @@ def creation_date(self) -> datetime.datetime: def __repr__(self): """String representation of the label.""" return ( - f"Label({self.id_}, name={self.name}, hotkey={self.hotkey}, " - f"domain={self.domain}, color={self.color.hex_str}, is_anomalous={self.is_anomalous})" + f"Label({self.id_}, name={self.name}, hotkey={self.hotkey}, domain={self.domain}, " + f"color={self.color.hex_str}, is_anomalous={self.is_anomalous}, is_background={self.is_background})" ) def __eq__(self, other: object) -> bool: diff --git a/interactive_ai/libs/iai_core_py/iai_core/repos/mappers/mongodb_mappers/label_mapper.py b/interactive_ai/libs/iai_core_py/iai_core/repos/mappers/mongodb_mappers/label_mapper.py index 3bcca7992f..2c6435d139 100644 --- a/interactive_ai/libs/iai_core_py/iai_core/repos/mappers/mongodb_mappers/label_mapper.py +++ b/interactive_ai/libs/iai_core_py/iai_core/repos/mappers/mongodb_mappers/label_mapper.py @@ -52,6 +52,7 @@ def forward(instance: Label) -> dict: "creation_date": DatetimeToMongo.forward(instance.creation_date), "is_empty": instance.is_empty, "is_anomalous": instance.is_anomalous, + "is_background": instance.is_background, } @staticmethod @@ -70,6 +71,7 @@ def backward(instance: dict) -> Label: is_empty=instance.get("is_empty", False), id_=label_id, is_anomalous=instance.get("is_anomalous", False), + is_background=instance.get("is_background", False), ephemeral=False, ) diff --git a/interactive_ai/libs/iai_core_py/tests/repos/mappers/test_static_mappers.py b/interactive_ai/libs/iai_core_py/tests/repos/mappers/test_static_mappers.py index 61a6f22805..267f0eee20 100644 --- a/interactive_ai/libs/iai_core_py/tests/repos/mappers/test_static_mappers.py +++ b/interactive_ai/libs/iai_core_py/tests/repos/mappers/test_static_mappers.py @@ -327,6 +327,7 @@ def test_label_mapper(self) -> None: is_empty=True, id_=LabelRepo.generate_id(), is_anomalous=True, + is_background=True, ephemeral=False, ) verify_mongo_mapper(entity_to_map=label, mapper_class=LabelToMongo)