scripts/utils.py

# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
import copy
import json
import math
import os
import zipfile
from argparse import Namespace
from datetime import timedelta
from typing import Any, Sequence

import numpy as np
import skimage
import torch
import torch.distributed as dist
from monai.bundle import ConfigParser
from monai.config import DtypeLike, NdarrayOrTensor
from monai.data import CacheDataset, DataLoader, partition_dataset
from monai.transforms import Compose, EnsureTyped, Lambdad, LoadImaged, Orientationd
from monai.transforms.utils_morphological_ops import dilate, erode
from monai.utils import TransformBackends, convert_data_type, convert_to_dst_type, get_equivalent_dtype
from scipy import stats
from torch import Tensor


def unzip_dataset(dataset_dir):
    if dist.is_available() and dist.is_initialized():
        rank = dist.get_rank()
    else:
        rank = 0

    if rank == 0:
        if not os.path.exists(dataset_dir):
            zip_file_path = dataset_dir + ".zip"
            if not os.path.isfile(zip_file_path):
                raise ValueError(f"Please download {zip_file_path}.")
            with zipfile.ZipFile(zip_file_path, "r") as zip_ref:
                zip_ref.extractall(path=os.path.dirname(dataset_dir))
            print(f"Unzipped {zip_file_path} to {dataset_dir}.")

    if dist.is_available() and dist.is_initialized():
        dist.barrier()  # Synchronize all processes

    return


def add_data_dir2path(list_files: list, data_dir: str, fold: int = None) -> tuple[list, list]:
    """
    Read a list of data dictionary.

    Args:
        list_files (list): input data to load and transform to generate dataset for model.
        data_dir (str): directory of files.
        fold (int, optional): fold index for cross validation. Defaults to None.

    Returns:
        tuple[list, list]: A tuple of two arrays (training, validation).
    """
    new_list_files = copy.deepcopy(list_files)
    if fold is not None:
        new_list_files_train = []
        new_list_files_val = []
    for d in new_list_files:
        d["image"] = os.path.join(data_dir, d["image"])

        if "label" in d:
            d["label"] = os.path.join(data_dir, d["label"])

        if fold is not None:
            if d["fold"] == fold:
                new_list_files_val.append(copy.deepcopy(d))
            else:
                new_list_files_train.append(copy.deepcopy(d))

    if fold is not None:
        return new_list_files_train, new_list_files_val
    else:
        return new_list_files, []


def maisi_datafold_read(json_list, data_base_dir, fold=None):
    with open(json_list, "r") as f:
        filenames_train = json.load(f)["training"]
    # training data
    train_files, val_files = add_data_dir2path(filenames_train, data_base_dir, fold=fold)
    print(f"dataset: {data_base_dir}, num_training_files: {len(train_files)}, num_val_files: {len(val_files)}")
    return train_files, val_files


def remap_labels(mask, label_dict_remap_json):
    """
    Remap labels in the mask according to the provided label dictionary.

    This function reads a JSON file containing label mapping information and applies
    the mapping to the input mask.

    Args:
        mask (Tensor): The input mask tensor to be remapped.
        label_dict_remap_json (str): Path to the JSON file containing the label mapping dictionary.

    Returns:
        Tensor: The remapped mask tensor.
    """
    with open(label_dict_remap_json, "r") as f:
        mapping_dict = json.load(f)
    mapper = MapLabelValue(
        orig_labels=[pair[0] for pair in mapping_dict.values()],
        target_labels=[pair[1] for pair in mapping_dict.values()],
        dtype=torch.uint8,
    )
    return mapper(mask[0, ...])[None, ...].to(mask.device)


def get_index_arr(img):
    """
    Generate an index array for the given image.

    This function creates a 3D array of indices corresponding to the dimensions of the input image.

    Args:
        img (ndarray): The input image array.

    Returns:
        ndarray: A 3D array containing the indices for each dimension of the input image.
    """
    return np.moveaxis(
        np.moveaxis(
            np.stack(np.meshgrid(np.arange(img.shape[0]), np.arange(img.shape[1]), np.arange(img.shape[2]))), 0, 3
        ),
        0,
        1,
    )


