ngio 0.2.0a2__py3-none-any.whl → 0.5.0b4__py3-none-any.whl

ngio/tables/v1/_roi_table.py

@@ -4,16 +4,30 @@ This class follows the roi_table specification at:
 https://fractal-analytics-platform.github.io/fractal-tasks-core/tables/
 """
 
+import warnings
 from collections.abc import Iterable
 from typing import Literal
+from uuid import uuid4
 
 import pandas as pd
 from pydantic import BaseModel
 
-from ngio.common import WorldCooROI
-from ngio.tables._validators import validate_columns
-from ngio.tables.backends import ImplementedTableBackends
-from ngio.utils import NgioValueError, ZarrGroupHandler
+from ngio.common import Roi
+from ngio.tables._abstract_table import (
+    AbstractBaseTable,
+    TabularData,
+)
+from ngio.tables.backends import (
+    BackendMeta,
+    TableBackend,
+    convert_to_pandas,
+    normalize_pandas_df,
+)
+from ngio.utils import (
+    NgioTableValidationError,
+    NgioValueError,
+    ZarrGroupHandler,
+)
 
 REQUIRED_COLUMNS = [
     "x_micrometer",
@@ -24,100 +38,274 @@ REQUIRED_COLUMNS = [
     "len_z_micrometer",
 ]
 
+#####################
+# Optional columns are not validated at the moment
+# only a warning is raised if non optional columns are present
+#####################
+
+TIME_COLUMNS = [
+    "t_second",
+    "len_t_second",
+]
 
 ORIGIN_COLUMNS = [
     "x_micrometer_original",
     "y_micrometer_original",
+    "z_micrometer_original",
 ]
 
 TRANSLATION_COLUMNS = ["translation_x", "translation_y", "translation_z"]
 
-OPTIONAL_COLUMNS = ORIGIN_COLUMNS + TRANSLATION_COLUMNS
+PLATE_COLUMNS = [
+    "plate_name",
+    "row",
+    "column",
+    "path_in_well",
+    "path_in_plate",
+    "acquisition_id",
+    "acquisition_name",
+]
+
+INDEX_COLUMNS = [
+    "FieldIndex",
+    "label",
+]
+
+OPTIONAL_COLUMNS = ORIGIN_COLUMNS + TRANSLATION_COLUMNS + PLATE_COLUMNS + INDEX_COLUMNS
+
+
+def _check_optional_columns(col_name: str) -> None:
+    """Check if the column name is in the optional columns."""
+    if col_name not in OPTIONAL_COLUMNS + TIME_COLUMNS:
+        warnings.warn(
+            f"Column {col_name} is not in the optional columns.", stacklevel=2
+        )
 
 
-def _dataframe_to_rois(dataframe: pd.DataFrame) -> dict[str, WorldCooROI]:
+def _dataframe_to_rois(
+    dataframe: pd.DataFrame,
+    required_columns: list[str] = REQUIRED_COLUMNS,
+) -> dict[str, Roi]:
     """Convert a DataFrame to a WorldCooROI object."""
+    # Validate the columns of the DataFrame
+    _missing_columns = set(required_columns).difference(set(dataframe.columns))
+    if len(_missing_columns) != 0:
+        raise NgioTableValidationError(
+            f"Could not find required columns: {_missing_columns} in the table."
+        )
+
+    extra_columns = set(dataframe.columns).difference(
+        set(required_columns + TIME_COLUMNS)
+    )
+
+    for col in extra_columns:
+        _check_optional_columns(col)
+
+    label_is_index = True if dataframe.index.name == "label" else False
+
+    extras = {}
+
     rois = {}
-    for key, row in dataframe.iterrows():
+    for row in dataframe.itertuples(index=True):
         # check if optional columns are present
-        origin = {col: row.get(col, None) for col in ORIGIN_COLUMNS}
-        origin = dict(filter(lambda x: x[1] is not None, origin.items()))
-        translation = {col: row.get(col, None) for col in TRANSLATION_COLUMNS}
-        translation = dict(filter(lambda x: x[1] is not None, translation.items()))
-
-        roi = WorldCooROI(
-            name=str(key),
-            x=row["x_micrometer"],
-            y=row["y_micrometer"],
-            z=row["z_micrometer"],
-            x_length=row["len_x_micrometer"],
-            y_length=row["len_y_micrometer"],
-            z_length=row["len_z_micrometer"],
-            unit="micrometer",  # type: ignore
-            **origin,
-            **translation,
+        if len(extra_columns) > 0:
+            extras = {col: getattr(row, col, None) for col in extra_columns}
+
+        z_micrometer = getattr(row, "z_micrometer", None)
+        z_length_micrometer = getattr(row, "len_z_micrometer", None)
+
+        t_second = getattr(row, "t_second", None)
+        t_length_second = getattr(row, "len_t_second", None)
+
+        if label_is_index:
+            label = int(row.Index)  # type: ignore (type can not be known here, but should be castable to int)
+        else:
+            label = getattr(row, "label", None)
+
+        slices = {
+            "x": (row.x_micrometer, row.len_x_micrometer),
+            "y": (row.y_micrometer, row.len_y_micrometer),
+            "z": (z_micrometer, z_length_micrometer),
+        }
+        if t_second is not None or t_length_second is not None:
+            slices["t"] = (t_second, t_length_second)
+        roi = Roi.from_values(
+            name=str(row.Index),
+            slices=slices,
+            space="world",
+            label=label,
+            **extras,
         )
         rois[roi.name] = roi
     return rois
 
 
-def _rois_to_dataframe(rois: dict[str, WorldCooROI], index_key: str) -> pd.DataFrame:
+def _rois_to_dataframe(rois: dict[str, Roi], index_key: str | None) -> pd.DataFrame:
     """Convert a list of WorldCooROI objects to a DataFrame."""
     data = []
     for roi in rois.values():
+        # This normalization is necessary for backward compatibility
+        if roi.space != "world":
+            raise NotImplementedError(
+                "Only ROIs in world coordinates can be serialized."
+            )
+
+        z_slice = roi.get("z")
+        if z_slice is None:
+            z_micrometer = 0.0
+            len_z_micrometer = 1.0
+        else:
+            z_micrometer = z_slice.start if z_slice.start is not None else 0.0
+            len_z_micrometer = z_slice.length if z_slice.length is not None else 1.0
+
+        x_slice = roi.get("x")
+        if x_slice is None:
+            raise NgioValueError("ROI is missing 'x' slice.")
+        y_slice = roi.get("y")
+        if y_slice is None:
+            raise NgioValueError("ROI is missing 'y' slice.")
         row = {
-            index_key: roi.name,
-            "x_micrometer": roi.x,
-            "y_micrometer": roi.y,
-            "z_micrometer": roi.z,
-            "len_x_micrometer": roi.x_length,
-            "len_y_micrometer": roi.y_length,
-            "len_z_micrometer": roi.z_length,
+            index_key: roi.get_name(),
+            "x_micrometer": x_slice.start if x_slice.start is not None else 0.0,
+            "y_micrometer": y_slice.start if y_slice.start is not None else 0.0,
+            "z_micrometer": z_micrometer,
+            "len_x_micrometer": x_slice.length if x_slice.length is not None else 1.0,
+            "len_y_micrometer": y_slice.length if y_slice.length is not None else 1.0,
+            "len_z_micrometer": len_z_micrometer,
         }
 
-        extra = roi.model_extra or {}
-        for col in ORIGIN_COLUMNS:
-            if col in extra:
-                row[col] = extra[col]
+        t_slice = roi.get("t")
+        if t_slice is not None:
+            row["t_second"] = t_slice.start if t_slice.start is not None else 0.0
+            row["len_t_second"] = t_slice.length if t_slice.length is not None else 1.0
 
-        for col in TRANSLATION_COLUMNS:
-            if col in extra:
-                row[col] = extra[col]
+        if roi.label is not None and index_key != "label":
+            row["label"] = roi.label
+
+        extra = roi.model_extra or {}
+        for col in extra:
+            _check_optional_columns(col)
+            row[col] = extra[col]
         data.append(row)
+
     dataframe = pd.DataFrame(data)
-    dataframe = dataframe.set_index(index_key)
+    dataframe = normalize_pandas_df(dataframe, index_key=index_key)
     return dataframe
 
 
-class ROITableV1Meta(BaseModel):
-    """Metadata for the ROI table."""
+class RoiDictWrapper:
+    """A wrapper for a dictionary of ROIs to provide a consistent interface."""
 
-    fractal_table_version: Literal["1"] = "1"
-    type: Literal["roi_table"] = "roi_table"
-    backend: str | None = None
+    def __init__(self, rois: Iterable[Roi]) -> None:
+        self._rois_by_name = {}
+        self._rois_by_label = {}
+        for roi in rois:
+            name = roi.name
+            if name in self._rois_by_name:
+                name = f"{name}_{uuid4().hex[:8]}"
+            self._rois_by_name[name] = roi
+            if roi.label is not None:
+                self._rois_by_label[roi.label] = roi
 
+    def get_by_name(self, name: str, default: Roi | None = None) -> Roi | None:
+        """Get an ROI by its name."""
+        return self._rois_by_name.get(name, default)
 
-class RoiTableV1:
-    """Class to handle fractal ROI tables.
+    def get_by_label(self, label: int, default: Roi | None = None) -> Roi | None:
+        """Get an ROI by its label."""
+        return self._rois_by_label.get(label, default)
 
