SurfaceLoss used for 3D data? about boundary-loss HOT 2 CLOSED

class SurfaceLoss():
    def __init__(self, **kwargs):
        # Self.idc is used to filter out some classes of the target mask. Use fancy indexing
        self.idc: List[int] = kwargs["idc"]
        self.nd: str = kwargs["nd"]  # 'wh' for 2d case, 'whd' for 3d case
        print(f"Initialized {self.__class__.__name__} with {kwargs}")

    def __call__(self, probs: Tensor, dist_maps: Tensor, _: Tensor) -> Tensor:
        assert simplex(probs)
        assert not one_hot(dist_maps)

        pc = probs[:, self.idc, ...].type(torch.float32)
        dc = dist_maps[:, self.idc, ...].type(torch.float32)

        multipled = einsum(f"bk{self.nd},bk{self.nd}->bk{self.nd}", pc, dc)

        loss = multipled.mean()

        return loss

def one_hot2dist(seg: np.ndarray, resolution: Tuple[float, float, float] = None) -> np.ndarray:
    assert one_hot(torch.tensor(seg), axis=0)
    K: int = len(seg)

    res = np.zeros_like(seg)
    for k in range(K):
        posmask = seg[k].astype(np.bool)

        if posmask.any():
            negmask = ~posmask
            res[k] = distance(negmask, sampling=resolution) * negmask \
                - (distance(posmask, sampling=resolution) - 1) * posmask
        # The idea is to leave blank the negative classes
        # since this is one-hot encoded, another class will supervise that pixel

    return res

class SliceDataset(Dataset):
    def __init__(self, filenames: List[str], folders: List[Path], are_hots: List[bool],
                 bounds_generators: List[Callable], transforms: List[Callable], debug=False, quiet=False,
                 K=4, in_memory: bool = False, spacing_dict: Dict[str, Tuple[float, ...]] = None,
                 augment: Optional[Callable] = None, ignore_norm: bool = False,
                 dimensions: int = 2, debug_size: int = 10) -> None:
        self.folders: List[Path] = folders
        self.transforms: List[Callable[[Tuple, int], Callable[[D], Tensor]]] = transforms
        assert len(self.transforms) == len(self.folders)

        self.filenames: List[str] = filenames
        self.K: int = K  # Number of classes

        self.spacing_dict: Optional[Dict[str, Tuple[float, ...]]] = spacing_dict
        if self.spacing_dict:
            assert len(self.spacing_dict) == len(self.filenames)
            print(f"> Spacing dictionnary loaded correctly")

    def __getitem__(self, index: int) -> Dict[str, Union[str, Tensor, List[Tensor], List[Tuple[slice, ...]]]]:
        filename: str = self.filenames[index]
        path_name: Path = Path(filename)
        images: List[D]

        if path_name.suffix == ".png":
            images = [Image.open(files[index]) for files in self.files]
        elif path_name.suffix == ".npy":
            images = [np.load(files[index]) for files in self.files]
        else:
            raise ValueError(filename)

        resolution: Tuple[float, ...]
        if self.spacing_dict:
            resolution = self.spacing_dict[path_name.stem]
        else:
            resolution = tuple([1] * self.dimensions)

        # Final transforms and assertions
        assert len(images) == len(self.folders) == len(self.transforms)
        t_tensors: List[Tensor] = [tr(resolution, self.K)(e) for (tr, e) in zip(self.transforms, images)]
        _, *img_shape = t_tensors[0].shape

        final_tensors = t_tensors

        # [...]

        return {'filenames': filename,
                'images': final_tensors[0],
                'gt': final_tensors[1],
                'labels': final_tensors[2:],
                'spacings': torch.tensor(resolution)}