def supress_non_largest_components(img, target_label, default_val=0):
    """
    Suppress all components except the largest one(s) for specified target labels.

    This function identifies the largest component(s) for each target label and
    suppresses all other smaller components.

    Args:
        img (ndarray): The input image array.
        target_label (list): List of label values to process.
        default_val (int, optional): Value to assign to suppressed voxels. Defaults to 0.

    Returns:
        tuple: A tuple containing:
            - ndarray: Modified image with non-largest components suppressed.
            - int: Number of voxels that were changed.
    """
    index_arr = get_index_arr(img)
    img_mod = copy.deepcopy(img)
    new_background = np.zeros(img.shape, dtype=np.bool_)
    for label in target_label:
        label_cc = skimage.measure.label(img == label, connectivity=3)
        uv, uc = np.unique(label_cc, return_counts=True)
        dominant_vals = uv[np.argsort(uc)[::-1][:2]]
        if len(dominant_vals) >= 2:  # Case: no predictions
            new_background = np.logical_or(
                new_background,
                np.logical_not(np.logical_or(label_cc == dominant_vals[0], label_cc == dominant_vals[1])),
            )

    for voxel in index_arr[new_background]:
        img_mod[tuple(voxel)] = default_val
    diff = np.sum((img - img_mod) > 0)

    return img_mod, diff


def erode_one_img(mask_t: Tensor, filter_size: int | Sequence[int] = 3, pad_value: float = 1.0) -> Tensor:
    """
    Erode 2D/3D binary mask with data type as torch tensor.

    Args:
        mask_t: input 2D/3D binary mask, [M,N] or [M,N,P] torch tensor.
        filter_size: erosion filter size, has to be odd numbers, default to be 3.
        pad_value: the filled value for padding. We need to pad the input before filtering
                   to keep the output with the same size as input. Usually use default value
                   and not changed.

    Return:
        Tensor: eroded mask, same shape as input.
    """
    return erode(mask_t.float().unsqueeze(0).unsqueeze(0), filter_size, pad_value=pad_value).squeeze(0).squeeze(0)


def dilate_one_img(mask_t: Tensor, filter_size: int | Sequence[int] = 3, pad_value: float = 0.0) -> Tensor:
    """
    Dilate 2D/3D binary mask with data type as torch tensor.

    Args:
        mask_t: input 2D/3D binary mask, [M,N] or [M,N,P] torch tensor.
        filter_size: dilation filter size, has to be odd numbers, default to be 3.
        pad_value: the filled value for padding. We need to pad the input before filtering
                   to keep the output with the same size as input. Usually use default value
                   and not changed.

    Return:
        Tensor: dilated mask, same shape as input.
    """
    return dilate(mask_t.float().unsqueeze(0).unsqueeze(0), filter_size, pad_value=pad_value).squeeze(0).squeeze(0)


def binarize_labels(x: Tensor, bits: int = 8) -> Tensor:
    """
    Convert input tensor to binary representation.

    This function takes an input tensor and converts it to a binary representation
    using the specified number of bits.

    Args:
        x (Tensor): Input tensor with shape (B, 1, H, W, D).
        bits (int, optional): Number of bits to use for binary representation. Defaults to 8.

    Returns:
        Tensor: Binary representation of the input tensor with shape (B, bits, H, W, D).
    """
    mask = 2 ** torch.arange(bits).to(x.device, x.dtype)
    return x.unsqueeze(-1).bitwise_and(mask).ne(0).byte().squeeze(1).permute(0, 4, 1, 2, 3)


def setup_ddp(rank: int, world_size: int) -> torch.device:
    """
    Initialize the distributed process group.

    Args:
        rank (int): rank of the current process.
        world_size (int): number of processes participating in the job.

     Returns:
        torch.device: device of the current process.
    """
    dist.init_process_group(
        backend="nccl", init_method="env://", timeout=timedelta(seconds=36000), rank=rank, world_size=world_size
    )
    dist.barrier()
    device = torch.device(f"cuda:{rank}")
    return device