-    To know more about the ROI table format, please refer to the
-    specification at:
-    https://fractal-analytics-platform.github.io/fractal-tasks-core/tables/
-    """
+    def _add_roi(self, roi: Roi, overwrite: bool = False) -> None:
+        """Add an ROI to the wrapper."""
+        if roi.name in self._rois_by_name and not overwrite:
+            raise NgioValueError(f"ROI with name {roi.name} already exists.")
 
-    def __init__(self, rois: Iterable[WorldCooROI] | None = None) -> None:
-        """Create a new ROI table."""
-        self._meta = ROITableV1Meta()
-        self._table_backend = None
+        self._rois_by_name[roi.name] = roi
+        if roi.label is not None:
+            self._rois_by_label[roi.label] = roi
+
+    def add_rois(self, rois: Roi | Iterable[Roi], overwrite: bool = False) -> None:
+        """Add ROIs to the wrapper."""
+        if isinstance(rois, Roi):
+            rois = [rois]
+
+        for roi in rois:
+            self._add_roi(roi, overwrite=overwrite)
+
+    def to_list(self) -> list[Roi]:
+        """Return the list of ROIs."""
+        return list(self._rois_by_name.values())
+
+    def to_dataframe(self, index_key: str | None = None) -> pd.DataFrame:
+        """Convert the ROIs to a DataFrame."""
+        return _rois_to_dataframe(self._rois_by_name, index_key=index_key)
 
-        self._rois = {}
+    @classmethod
+    def from_dataframe(
+        cls, dataframe: pd.DataFrame, required_columns: list[str] = REQUIRED_COLUMNS
+    ) -> "RoiDictWrapper":
+        """Create a RoiDictWrapper from a DataFrame."""
+        rois = _dataframe_to_rois(dataframe, required_columns=required_columns)
+        return cls(rois.values())
+
+
+def _table_to_rois(
+    table: TabularData,
+    index_key: str | None = None,
+    index_type: Literal["int", "str"] | None = None,
+    required_columns: list[str] = REQUIRED_COLUMNS,
+) -> tuple[pd.DataFrame, RoiDictWrapper]:
+    """Convert a table to a dictionary of ROIs.
+
+    Args:
+        table: The table to convert.
+        index_key: The column name to use as the index of the DataFrame.
+        index_type: The type of the index column in the DataFrame.
+        required_columns: The required columns in the DataFrame.
+
+    Returns:
+        A tuple containing the DataFrame and a RoiDictWrapper with the ROIs.
+    """
+    dataframe = convert_to_pandas(
+        table,
+        index_key=index_key,
+        index_type=index_type,
+    )
+    roi_dict_wrapper = RoiDictWrapper.from_dataframe(
+        dataframe, required_columns=required_columns
+    )
+    return dataframe, roi_dict_wrapper
+
+
+class GenericRoiTableV1(AbstractBaseTable):
+    def __init__(
+        self,
+        *,
+        rois: Iterable[Roi] | None = None,
+        meta: BackendMeta,
+        required_columns: list[str] = REQUIRED_COLUMNS,
+    ) -> None:
+        table = None
+
+        self._rois: RoiDictWrapper | None = None
         if rois is not None:
-            self.add(rois)
+            self._rois = RoiDictWrapper(rois)
+            table = self._rois.to_dataframe(index_key=meta.index_key)
+
+        self._required_columns = required_columns
+        super().__init__(table_data=table, meta=meta)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the table."""
+        rois = self.rois()
+        prop = f"num_rois={len(rois)}"
+        class_name = self.__class__.__name__
+        return f"{class_name}({prop})"
 
     @staticmethod
