Comp2Comp-main/comp2comp/spine/spine.py

"""
@author: louisblankemeier
"""

import math
import os
import shutil
import zipfile
from pathlib import Path
from time import time
from typing import Union

import nibabel as nib
import numpy as np
import pandas as pd
import wget
from PIL import Image
from totalsegmentatorv2.python_api import totalsegmentator

from comp2comp.inference_class_base import InferenceClass
from comp2comp.io import io_utils
from comp2comp.models.models import Models
from comp2comp.spine import spine_utils
from comp2comp.visualization.dicom import to_dicom

# from totalsegmentator.libs import (
#     download_pretrained_weights,
#     nostdout,
#     setup_nnunet,
# )




class SpineSegmentation(InferenceClass):
    """Spine segmentation."""

    def __init__(self, model_name, save=True):
        super().__init__()
        self.model_name = model_name
        self.save_segmentations = save

    def __call__(self, inference_pipeline):
        # inference_pipeline.dicom_series_path = self.input_path
        self.output_dir = inference_pipeline.output_dir
        self.output_dir_segmentations = os.path.join(self.output_dir, "segmentations/")
        if not os.path.exists(self.output_dir_segmentations):
            os.makedirs(self.output_dir_segmentations)

        self.model_dir = inference_pipeline.model_dir

        # seg, mv = self.spine_seg(
        #     os.path.join(self.output_dir_segmentations, "converted_dcm.nii.gz"),
        #     self.output_dir_segmentations + "spine.nii.gz",
        #     inference_pipeline.model_dir,
        # )
        os.environ["TOTALSEG_WEIGHTS_PATH"] = self.model_dir

        seg = totalsegmentator(
            input=os.path.join(self.output_dir_segmentations, "converted_dcm.nii.gz"),
            output=os.path.join(self.output_dir_segmentations, "segmentation.nii"),
            task_ids=[292],
            ml=True,
            nr_thr_resamp=1,
            nr_thr_saving=6,
            fast=False,
            nora_tag="None",
            preview=False,
            task="total",
            # roi_subset=[
            #     "vertebrae_T12",
            #     "vertebrae_L1",
            #     "vertebrae_L2",
            #     "vertebrae_L3",
            #     "vertebrae_L4",
            #     "vertebrae_L5",
            # ],
            roi_subset=None,
            statistics=False,
            radiomics=False,
            crop_path=None,
            body_seg=False,
            force_split=False,
            output_type="nifti",
            quiet=False,
            verbose=False,
            test=0,
            skip_saving=True,
            device="gpu",
            license_number=None,
            statistics_exclude_masks_at_border=True,
            no_derived_masks=False,
            v1_order=False,
        )
        mv = nib.load(
            os.path.join(self.output_dir_segmentations, "converted_dcm.nii.gz")
        )

        # inference_pipeline.segmentation = nib.load(
        #     os.path.join(self.output_dir_segmentations, "segmentation.nii")
        # )
        inference_pipeline.segmentation = seg
        inference_pipeline.medical_volume = mv
        inference_pipeline.save_segmentations = self.save_segmentations
        return {}

    def setup_nnunet_c2c(self, model_dir: Union[str, Path]):
        """Adapted from TotalSegmentator."""

        model_dir = Path(model_dir)
        config_dir = model_dir / Path("." + self.model_name)
        (config_dir / "nnunet/results/nnUNet/3d_fullres").mkdir(
            exist_ok=True, parents=True
        )
        (config_dir / "nnunet/results/nnUNet/2d").mkdir(exist_ok=True, parents=True)
        weights_dir = config_dir / "nnunet/results"
        self.weights_dir = weights_dir

        os.environ["nnUNet_raw_data_base"] = str(
            weights_dir
        )  # not needed, just needs to be an existing directory
        os.environ["nnUNet_preprocessed"] = str(
            weights_dir
        )  # not needed, just needs to be an existing directory
        os.environ["RESULTS_FOLDER"] = str(weights_dir)

    def download_spine_model(self, model_dir: Union[str, Path]):
        download_dir = Path(
            os.path.join(
                self.weights_dir,
                "nnUNet/3d_fullres/Task252_Spine/nnUNetTrainerV2_ep4000_nomirror__nnUNetPlansv2.1",
            )
        )
        fold_0_path = download_dir / "fold_0"
        if not os.path.exists(fold_0_path):
            download_dir.mkdir(parents=True, exist_ok=True)
            wget.download(
                "https://huggingface.co/louisblankemeier/spine_v1/resolve/main/fold_0.zip",
                out=os.path.join(download_dir, "fold_0.zip"),
            )
            with zipfile.ZipFile(
                os.path.join(download_dir, "fold_0.zip"), "r"
            ) as zip_ref:
                zip_ref.extractall(download_dir)
            os.remove(os.path.join(download_dir, "fold_0.zip"))
            wget.download(
                "https://huggingface.co/louisblankemeier/spine_v1/resolve/main/plans.pkl",
                out=os.path.join(download_dir, "plans.pkl"),
            )
            print("Spine model downloaded.")
        else:
            print("Spine model already downloaded.")

    def spine_seg(
        self, input_path: Union[str, Path], output_path: Union[str, Path], model_dir
    ):
        """Run spine segmentation.

        Args:
            input_path (Union[str, Path]): Input path.
            output_path (Union[str, Path]): Output path.
        """

        print("Segmenting spine...")
        st = time()
        os.environ["SCRATCH"] = self.model_dir
        os.environ["TOTALSEG_WEIGHTS_PATH"] = self.model_dir

        # Setup nnunet
        model = "3d_fullres"
        folds = [0]
        trainer = "nnUNetTrainerV2_ep4000_nomirror"
        crop_path = None
        task_id = [252]

        if self.model_name == "ts_spine":
            setup_nnunet()
            download_pretrained_weights(task_id[0])
        elif self.model_name == "stanford_spine_v0.0.1":
            self.setup_nnunet_c2c(model_dir)
            self.download_spine_model(model_dir)
        else:
            raise ValueError("Invalid model name.")

        if not self.save_segmentations:
            output_path = None

        from totalsegmentator.nnunet import nnUNet_predict_image

        with nostdout():
            img, seg = nnUNet_predict_image(
                input_path,
                output_path,
                task_id,
                model=model,
                folds=folds,
                trainer=trainer,
                tta=False,
                multilabel_image=True,
                resample=1.5,
                crop=None,
                crop_path=crop_path,
                task_name="total",
                nora_tag="None",
                preview=False,
                nr_threads_resampling=1,
                nr_threads_saving=6,
                quiet=False,
                verbose=False,
                test=0,
            )
        end = time()

        # Log total time for spine segmentation
        print(f"Total time for spine segmentation: {end-st:.2f}s.")

        if self.model_name == "stanford_spine_v0.0.1":
            seg_data = seg.get_fdata()
            # subtract 17 from seg values except for 0
            seg_data = np.where(seg_data == 0, 0, seg_data - 17)
            seg = nib.Nifti1Image(seg_data, seg.affine, seg.header)

        return seg, img