def define_instance(args: Namespace, instance_def_key: str) -> Any:
    """
    Define and instantiate an object based on the provided arguments and instance definition key.

    This function uses a ConfigParser to parse the arguments and instantiate an object
    defined by the instance_def_key.

    Args:
        args: An object containing the arguments to be parsed.
        instance_def_key (str): The key used to retrieve the instance definition from the parsed content.

    Returns:
        The instantiated object as defined by the instance_def_key in the parsed configuration.
    """
    parser = ConfigParser(vars(args))
    parser.parse(True)
    return parser.get_parsed_content(instance_def_key, instantiate=True)


def prepare_maisi_controlnet_json_dataloader(
    json_data_list: list | str,
    data_base_dir: list | str,
    batch_size: int = 1,
    fold: int = 0,
    cache_rate: float = 0.0,
    rank: int = 0,
    world_size: int = 1,
) -> tuple[DataLoader, DataLoader]:
    """
    Prepare dataloaders for training and validation.

    Args:
        json_data_list (list | str): the name of JSON files listing the data.
        data_base_dir (list | str): directory of files.
        batch_size (int, optional): how many samples per batch to load . Defaults to 1.
        fold (int, optional): fold index for cross validation. Defaults to 0.
        cache_rate (float, optional): percentage of cached data in total. Defaults to 0.0.
        rank (int, optional): rank of the current process. Defaults to 0.
        world_size (int, optional): number of processes participating in the job. Defaults to 1.

    Returns:
        tuple[DataLoader, DataLoader]:  A tuple of two dataloaders (training, validation).
    """
    use_ddp = world_size > 1
    if isinstance(json_data_list, list):
        assert isinstance(data_base_dir, list)
        list_train = []
        list_valid = []
        for data_list, data_root in zip(json_data_list, data_base_dir):
            with open(data_list, "r") as f:
                json_data = json.load(f)["training"]
            train, val = add_data_dir2path(json_data, data_root, fold)
            list_train += train
            list_valid += val
    else:
        with open(json_data_list, "r") as f:
            json_data = json.load(f)["training"]
        list_train, list_valid = add_data_dir2path(json_data, data_base_dir, fold)

    common_transform = [
        LoadImaged(keys=["image", "label"], image_only=True, ensure_channel_first=True),
        Orientationd(keys=["label"], axcodes="RAS"),
        EnsureTyped(keys=["label"], dtype=torch.uint8, track_meta=True),
        Lambdad(keys="top_region_index", func=lambda x: torch.FloatTensor(x)),
        Lambdad(keys="bottom_region_index", func=lambda x: torch.FloatTensor(x)),
        Lambdad(keys="spacing", func=lambda x: torch.FloatTensor(x)),
        Lambdad(keys=["top_region_index", "bottom_region_index", "spacing"], func=lambda x: x * 1e2),
    ]
    train_transforms, val_transforms = Compose(common_transform), Compose(common_transform)

    train_loader = None

    if use_ddp:
        list_train = partition_dataset(data=list_train, shuffle=True, num_partitions=world_size, even_divisible=True)[
            rank
        ]
    train_ds = CacheDataset(data=list_train, transform=train_transforms, cache_rate=cache_rate, num_workers=8)
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)
    if use_ddp:
        list_valid = partition_dataset(data=list_valid, shuffle=True, num_partitions=world_size, even_divisible=False)[
            rank
        ]
    val_ds = CacheDataset(data=list_valid, transform=val_transforms, cache_rate=cache_rate, num_workers=8)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=False)
    return train_loader, val_loader


def organ_fill_by_closing(data, target_label, device, close_times=2, filter_size=3, pad_value=0.0):
    """
    Fill holes in an organ mask using morphological closing operations.

    This function performs a series of dilation and erosion operations to fill holes
    in the organ mask identified by the target label.

    Args:
        data (ndarray): The input data containing organ labels.
        target_label (int): The label of the organ to be processed.
        device (str): The device to perform the operations on (e.g., 'cuda:0').
        close_times (int, optional): Number of times to perform the closing operation. Defaults to 2.
        filter_size (int, optional): Size of the filter for dilation and erosion. Defaults to 3.
        pad_value (float, optional): Value used for padding in dilation and erosion. Defaults to 0.0.

    Returns:
        ndarray: Boolean mask of the filled organ.
    """
    mask = (data == target_label).astype(np.uint8)
    mask = torch.from_numpy(mask).to(device)
    for _ in range(close_times):
        mask = dilate_one_img(mask, filter_size=filter_size, pad_value=pad_value)
        mask = erode_one_img(mask, filter_size=filter_size, pad_value=pad_value)
    return mask.cpu().numpy().astype(np.bool_)


