vista3d_pipeline.py

import copy
import json
import logging
import os
import pathlib
from typing import Sequence

import numpy as np
import torch
from monai.apps.vista3d.transforms import VistaPostTransformd, VistaPreTransformd
from monai.data.utils import decollate_batch, list_data_collate
from monai.networks.utils import eval_mode, train_mode
from monai.transforms import (
    CastToTyped,
    Compose,
    CropForegroundd,
    EnsureChannelFirstd,
    EnsureTyped,
    Invertd,
    Lambdad,
    LoadImaged,
    Orientationd,
    SaveImaged,
    ScaleIntensityRanged,
    Spacingd,
    reset_ops_id,
)
from monai.utils import ForwardMode, optional_import, set_determinism
from monai.utils.enums import CommonKeys as Keys
from monai.utils.module import look_up_option
from scripts.inferer import Vista3dInferer
from transformers import AutoModel, Pipeline
from transformers.pipelines import PIPELINE_REGISTRY

rearrange, _ = optional_import("einops", name="rearrange")

FILE_PATH = os.path.dirname(__file__)


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)


class VISTA3DPipeline(Pipeline):
    """Define the VISTA3D pipeline."""

    PREPROCESSING_EXTRA_ARGS = [
        "image_key",
        "resample_spacing",
        "metadata_path",
    ]
    INFERENCE_EXTRA_ARGS = [
        "mode",
        "amp",
        "hyper_kwargs",
        "roi_size",
        "overlap",
        "sw_batch_size",
        "use_point_window",
    ]
    POSTPROCESSING_EXTRA_ARGS = [
        "pred_key",
        "image_key",
        "output_dir",
        "output_ext",
        "output_postfix",
        "separate_folder",
        "save_output",
    ]
    EVERYTHING_LABEL = list(
        set([i + 1 for i in range(132)])
        - set([2, 16, 18, 20, 21, 23, 24, 25, 26, 27, 128, 129, 130, 131, 132])
    )

    def __init__(self, model, **kwargs):
        super().__init__(model, **kwargs)
        self.preprocessing_transforms = self._init_preprocessing_transforms(
            **self._preprocess_params
        )
        self.inferer = self._init_inferer(**self._forward_params)
        self.postprocessing_transforms = self._init_postprocessing_transforms(
            **self._postprocess_params
        )

    def _init_inferer(
        self,
        roi_size: Sequence = (128, 128, 128),
        overlap: float = 0.3,
        sw_batch_size: int = 1,
        use_point_window: bool = True,
    ):
        return Vista3dInferer(
            roi_size=roi_size,
            overlap=overlap,
            use_point_window=use_point_window,
            sw_batch_size=sw_batch_size,
        )

    def _init_preprocessing_transforms(
        self,
        image_key: str = "image",
        resample_spacing: Sequence = (1.5, 1.5, 1.5),
        metadata_path: str = os.path.join(FILE_PATH, "metadata.json"),
    ):
        device = self.device
        subclass = {
            "2": [14, 5],
            "20": [28, 29, 30, 31, 32],
            "21": list(range(33, 57)) + list(range(63, 98)) + [114, 120, 122],
        }
        metadata = json.loads(pathlib.Path(metadata_path).read_text())
        labels_dict = metadata["network_data_format"]["outputs"]["pred"]["channel_def"]
        preprocessing_transforms = Compose(
            [
                LoadImaged(keys=image_key, image_only=True),
                EnsureChannelFirstd(keys=image_key),
                EnsureTyped(keys=image_key, device=device, track_meta=True),
                Spacingd(keys=image_key, pixdim=resample_spacing, mode="bilinear"),
                CropForegroundd(
                    keys=image_key, allow_smaller=True, margin=10, source_key=image_key
                ),
                VistaPreTransformd(
                    keys=image_key, subclass=subclass, labels_dict=labels_dict
                ),
                ScaleIntensityRanged(
                    keys=image_key,
                    a_min=-963.8247715525971,
                    a_max=1053.678477684517,
                    b_min=0,
                    b_max=1,
                    clip=True,
                ),
                Orientationd(keys=image_key, axcodes="RAS"),
                CastToTyped(keys=image_key, dtype=torch.float32),
            ]
        )
        return preprocessing_transforms

    def _init_postprocessing_transforms(
        self,
        pred_key: str = "pred",
        image_key: str = "image",
        output_dir: str = "output_directory",
        output_ext: str = ".nii.gz",
        output_dtype: torch.dtype = torch.float32,
        output_postfix: str = "seg",
        separate_folder: bool = True,
        save_output: bool = True,
    ):
        transforms = [
            VistaPostTransformd(keys=pred_key),
            Invertd(
                keys=pred_key,
                transform=copy.deepcopy(self.preprocessing_transforms),
                orig_keys=image_key,
                nearest_interp=True,
                to_tensor=True,
            ),
            Lambdad(keys=pred_key, func=lambda x: torch.nan_to_num(x, nan=255)),
        ]
        if save_output:
            transforms.append(
                SaveImaged(
                    keys=pred_key,
                    resample=False,
                    output_dir=output_dir,
                    output_ext=output_ext,
                    output_dtype=output_dtype,
                    output_postfix=output_postfix,
                    separate_folder=separate_folder,
                ),
            )
        postprocessing_transforms = Compose(transforms=transforms)
        return postprocessing_transforms

    def _sanitize_parameters(self, **kwargs):
        """
        _sanitize_parameters exists to allow users to pass any parameters whenever they wish,
        be it at initialization time pipeline(...., maybe_arg=4) or at call time pipe = pipeline(...); output = pipe(...., maybe_arg=4).
        The returns of _sanitize_parameters are the 3 dicts of kwargs that will be passed directly to preprocess, _forward and postprocess.
        Don't fill anything if the caller didn't call with any extra parameter. That allows to keep the default arguments in the function
        definition which is always more “natural”."""

        vista3d_preprocessing_kwargs = {}
        vista3d_infer_kwargs = {}
        vista3d_postprocessing_kwargs = {}
        for key in self.INFERENCE_EXTRA_ARGS:
            if key in kwargs:
                vista3d_infer_kwargs[key] = kwargs[key]

        for key in self.PREPROCESSING_EXTRA_ARGS:
            if key in kwargs:
                vista3d_preprocessing_kwargs[key] = kwargs[key]

        for key in self.POSTPROCESSING_EXTRA_ARGS:
            if key in kwargs:
                vista3d_postprocessing_kwargs[key] = kwargs[key]

        return (
            vista3d_preprocessing_kwargs,
            vista3d_infer_kwargs,
            vista3d_postprocessing_kwargs,
        )

    def check_prompts_format(self, label_prompt, points, point_labels):
        """check the format of user prompts
        label_prompt: [1,2,3,4,...,B] List of tensors
        points: [[[x,y,z], [x,y,z], ...]] List of coordinates of a single object
        point_labels: [[1,1,0,...]] List of scalar that matches number of points
        """
        # check prompt is given
        if label_prompt is None and points is None:
            everything_labels = self.hyper_kwargs.get("everything_labels", None)
            if everything_labels is not None:
                label_prompt = [torch.tensor(_) for _ in everything_labels]
                return label_prompt, points, point_labels
            else:
                raise ValueError("Prompt must be given for inference.")
        # check label_prompt
        if label_prompt is not None:
            if isinstance(label_prompt, list):
                if not np.all([len(_) == 1 for _ in label_prompt]):
                    raise ValueError(
                        "Label prompt must be a list of single scalar, [1,2,3,4,...,]."
                    )
                if isinstance(label_prompt[0], list):
                    for prompt in label_prompt:
                        if not np.all([(x < 255).item() for x in prompt]):
                            raise ValueError(
                                "Current bundle only supports label prompt smaller than 255."
                            )
                else:
                    if not np.all([(x < 255).item() for x in label_prompt]):
                        raise ValueError(
                            "Current bundle only supports label prompt smaller than 255."
                        )
                if points is None:
                    supported_list = list(
                        {i + 1 for i in range(132)} - {16, 18, 129, 130, 131}
                    )
                    if isinstance(label_prompt[0], list):
                        for prompt in label_prompt:
                            if not np.all([(x < 255).item() for x in prompt]):
                                raise ValueError(
                                    "Current bundle only supports label prompt smaller than 255."
                                )
                    else:
                        if not np.all([x in supported_list for x in label_prompt]):
                            raise ValueError(
                                "Undefined label prompt detected. Provide point prompts for zero-shot."
                            )
            else:
                raise ValueError("Label prompt must be a list, [1,2,3,4,...,].")
        # check points
        if points is not None:
            if point_labels is None:
                raise ValueError("Point labels must be given if points are given.")
            if not np.all([len(_) == 3 for _ in points]):
                raise ValueError(
                    "Points must be three dimensional (x,y,z) in the shape of [[x,y,z],...,[x,y,z]]."
                )
            if len(points) != len(point_labels):
                raise ValueError("Points must match point labels.")
            if not np.all([_ in [-1, 0, 1, 2, 3] for _ in point_labels]):
                raise ValueError(
                    "Point labels can only be -1,0,1 and 2,3 for special flags."
                )
        if label_prompt is not None and points is not None:
            if len(label_prompt) != 1:
                raise ValueError(
                    "Label prompt can only be a single object if provided with point prompts."
                )
        # check point_labels
        if point_labels is not None:
            if points is None:
                raise ValueError("Points must be given if point labels are given.")
        return label_prompt, points, point_labels

    def transform_points(self, point, affine):
        """transform point to the coordinates of the transformed image
        point: numpy array [bs, N, 3]
        """
        bs, n = point.shape[:2]
        point = np.concatenate((point, np.ones((bs, n, 1))), axis=-1)
        point = rearrange(point, "b n d -> d (b n)")
        point = affine @ point
        point = rearrange(point, "d (b n)-> b n d", b=bs)[:, :, :3]
        return point

    def preprocess(
        self,
        inputs,
        **kwargs,
    ):
        for key, value in kwargs.items():
            if key in self._preprocess_params and value != self._preprocess_params[key]:
                logging.warning(
                    f"Please set the parameter {key} during initialization."
                )

            if key not in self.PREPROCESSING_EXTRA_ARGS:
                logging.warning(f"Cannot set parameter {key} for preprocessing.")
        inputs = self.preprocessing_transforms(inputs)
        inputs = list_data_collate([inputs])
        return inputs

    def _forward(
        self,
        inputs,
        mode: str = ForwardMode.EVAL,
        amp: bool = True,
        hyper_kwargs: dict = {"user_prompt": 1, "everything_labels": 1},
    ):
        set_determinism(seed=123)

        if inputs is None:
            raise ValueError("Must provide input data for inference.")
        self.hyper_kwargs = hyper_kwargs

        label_set = hyper_kwargs.get("label_set", None)
        # this validation label set should be consistent with 'labels.unique()', used to generate fg/bg points
        val_label_set = hyper_kwargs.get("val_label_set", label_set)
        # If user provide prompts in the inference, input image must contain original affine.
        # the point coordinates are from the original_affine space, while image here is after preprocess transforms.
        if hyper_kwargs["user_prompt"]:
            inputs, label_prompt, points, point_labels = (
                inputs["image"],
                inputs.get("label_prompt", None),
                inputs.get("points", None),
                inputs.get("point_labels", None),
            )
            labels = None
            label_prompt, points, point_labels = self.check_prompts_format(
                label_prompt, points, point_labels
            )
            inputs = inputs.to(self.device)
            # For N foreground object, label_prompt is [1, N], but the batch number 1 needs to be removed. Convert to [N, 1]
            label_prompt = (
                torch.as_tensor([label_prompt]).to(inputs.device)[0].unsqueeze(-1)
                if label_prompt is not None
                else None
            )
            # For points, the size can only be [1, K, 3], where K is the number of points for this single foreground object.
            if points is not None:
                points = torch.as_tensor([points])
                points = self.transform_points(
                    points,
                    np.linalg.inv(inputs.affine[0])
                    @ inputs.meta["original_affine"][0].numpy(),
                )
                points = torch.from_numpy(points).to(inputs.device)
            point_labels = (
                torch.as_tensor([point_labels]).to(inputs.device)
                if point_labels is not None
                else None
            )

        # If validation with ground truth label available.
        else:
            # TODO add these as attribute.
            inputs, labels = inputs["image"], inputs["label"]
            # create label prompt, this should be consistent with the label prompt used for training.
            if label_set is None:
                output_classes = hyper_kwargs.get("output_classes", None)
                label_set = np.arange(output_classes).tolist()
            label_prompt = torch.tensor(label_set).to(self.device).unsqueeze(-1)
            # point prompt is generated withing vista3d, provide empty points
            points = torch.zeros(label_prompt.shape[0], 1, 3).to(inputs.device)
            point_labels = -1 + torch.zeros(label_prompt.shape[0], 1).to(inputs.device)
            # validation for either auto or point.
            if hyper_kwargs.get("val_head", "auto") == "auto":
                # automatic only validation
                # remove val_label_set, vista3d will not sample points from gt labels.
                val_label_set = None
            else:
                # point only validation
                label_prompt = None

        # put iteration outputs into outputs TODO need to align with the customized inputs
        outputs = {Keys.IMAGE: inputs, Keys.LABEL: labels}
        mode = look_up_option(mode, ForwardMode)
        if mode == ForwardMode.EVAL:
            mode = eval_mode
        elif mode == ForwardMode.TRAIN:
            mode = train_mode
        else:
            raise ValueError(f"unsupported mode: {mode}, should be 'eval' or 'train'.")

        # execute forward computation
        self.model.network.to(self.device)
        with mode(self.model):
            if amp:
                with torch.autocast("cuda"):
                    outputs[Keys.PRED] = self.inferer(
                        inputs=inputs,
                        network=self.model.network,
                        point_coords=points,
                        point_labels=point_labels,
                        class_vector=label_prompt,
                        labels=labels,
                        label_set=val_label_set,
                    )
            else:
                outputs[Keys.PRED] = self.inferer(
                    inputs=inputs,
                    network=self.model.network,
                    point_coords=points,
                    point_labels=point_labels,
                    class_vector=label_prompt,
                    labels=labels,
                    label_set=val_label_set,
                )
        inputs = reset_ops_id(inputs)
        # Add dim 0 for decollate batch
        outputs["label_prompt"] = (
            label_prompt.unsqueeze(0) if label_prompt is not None else None
        )
        outputs["points"] = points.unsqueeze(0) if points is not None else None
        outputs["point_labels"] = (
            point_labels.unsqueeze(0) if point_labels is not None else None
        )
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        return outputs

    def postprocess(self, outputs, **kwargs):
        for key, value in kwargs.items():
            if (
                key in self._postprocess_params
                and value != self._postprocess_params[key]
            ):
                logging.warning(
                    f"Please set the parameter {key} during initialization."
                )

            if key not in self.POSTPROCESSING_EXTRA_ARGS:
                logging.warning(f"Cannot set parameter {key} for postprocessing.")
        outputs = self.postprocessing_transforms(decollate_batch(outputs))
        return outputs


def register_simple_pipeline():
    PIPELINE_REGISTRY.register_pipeline(
        "vista3d",
        pipeline_class=VISTA3DPipeline,
        pt_model=AutoModel,
        default={"pt": (os.path.join(FILE_PATH, "vista3d_pretrained_model"), "")},
        type="image",  # current support type: text, audio, image, multimodal
    )