class AxialCropper(InferenceClass):
    """Crop the CT image (medical_volume) and segmentation based on user-specified
    lower and upper levels of the spine.
    """

    def __init__(self, lower_level: str = "L5", upper_level: str = "L1", save=True):
        """
        Args:
            lower_level (str, optional): Lower level of the spine. Defaults to "L5".
            upper_level (str, optional): Upper level of the spine. Defaults to "L1".
            save (bool, optional): Save cropped image and segmentation. Defaults to True.

        Raises:
            ValueError: If lower_level or upper_level is not a valid spine level.
        """
        super().__init__()
        self.lower_level = lower_level
        self.upper_level = upper_level
        ts_spine_full_model = Models.model_from_name("ts_spine_full")
        categories = ts_spine_full_model.categories
        try:
            self.lower_level_index = categories[self.lower_level]
            self.upper_level_index = categories[self.upper_level]
        except KeyError:
            raise ValueError("Invalid spine level.") from None
        self.save = save

    def __call__(self, inference_pipeline):
        """
        First dim goes from L to R.
        Second dim goes from P to A.
        Third dim goes from I to S.
        """
        segmentation = inference_pipeline.segmentation
        segmentation_data = segmentation.get_fdata()
        upper_level_index = np.where(segmentation_data == self.upper_level_index)[
            2
        ].max()
        lower_level_index = np.where(segmentation_data == self.lower_level_index)[
            2
        ].min()
        segmentation = segmentation.slicer[:, :, lower_level_index:upper_level_index]
        inference_pipeline.segmentation = segmentation

        medical_volume = inference_pipeline.medical_volume
        medical_volume = medical_volume.slicer[
            :, :, lower_level_index:upper_level_index
        ]
        inference_pipeline.medical_volume = medical_volume

        if self.save:
            nib.save(
                segmentation,
                os.path.join(
                    inference_pipeline.output_dir, "segmentations", "spine.nii.gz"
                ),
            )
            nib.save(
                medical_volume,
                os.path.join(
                    inference_pipeline.output_dir,
                    "segmentations",
                    "converted_dcm.nii.gz",
                ),
            )
        return {}