def organ_fill_by_removed_mask(data, target_label, remove_mask, device):
    """
    Fill an organ mask in regions where it was previously removed.

    Args:
        data (ndarray): The input data containing organ labels.
        target_label (int): The label of the organ to be processed.
        remove_mask (ndarray): Boolean mask indicating regions where the organ was removed.
        device (str): The device to perform the operations on (e.g., 'cuda:0').

    Returns:
        ndarray: Boolean mask of the filled organ in previously removed regions.
    """
    mask = (data == target_label).astype(np.uint8)
    mask = dilate_one_img(torch.from_numpy(mask).to(device), filter_size=3, pad_value=0.0)
    mask = dilate_one_img(mask, filter_size=3, pad_value=0.0)
    roi_oragn_mask = dilate_one_img(mask, filter_size=3, pad_value=0.0).cpu().numpy()
    return (roi_oragn_mask * remove_mask).astype(np.bool_)


def get_body_region_index_from_mask(input_mask):
    """
    Determine the top and bottom body region indices from an input mask.

    Args:
        input_mask (Tensor): Input mask tensor containing body region labels.

    Returns:
        tuple: Two lists representing the top and bottom region indices.
    """
    region_indices = {}
    # head and neck
    region_indices["region_0"] = [22, 120]
    # thorax
    region_indices["region_1"] = [28, 29, 30, 31, 32]
    # abdomen
    region_indices["region_2"] = [1, 2, 3, 4, 5, 14]
    # pelvis and lower
    region_indices["region_3"] = [93, 94]

    nda = input_mask.cpu().numpy().squeeze()
    unique_elements = np.lib.arraysetops.unique(nda)
    unique_elements = list(unique_elements)
    # print(f"nda: {nda.shape} {unique_elements}.")
    overlap_array = np.zeros(len(region_indices), dtype=np.uint8)
    for _j in range(len(region_indices)):
        overlap = any(element in region_indices[f"region_{_j}"] for element in unique_elements)
        overlap_array[_j] = np.uint8(overlap)
    overlap_array_indices = np.nonzero(overlap_array)[0]
    top_region_index = np.eye(len(region_indices), dtype=np.uint8)[np.amin(overlap_array_indices), ...]
    top_region_index = list(top_region_index)
    top_region_index = [int(_k) for _k in top_region_index]
    bottom_region_index = np.eye(len(region_indices), dtype=np.uint8)[np.amax(overlap_array_indices), ...]
    bottom_region_index = list(bottom_region_index)
    bottom_region_index = [int(_k) for _k in bottom_region_index]
    # print(f"{top_region_index} {bottom_region_index}")
    return top_region_index, bottom_region_index


