import nibabel as nib
from core.node import Node
from core.contracts_radiomics import MaskDataContract
from pathlib import Path


class MaskReader(Node):
    def __init__(self, path: str, geometry_id: str):
        self.path = Path(path)

        img = nib.load(str(self.path))
        dim = "3D" if img.ndim == 3 else "2D"

        super().__init__(
            name="MaskReader",
            outputs={
                "mask": MaskDataContract(
                    modality="CT",
                    dim=dim,
                    geometry_id=geometry_id
                )
            }
        )

    def load_data(self):
        img = nib.load(str(self.path))
        return img.get_fdata()