class SpineComputeROIs(InferenceClass):
    def __init__(self, spine_model):
        super().__init__()
        self.spine_model_name = spine_model
        self.spine_model_type = Models.model_from_name(self.spine_model_name)

    def __call__(self, inference_pipeline):
        # Compute ROIs
        inference_pipeline.spine_model_type = self.spine_model_type

        (spine_hus, rois, segmentation_hus, centroids_3d) = spine_utils.compute_rois(
            inference_pipeline.segmentation,
            inference_pipeline.medical_volume,
            self.spine_model_type,
        )

        inference_pipeline.spine_hus = spine_hus
        inference_pipeline.segmentation_hus = segmentation_hus
        inference_pipeline.rois = rois
        inference_pipeline.centroids_3d = centroids_3d

        return {}


class SpineMetricsSaver(InferenceClass):
    """Save metrics to a CSV file."""

    def __init__(self):
        super().__init__()

    def __call__(self, inference_pipeline):
        """Save metrics to a CSV file."""
        self.spine_hus = inference_pipeline.spine_hus
        self.seg_hus = inference_pipeline.segmentation_hus
        self.output_dir = inference_pipeline.output_dir
        self.csv_output_dir = os.path.join(self.output_dir, "metrics")
        if not os.path.exists(self.csv_output_dir):
            os.makedirs(self.csv_output_dir, exist_ok=True)
        self.save_results()
        if hasattr(inference_pipeline, "dicom_ds"):
            if not os.path.exists(os.path.join(self.output_dir, "dicom_metadata.csv")):
                io_utils.write_dicom_metadata_to_csv(
                    inference_pipeline.dicom_ds,
                    os.path.join(self.output_dir, "dicom_metadata.csv"),
                )

        return {}

    def save_results(self):
        """Save results to a CSV file."""
        df = pd.DataFrame(columns=["Level", "ROI HU", "Seg HU"])
        for i, level in enumerate(self.spine_hus):
            hu = self.spine_hus[level]
            seg_hu = self.seg_hus[level]
            row = [level, hu, seg_hu]
            df.loc[i] = row
        df = df.iloc[::-1]
        df.to_csv(os.path.join(self.csv_output_dir, "spine_metrics.csv"), index=False)


class SpineFindDicoms(InferenceClass):
    def __init__(self):
        super().__init__()

    def __call__(self, inference_pipeline):
        inferior_superior_centers = spine_utils.find_spine_dicoms(
            inference_pipeline.centroids_3d,
        )

        spine_utils.save_nifti_select_slices(
            inference_pipeline.output_dir, inferior_superior_centers
        )
        inference_pipeline.dicom_file_paths = [
            str(center) for center in inferior_superior_centers
        ]
        inference_pipeline.names = list(inference_pipeline.rois.keys())
        inference_pipeline.dicom_file_names = list(inference_pipeline.rois.keys())
        inference_pipeline.inferior_superior_centers = inferior_superior_centers

        return {}