def general_mask_generation_post_process(volume_t, target_tumor_label=None, device="cuda:0"):
    """
    Perform post-processing on a generated mask volume.

    This function applies various refinement steps to improve the quality of the generated mask,
    including body mask refinement, tumor prediction refinement, and organ-specific processing.

    Args:
        volume_t (ndarray): Input volume containing organ and tumor labels.
        target_tumor_label (int, optional): Label of the target tumor. Defaults to None.
        device (str, optional): Device to perform operations on. Defaults to "cuda:0".

    Returns:
        ndarray: Post-processed volume with refined organ and tumor labels.
    """
    # assume volume_t is np array with shape (H,W,D)
    hepatic_vessel = volume_t == 25
    airway = volume_t == 132

    # ------------ refine body mask pred
    body_region_mask = (
        erode_one_img(torch.from_numpy((volume_t > 0)).to(device), filter_size=3, pad_value=0.0).cpu().numpy()
    )
    body_region_mask, _ = supress_non_largest_components(body_region_mask, [1])
    body_region_mask = (
        dilate_one_img(torch.from_numpy(body_region_mask).to(device), filter_size=3, pad_value=0.0)
        .cpu()
        .numpy()
        .astype(np.uint8)
    )
    volume_t = volume_t * body_region_mask

    # ------------ refine tumor pred
    tumor_organ_dict = {23: 28, 24: 4, 26: 1, 27: 62, 128: 200}
    for t in [23, 24, 26, 27, 128]:
        if t != target_tumor_label:
            volume_t[volume_t == t] = tumor_organ_dict[t]
        else:
            volume_t[organ_fill_by_closing(volume_t, target_label=t, device=device)] = t
            volume_t[organ_fill_by_closing(volume_t, target_label=t, device=device)] = t
    # we only keep the largest connected componet for tumors except hepatic tumor and bone lesion
    if target_tumor_label != 26 and target_tumor_label != 128:
        volume_t, _ = supress_non_largest_components(volume_t, [target_tumor_label], default_val=200)
    target_tumor = volume_t == target_tumor_label

    # ------------ remove undesired organ pred
    # general post-process non-largest components suppression
    # process 4 ROI organs + spleen + 2 kidney + 5 lung lobes + duodenum + inferior vena cava
    oran_list = [1, 4, 10, 12, 3, 28, 29, 30, 31, 32, 5, 14, 13, 6, 7, 8, 9, 10]
    if target_tumor_label != 128:
        oran_list += list(range(33, 60))  # + list(range(63,87))
    data, _ = supress_non_largest_components(volume_t, oran_list, default_val=200)  # 200 is body region
    organ_remove_mask = (volume_t - data).astype(np.bool_)
    # process intestinal system (stomach 12, duodenum 13, small bowel 19, colon 62)
    intestinal_mask_ = (
        (data == 12).astype(np.uint8)
        + (data == 13).astype(np.uint8)
        + (data == 19).astype(np.uint8)
        + (data == 62).astype(np.uint8)
    )
    intestinal_mask, _ = supress_non_largest_components(intestinal_mask_, [1], default_val=0)
    # process small bowel 19
    small_bowel_remove_mask = (data == 19).astype(np.uint8) - (data == 19).astype(np.uint8) * intestinal_mask
    # process colon 62
    colon_remove_mask = (data == 62).astype(np.uint8) - (data == 62).astype(np.uint8) * intestinal_mask
    intestinal_remove_mask = (small_bowel_remove_mask + colon_remove_mask).astype(np.bool_)
    data[intestinal_remove_mask] = 200

    # ------------ full correponding organ in removed regions
    for organ_label in oran_list:
        data[organ_fill_by_closing(data, target_label=organ_label, device=device)] = organ_label

    if target_tumor_label == 23 and np.sum(target_tumor) > 0:
        # speical process for cases with lung tumor
        dia_lung_tumor_mask = (
            dilate_one_img(torch.from_numpy((data == 23)).to(device), filter_size=3, pad_value=0.0).cpu().numpy()
        )
        tmp = (
            (data * (dia_lung_tumor_mask.astype(np.uint8) - (data == 23).astype(np.uint8))).astype(np.float32).flatten()
        )
        tmp[tmp == 0] = float("nan")
        mode = int(stats.mode(tmp.flatten(), nan_policy="omit")[0])
        if mode in [28, 29, 30, 31, 32]:
            dia_lung_tumor_mask = (
                dilate_one_img(torch.from_numpy(dia_lung_tumor_mask).to(device), filter_size=3, pad_value=0.0)
                .cpu()
                .numpy()
            )
            lung_remove_mask = dia_lung_tumor_mask.astype(np.uint8) - (data == 23).astype(np.uint8).astype(np.uint8)
            data[organ_fill_by_removed_mask(data, target_label=mode, remove_mask=lung_remove_mask, device=device)] = (
                mode
            )
        dia_lung_tumor_mask = (
            dilate_one_img(torch.from_numpy(dia_lung_tumor_mask).to(device), filter_size=3, pad_value=0.0).cpu().numpy()
        )
        data[
            organ_fill_by_removed_mask(
                data, target_label=23, remove_mask=dia_lung_tumor_mask * organ_remove_mask, device=device
            )
        ] = 23
        for organ_label in [28, 29, 30, 31, 32]:
            data[organ_fill_by_closing(data, target_label=organ_label, device=device)] = organ_label
            data[organ_fill_by_closing(data, target_label=organ_label, device=device)] = organ_label
            data[organ_fill_by_closing(data, target_label=organ_label, device=device)] = organ_label

    if target_tumor_label == 26 and np.sum(target_tumor) > 0:
        # speical process for cases with hepatic tumor
        # process liver 1
        data[organ_fill_by_removed_mask(data, target_label=1, remove_mask=intestinal_remove_mask, device=device)] = 1
        data[organ_fill_by_removed_mask(data, target_label=1, remove_mask=intestinal_remove_mask, device=device)] = 1
        # process spleen 2
        data[organ_fill_by_removed_mask(data, target_label=3, remove_mask=organ_remove_mask, device=device)] = 3
        data[organ_fill_by_removed_mask(data, target_label=3, remove_mask=organ_remove_mask, device=device)] = 3
        dia_tumor_mask = (
            dilate_one_img(torch.from_numpy((data == target_tumor_label)).to(device), filter_size=3, pad_value=0.0)
            .cpu()
            .numpy()
        )
        dia_tumor_mask = (
            dilate_one_img(torch.from_numpy(dia_tumor_mask).to(device), filter_size=3, pad_value=0.0).cpu().numpy()
        )
        data[
            organ_fill_by_removed_mask(
                data, target_label=target_tumor_label, remove_mask=dia_tumor_mask * organ_remove_mask, device=device
            )
        ] = target_tumor_label
        # refine hepatic tumor
        hepatic_tumor_vessel_liver_mask_ = (
            (data == 26).astype(np.uint8) + (data == 25).astype(np.uint8) + (data == 1).astype(np.uint8)
        )
        hepatic_tumor_vessel_liver_mask_ = (hepatic_tumor_vessel_liver_mask_ > 1).astype(np.uint8)
        hepatic_tumor_vessel_liver_mask, _ = supress_non_largest_components(
            hepatic_tumor_vessel_liver_mask_, [1], default_val=0
        )
        removed_region = (hepatic_tumor_vessel_liver_mask_ - hepatic_tumor_vessel_liver_mask).astype(np.bool_)
        data[removed_region] = 200
        target_tumor = (target_tumor * hepatic_tumor_vessel_liver_mask).astype(np.bool_)
        # refine liver
        data[organ_fill_by_closing(data, target_label=1, device=device)] = 1
        data[organ_fill_by_closing(data, target_label=1, device=device)] = 1
        data[organ_fill_by_closing(data, target_label=1, device=device)] = 1

    if target_tumor_label == 27 and np.sum(target_tumor) > 0:
        # speical process for cases with colon tumor
        dia_tumor_mask = (
            dilate_one_img(torch.from_numpy((data == target_tumor_label)).to(device), filter_size=3, pad_value=0.0)
            .cpu()
            .numpy()
        )
        dia_tumor_mask = (
            dilate_one_img(torch.from_numpy(dia_tumor_mask).to(device), filter_size=3, pad_value=0.0).cpu().numpy()
        )
        data[
            organ_fill_by_removed_mask(
                data, target_label=target_tumor_label, remove_mask=dia_tumor_mask * organ_remove_mask, device=device
            )
        ] = target_tumor_label

    if target_tumor_label == 129 and np.sum(target_tumor) > 0:
        # speical process for cases with kidney tumor
        for organ_label in [5, 14]:
            data[organ_fill_by_closing(data, target_label=organ_label, device=device)] = organ_label
            data[organ_fill_by_closing(data, target_label=organ_label, device=device)] = organ_label
            data[organ_fill_by_closing(data, target_label=organ_label, device=device)] = organ_label
    # TODO: current model does not support hepatic vessel by size control.
    # we treat it as liver for better visiaulization
    print(
        "Current model does not support hepatic vessel by size control, "
        "so we treat generated hepatic vessel as part of liver for better visiaulization."
    )
    data[hepatic_vessel] = 1
    data[airway] = 132
    if target_tumor_label is not None:
        data[target_tumor] = target_tumor_label

    return data