-    def type() -> Literal["roi_table"]:
+    def table_type() -> str:
         """Return the type of the table."""
-        return "roi_table"
+        return "generic_roi_table"
 
     @staticmethod
     def version() -> Literal["1"]:
@@ -125,102 +313,263 @@ class RoiTableV1:
         return "1"
 
     @property
-    def backend_name(self) -> str | None:
-        """Return the name of the backend."""
-        if self._table_backend is None:
-            return None
-        return self._table_backend.backend_name()
+    def table_data(self) -> TabularData:
+        """Return the table."""
+        if self._rois is None:
+            return super().table_data
 
-    @classmethod
-    def _from_handler(
-        cls, handler: ZarrGroupHandler, backend_name: str | None = None
-    ) -> "RoiTableV1":
-        """Create a new ROI table from a Zarr store."""
-        meta = ROITableV1Meta(**handler.load_attrs())
-
-        if backend_name is None:
-            backend = ImplementedTableBackends().get_backend(
-                backend_name=meta.backend,
-                group_handler=handler,
-                index_key="FieldIndex",
-                index_type="str",
-            )
-        else:
-            backend = ImplementedTableBackends().get_backend(
-                backend_name=backend_name,
-                group_handler=handler,
-                index_key="FieldIndex",
-                index_type="str",
+        if len(self.rois()) > 0:
+            self._table_data = self._rois.to_dataframe(index_key=self.meta.index_key)
+        return super().table_data
+
+    def set_table_data(
+        self, table_data: TabularData | None = None, refresh: bool = False
+    ) -> None:
+        if table_data is not None:
+            if not isinstance(table_data, TabularData):
+                raise NgioValueError(
+                    "The table must be a pandas DataFrame, polars LazyFrame, "
+                    " or AnnData object."
+                )
+
+            table_data, rois = _table_to_rois(
+                table_data,
+                index_key=self.index_key,
+                index_type=self.index_type,
+                required_columns=REQUIRED_COLUMNS,
             )
-            meta.backend = backend_name
+            self._table_data = table_data
+            self._rois = rois
+            return None
 
-        if not backend.implements_dataframe:
+        if self._table_data is not None and not refresh:
+            return None
+
+        if self._table_backend is None:
             raise NgioValueError(
-                "The backend does not implement the dataframe protocol."
+                "The table does not have a DataFrame in memory nor a backend."
             )
 
-        table = cls()
-        table._meta = meta
-        table._table_backend = backend
+        table_data, rois = _table_to_rois(
+            self._table_backend.load(),
+            index_key=self.index_key,
+            index_type=self.index_type,
+            required_columns=REQUIRED_COLUMNS,
+        )
+        self._table_data = table_data
+        self._rois = rois
+
+    def _check_rois(self) -> None:
+        """Load the ROIs from the table.
+
+        If the ROIs are already loaded, do nothing.
+        If the ROIs are not loaded, load them from the table.
+        """
+        if self._rois is None:
+            self._rois = RoiDictWrapper.from_dataframe(
+                self.dataframe, required_columns=self._required_columns
+            )
+
+    def rois(self) -> list[Roi]:
+        """List all ROIs in the table."""
+        self._check_rois()
+        if self._rois is None:
+            return []
+        return self._rois.to_list()
+
+    def add(self, roi: Roi | Iterable[Roi], overwrite: bool = False) -> None:
+        """Append ROIs to the current table.
+
+        Args:
+            roi: A single ROI or a list of ROIs to add to the table.
+            overwrite: If True, overwrite existing ROIs with the same name.
+        """
+        if isinstance(roi, Roi):
+            roi = [roi]
+
+        self._check_rois()
+        if self._rois is None:
+            self._rois = RoiDictWrapper([])
 
-        dataframe = backend.load_as_dataframe()
-        dataframe = validate_columns(
-            dataframe,
+        self._rois.add_rois(roi, overwrite=overwrite)
+
+    def get(self, roi_name: str) -> Roi:
+        """Get an ROI from the table."""
+        self._check_rois()
+        if self._rois is None:
+            self._rois = RoiDictWrapper([])
+
+        roi = self._rois.get_by_name(roi_name)
+        if roi is None:
+            raise NgioValueError(f"ROI with name {roi_name} not found in the table.")
+        return roi
+
+    @classmethod
+    def from_table_data(
+        cls, table_data: TabularData, meta: BackendMeta
+    ) -> "GenericRoiTableV1":
+        """Create a new ROI table from a table data."""
+        _, rois = _table_to_rois(
+            table=table_data,
+            index_key=meta.index_key,
+            index_type=meta.index_type,
             required_columns=REQUIRED_COLUMNS,
-            optional_columns=OPTIONAL_COLUMNS,
         )
-        table._rois = _dataframe_to_rois(dataframe)
-        return table
+        return cls(rois=rois.to_list(), meta=meta)
 
-    def _set_backend(
-        self,
-        handler: ZarrGroupHandler,
-        backend_name: str | None = None,
+
+class RoiTableV1Meta(BackendMeta):
+    """Metadata for the ROI table."""
+
+    table_version: Literal["1"] = "1"
+    type: Literal["roi_table"] = "roi_table"
+    index_key: str | None = "FieldIndex"
+    index_type: Literal["str", "int"] | None = "str"
+
+
+class RoiTableV1(GenericRoiTableV1):
+    """Class to handle fractal ROI tables.
+
+    To know more about the ROI table format, please refer to the
+    specification at:
+    https://fractal-analytics-platform.github.io/fractal-tasks-core/tables/
+    """
+
+    def __init__(
+        self, rois: Iterable[Roi] | None = None, *, meta: RoiTableV1Meta | None = None
     ) -> None:
-        """Set the backend of the table."""
-        backend = ImplementedTableBackends().get_backend(
-            backend_name=backend_name,
-            group_handler=handler,
-            index_key="FieldIndex",
-            index_type="str",
+        """Create a new ROI table."""
+        if meta is None:
+            meta = RoiTableV1Meta()
+
+        if meta.index_key is None:
+            meta.index_key = "FieldIndex"
+
+        if meta.index_type is None:
+            meta.index_type = "str"
+        super().__init__(meta=meta, rois=rois)
+
+    @classmethod
+    def from_handler(
+        cls,
+        handler: ZarrGroupHandler,
+        backend: TableBackend | None = None,
+    ) -> "RoiTableV1":
+        table = cls._from_handler(
+            handler=handler,
+            backend=backend,
+            meta_model=RoiTableV1Meta,
         )
-        self._meta.backend = backend_name
-        self._table_backend = backend
+        return table
 
-    def rois(self) -> list[WorldCooROI]:
-        """List all ROIs in the table."""
-        return list(self._rois.values())
+    @staticmethod
+    def table_type() -> Literal["roi_table"]:
+        """Return the type of the table."""
+        return "roi_table"
 
-    def get(self, roi_name: str) -> WorldCooROI:
-        """Get an ROI from the table."""
-        if roi_name not in self._rois:
-            raise NgioValueError(f"ROI {roi_name} not found in the table.")
-        return self._rois[roi_name]
 
-    def add(self, roi: WorldCooROI | Iterable[WorldCooROI]) -> None:
-        """Append ROIs to the current table."""
-        if isinstance(roi, WorldCooROI):
-            roi = [roi]
+class RegionMeta(BaseModel):
+    """Metadata for the region."""
 
-        for _roi in roi:
-            if _roi.name in self._rois:
-                raise NgioValueError(f"ROI {_roi.name} already exists in the table.")
-            self._rois[_roi.name] = _roi
+    path: str
 
-    def consolidate(self) -> None:
-        """Write the current state of the table to the Zarr file."""
-        if self._table_backend is None:
+
+class MaskingRoiTableV1Meta(BackendMeta):
+    """Metadata for the ROI table."""
+
+    table_version: Literal["1"] = "1"
+    type: Literal["masking_roi_table"] = "masking_roi_table"
+    region: RegionMeta | None = None
+    instance_key: str = "label"
+    index_key: str | None = "label"
+    index_type: Literal["int", "str"] | None = "int"
+
+
+class MaskingRoiTableV1(GenericRoiTableV1):
+    """Class to handle fractal ROI tables.
+
+    To know more about the ROI table format, please refer to the
+    specification at:
+    https://fractal-analytics-platform.github.io/fractal-tasks-core/tables/
+    """
+
+    def __init__(
+        self,
+        rois: Iterable[Roi] | None = None,
+        *,
+        reference_label: str | None = None,
+        meta: MaskingRoiTableV1Meta | None = None,
+    ) -> None:
+        """Create a new ROI table."""
+        if meta is None:
+            meta = MaskingRoiTableV1Meta()
+
+        if reference_label is not None:
+            path = f"../labels/{reference_label}"
+            meta.region = RegionMeta(path=path)
+
+        if meta.index_key is None:
+            meta.index_key = "label"
+
+        if meta.index_type is None:
+            meta.index_type = "int"
+        meta.instance_key = meta.index_key
+        super().__init__(meta=meta, rois=rois)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the table."""
+        rois = self.rois()
+        if self.reference_label is not None:
+            prop = f"num_rois={len(rois)}, reference_label={self.reference_label}"
+        else:
+            prop = f"num_rois={len(rois)}"
+        return f"MaskingRoiTableV1({prop})"
+
+    @classmethod
+    def from_handler(
+        cls,
+        handler: ZarrGroupHandler,
+        backend: TableBackend | None = None,
+    ) -> "MaskingRoiTableV1":
+        table = cls._from_handler(
+            handler=handler,
+            backend=backend,
+            meta_model=MaskingRoiTableV1Meta,
+        )
+        return table
+
+    @staticmethod
+    def table_type() -> Literal["masking_roi_table"]:
+        """Return the type of the table."""
+        return "masking_roi_table"
+
+    @property
+    def meta(self) -> MaskingRoiTableV1Meta:
+        """Return the metadata of the table."""
+        if not isinstance(self._meta, MaskingRoiTableV1Meta):
             raise NgioValueError(
-                "No backend set for the table. "
-                "Please add the table to a OME-Zarr Image before calling consolidate."
+                "The metadata of the table is not of type MaskingRoiTableV1Meta."
             )
+        return self._meta
 
-        dataframe = _rois_to_dataframe(self._rois, index_key="FieldIndex")
-        dataframe = validate_columns(
-            dataframe,
-            required_columns=REQUIRED_COLUMNS,
-            optional_columns=OPTIONAL_COLUMNS,
-        )
-        self._table_backend.write_from_dataframe(
-            dataframe, metadata=self._meta.model_dump(exclude_none=True)
-        )
+    @property
+    def reference_label(self) -> str | None:
+        """Return the reference label."""
+        path = self.meta.region
+        if path is None:
+            return None
+
+        path = path.path
+        path = path.split("/")[-1]
+        return path
+
+    def get_label(self, label: int) -> Roi:
+        """Get an ROI by label."""
+        self._check_rois()
+        if self._rois is None:
+            self._rois = RoiDictWrapper([])
+        roi = self._rois.get_by_label(label)
+
+        if roi is None:
+            raise NgioValueError(f"ROI with label {label} not found.")
+        return roi