Source code for ddd_pylib.labels_manip._kill_border_labels

from ._replace_values_operator import ReplaceValuesOperator
from ._list_border_labels_operator import ListBorderValuesOperator
from typing import List, Dict
import pandas as pd




class KillBorderLabelsOperator(ReplaceValuesOperator):
    """
    Given a list of axes and a border thickness, this operator identifies all unique label 
    values that touch the borders of the image along those axes and replaces them with 0.
    This operator works with tracked labels: if a label is removed at time t, it will also 
    be removed at all other time points.
    """
    def __init__(self):
        super().__init__()
        self._borderAxes = self.defaultBorderAxes()
        self._borderSize = self.defaultBorderSize()



    @staticmethod
    def defaultBorderSize() -> int:
        return 1




    @staticmethod
    def defaultBorderAxes() -> List[str]:
        return ["Z", "X", "Y"]




    def setBorderSize(self, size: int):
        if size < 1:
            raise ValueError("Border size must be at least 1.")
        self._borderSize = size




    def setBorderAxes(self, axes: List[str]):
        if not axes:
            raise ValueError("Border axes list cannot be empty.")
        self._borderAxes = axes




    def getPrefix(self):
        return "KillBorders-"


    def _makeKillDict(self, border_values_df: pd.DataFrame) -> Dict[int, int]:
        """
        Create a dictionary mapping border label values to 0 (to be removed).
        """
        img = self.getMainImage()
        d = img.dtype.type
        kill_dict = {}
        columns_list = [col for col in border_values_df.columns if col.startswith(('+', '-'))]
        for col in columns_list:
            unique_labels = [d(lbl) for lbl in border_values_df[col].dropna().unique()]
            for label in unique_labels:
                if label != 0:
                    kill_dict[label] = 0
        return kill_dict

    def _run(self):
        op = ListBorderValuesOperator()
        op.setMainImage(self.getMainImage())
        op.setBorderSize(self._borderSize)
        op.setBorderAxes(self._borderAxes)
        yield from op._run()
        self.setValuesDict(self._makeKillDict(op.getResult()))
        yield from super()._run()