class MapLabelValue:
    """
    Utility to map label values to another set of values.
    For example, map [3, 2, 1] to [0, 1, 2], [1, 2, 3] -> [0.5, 1.5, 2.5], ["label3", "label2", "label1"] -> [0, 1, 2],
    [3.5, 2.5, 1.5] -> ["label0", "label1", "label2"], etc.
    The label data must be numpy array or array-like data and the output data will be numpy array.

    """

    backend = [TransformBackends.NUMPY, TransformBackends.TORCH]

    def __init__(self, orig_labels: Sequence, target_labels: Sequence, dtype: DtypeLike = np.float32) -> None:
        """
        Args:
            orig_labels: original labels that map to others.
            target_labels: expected label values, 1: 1 map to the `orig_labels`.
            dtype: convert the output data to dtype, default to float32.
                if dtype is from PyTorch, the transform will use the pytorch backend, else with numpy backend.

        """
        if len(orig_labels) != len(target_labels):
            raise ValueError("orig_labels and target_labels must have the same length.")

        self.orig_labels = orig_labels
        self.target_labels = target_labels
        self.pair = tuple((o, t) for o, t in zip(self.orig_labels, self.target_labels) if o != t)
        type_dtype = type(dtype)
        if getattr(type_dtype, "__module__", "") == "torch":
            self.use_numpy = False
            self.dtype = get_equivalent_dtype(dtype, data_type=torch.Tensor)
        else:
            self.use_numpy = True
            self.dtype = get_equivalent_dtype(dtype, data_type=np.ndarray)

    def __call__(self, img: NdarrayOrTensor):
        """
        Apply the label mapping to the input image.

        Args:
            img (NdarrayOrTensor): Input image to be remapped.

        Returns:
            NdarrayOrTensor: Remapped image.
        """
        if self.use_numpy:
            img_np, *_ = convert_data_type(img, np.ndarray)
            _out_shape = img_np.shape
            img_flat = img_np.flatten()
            try:
                out_flat = img_flat.astype(self.dtype)
            except ValueError:
                # can't copy unchanged labels as the expected dtype is not supported, must map all the label values
                out_flat = np.zeros(shape=img_flat.shape, dtype=self.dtype)
            for o, t in self.pair:
                out_flat[img_flat == o] = t
            out_t = out_flat.reshape(_out_shape)
        else:
            img_t, *_ = convert_data_type(img, torch.Tensor)
            out_t = img_t.detach().clone().to(self.dtype)  # type: ignore
            for o, t in self.pair:
                out_t[img_t == o] = t
        out, *_ = convert_to_dst_type(src=out_t, dst=img, dtype=self.dtype)
        return out