class SpineCoronalSagittalVisualizer(InferenceClass):
    def __init__(self, format="png"):
        super().__init__()
        self.format = format

    def __call__(self, inference_pipeline):
        output_path = inference_pipeline.output_dir
        spine_model_type = inference_pipeline.spine_model_type

        img_sagittal, img_coronal = spine_utils.visualize_coronal_sagittal_spine(
            inference_pipeline.segmentation.get_fdata(),
            list(inference_pipeline.rois.values()),
            inference_pipeline.medical_volume.get_fdata(),
            list(inference_pipeline.centroids_3d.values()),
            output_path,
            spine_hus=inference_pipeline.spine_hus,
            seg_hus=inference_pipeline.segmentation_hus,
            model_type=spine_model_type,
            pixel_spacing=inference_pipeline.pixel_spacing_list,
            format=self.format,
        )
        inference_pipeline.spine_vis_sagittal = img_sagittal
        inference_pipeline.spine_vis_coronal = img_coronal
        inference_pipeline.spine = True
        if not inference_pipeline.save_segmentations:
            shutil.rmtree(os.path.join(output_path, "segmentations"))
        return {}


class SpineReport(InferenceClass):
    def __init__(self, format="png"):
        super().__init__()
        self.format = format

    def __call__(self, inference_pipeline):
        sagittal_image = inference_pipeline.spine_vis_sagittal
        coronal_image = inference_pipeline.spine_vis_coronal
        # concatenate these numpy arrays laterally
        img = np.concatenate((coronal_image, sagittal_image), axis=1)
        output_path = os.path.join(
            inference_pipeline.output_dir, "images", "spine_report"
        )
        if self.format == "png":
            im = Image.fromarray(img)
            im.save(output_path + ".png")
        elif self.format == "dcm":
            to_dicom(img, output_path + ".dcm")
        return {}


class SpineMuscleAdiposeTissueReport(InferenceClass):
    """Spine muscle adipose tissue report class."""

    def __init__(self):
        super().__init__()
        self.image_files = [
            "spine_coronal.png",
            "spine_sagittal.png",
            "T12.png",
            "L1.png",
            "L2.png",
            "L3.png",
            "L4.png",
            "L5.png",
        ]

    def __call__(self, inference_pipeline):
        image_dir = Path(inference_pipeline.output_dir) / "images"
        self.generate_panel(image_dir)
        return {}

    def generate_panel(self, image_dir: Union[str, Path]):
        """Generate panel.
        Args:
            image_dir (Union[str, Path]): Path to the image directory.
        """
        image_files = [os.path.join(image_dir, path) for path in self.image_files]
        # construct a list which includes only the images that exist
        image_files = [path for path in image_files if os.path.exists(path)]

        im_cor = Image.open(image_files[0])
        im_sag = Image.open(image_files[1])
        im_cor_width = int(im_cor.width / im_cor.height * 512)
        num_muscle_fat_cols = math.ceil((len(image_files) - 2) / 2)
        width = (8 + im_cor_width + 8) + ((512 + 8) * num_muscle_fat_cols)
        height = 1048
        new_im = Image.new("RGB", (width, height))

        index = 2
        for j in range(8, height, 520):
            for i in range(8 + im_cor_width + 8, width, 520):
                try:
                    im = Image.open(image_files[index])
                    im.thumbnail((512, 512))
                    new_im.paste(im, (i, j))
                    index += 1
                    im.close()
                except Exception:
                    continue

        im_cor.thumbnail((im_cor_width, 512))
        new_im.paste(im_cor, (8, 8))
        im_sag.thumbnail((im_cor_width, 512))
        new_im.paste(im_sag, (8, 528))
        new_im.save(os.path.join(image_dir, "spine_muscle_adipose_tissue_report.png"))
        im_cor.close()
        im_sag.close()
        new_im.close()