Source code for ddd_pylib.tracking._objects_tracking_operator

from ddd_pylib._base import ImageOperator
import pandas as pd
import xarray as xr
import numpy as np
from._tracking_base import TrackingBase
from ddd_pylib.labels_manip import LabelsToPointsOperator
from ._points_tracking_operator import PointsTrackingOperator




class ObjectsTrackingOperator(TrackingBase, ImageOperator):
    """
    This operator performs tracking of nD+t labeled objects.
    It accepts an image with labeled objects and creates a new identical image with consistent labels 
    across time points, based on the tracking strategy provided.

    This operator is just a wrapper, you need to provide a tracking strategy to use it. 
    For now, the only tracking strategy available is the TrackpyStrategy, which uses the trackpy library to perform tracking of points.

    Example:
        .. literalinclude:: /_examples/objects_tracking_operator.py
           :language: python
    """
    def __init__(self):
        super().__init__()

    def _applyToFrame(self, frameData, t, nT):
        raise RuntimeError("This operator does not support frame-wise processing. Use the 'linkTracks' method instead.")



    def getPrefix(self) -> str:
        return "Tracked-"




    def relabelWithTracks(self, labels: xr.DataArray, linked_df: pd.DataFrame) -> xr.DataArray:
        if (linked_df is None) or ('track_id' not in linked_df.columns) or ('original label' not in linked_df.columns):
            raise ValueError("Linked DataFrame must contain 'track_id' and 'original label' columns.")

        out = []
        for t, g in linked_df.groupby("T"):
            frame = labels.isel(T=t)
            lab = frame.values

            orig = g["original label"].astype(int).values
            track = g["track_id"].astype(int).values

            lut = np.zeros(lab.max() + 1, dtype=np.uint16)
            lut[orig[orig != 0]] = track[orig != 0]  # skip orig_label 0

            tmp = frame.copy()
            tmp.values = lut[lab]
            out.append(tmp)

        res = xr.concat(out, dim="T")
        return res.transpose(*labels.dims)


    def _run(self):
        im = self.getMainImage()
        op = LabelsToPointsOperator()
        op.setMainImage(im)
        yield from op._run()
        df = op.getResult()

        s = self.getStrategy()
        op = PointsTrackingOperator()
        op.setCalibration(im.attrs)
        op.setMainDataFrame(df)
        op.setStrategy(s)
        yield from op._run()

        linked_df = op.getResult()
        self._result = self.relabelWithTracks(im, linked_df)
        self._result.name = self.makeResultName()
        yield 0