def dynamic_infer(inferer, model, images):
    """
    Perform dynamic inference using a model and an inferer, typically a monai SlidingWindowInferer.

    This function determines whether to use the model directly or to use the provided inferer
    (such as a sliding window inferer) based on the size of the input images.

    Args:
        inferer: An inference object, typically a monai SlidingWindowInferer, which handles patch-based inference.
        model (torch.nn.Module): The model used for inference.
        images (torch.Tensor): The input images for inference, shape [N,C,H,W,D] or [N,C,H,W].

    Returns:
        torch.Tensor: The output from the model or the inferer, depending on the input size.
    """
    if torch.numel(images[0:1, 0:1, ...]) < math.prod(inferer.roi_size):
        return model(images)
    else:
        # Extract the spatial dimensions from the images tensor (H, W, D)
        spatial_dims = images.shape[2:]
        orig_roi = inferer.roi_size

        # Check that roi has the same number of dimensions as spatial_dims
        if len(orig_roi) != len(spatial_dims):
            raise ValueError(f"ROI length ({len(orig_roi)}) does not match spatial dimensions ({len(spatial_dims)}).")

        # Iterate and adjust each ROI dimension
        adjusted_roi = [min(roi_dim, img_dim) for roi_dim, img_dim in zip(orig_roi, spatial_dims)]
        inferer.roi_size = adjusted_roi
        output = inferer(network=model, inputs=images)
        inferer.roi_size = orig_roi
        return output