app.py

import streamlit as st
import math
import numpy as np
import nibabel as nib
import torch
import torch.nn.functional as F
from transformers import AutoModel
import os
import tempfile
from pathlib import Path
from skimage.filters import threshold_otsu
import torchio as tio
# import psutil

def infer_full_vol(tensor, model):
    tensor = tensor.unsqueeze(0).unsqueeze(0)  # Shape: [1, 1, D, H, W] - adding batch and channel dims
    
    tensor = torch.movedim(tensor, -1, -3)
    tensor = tensor / tensor.max()
    
    sizes = tensor.shape[-3:]
    new_sizes = [math.ceil(s / 16) * 16 for s in sizes]
    total_pads = [new_size - s for s, new_size in zip(sizes, new_sizes)]
    pad_before = [pad // 2 for pad in total_pads]
    pad_after = [pad - pad_before[i] for i, pad in enumerate(total_pads)]
    padding = []
    for i in reversed(range(len(pad_before))):
        padding.extend([pad_before[i], pad_after[i]])
    tensor = F.pad(tensor, padding)
    
    with torch.no_grad():
        output = model(tensor) 
        if type(output) is tuple or type(output) is list:
            output = output[0]
        output = torch.sigmoid(output)
        
    slices = [slice(None)] * output.dim()
    for i in range(len(pad_before)):
        dim = -3 + i
        start = pad_before[i]
        size = sizes[i]
        end = start + size
        slices[dim] = slice(start, end)
    output = output[tuple(slices)]
    
    output = torch.movedim(output, -3, -1).type(tensor.type())
    return output.squeeze().detach().cpu().numpy()

def infer_patch_based(tensor, model, patch_size=64, stride_length=32, stride_width=32, stride_depth=16, batch_size=10, num_worker=2):
    test_subject = tio.Subject(img = tio.ScalarImage(tensor=tensor.unsqueeze(0))) # adding channel dim while creating the TorchIO subject
                
    overlap = np.subtract(patch_size, (stride_length, stride_width, stride_depth))
    
    def normaliser(batch):
        """
        Purpose: Normalise pixel intensities of each patch using the max values in the 3D patch
        :param batch: 5D array (batch_size x channel x width x depth x height)
        """
        for i in range(batch.shape[0]):
            batch[i] = batch[i] / batch[i].max()
        return batch
    
    with torch.no_grad():
        grid_sampler = tio.inference.GridSampler(
                            test_subject,
                            patch_size,
                            overlap,
                        )
        aggregator = tio.inference.GridAggregator(grid_sampler, overlap_mode="average")
        patch_loader = torch.utils.data.DataLoader(grid_sampler, batch_size=batch_size, shuffle=False, num_workers=num_worker)
        total_batches = len(patch_loader)
        progress_bar = st.progress(0)
        for i, patches_batch in enumerate(patch_loader):
            st.text(f"Processing batch {i + 1} of {total_batches}...")
            
            local_batch = normaliser(patches_batch['img'][tio.DATA].float())
            local_batch = local_batch / local_batch.max()
            locations = patches_batch[tio.LOCATION]

            local_batch = torch.movedim(local_batch, -1, -3)
            
            output = model(local_batch)
            if type(output) is tuple or type(output) is list:
                output = output[0]
            output = torch.sigmoid(output).detach().cpu()

            output = torch.movedim(output, -3, -1).type(local_batch.type())
            aggregator.add_batch(output, locations)

            progress_bar.progress((i + 1) / total_batches)
            # st.text(f"CPU usage: {psutil.cpu_percent()}% | RAM usage: {psutil.virtual_memory().percent}%")

        predicted = aggregator.get_output_tensor().squeeze().numpy()
        
    return predicted

# Set page configuration
st.set_page_config(
    page_title="DS6 | Segmenting vessels in 3D MRA-ToF (ideally, 7T)",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Sidebar content
with st.sidebar:
    st.title("Segmenting vessels in the brain from a 3D Magnetic Resonance Angiograph | DS6")
    st.markdown("""
    This application allows you to upload a 3D NIfTI file (dims: H x W x D, where the final dim is the slice dim in the axial plane), process it through a pre-trained 3D model (from DS6 and other related works), and download the output as a `.nii.gz` file containing the vessel segmentation.
    
    **Instructions**:
    - Upload your 3D NIfTI file (`.nii` or `.nii.gz`). The model was trained on `7T MRA-ToF` data, but it should work on other field strengths as well.
    - Select a pretrained model from the dropdown menu.
    - Select the inference mode (full volume or patch-based) from the dropdown menu.
    - Click the "Process" button to generate the latent factors.
    """)
    st.markdown("---")
    st.markdown("© 2024 Soumick Chatterjee")

# Main content
st.header("DS6, Deformation-Aware Semi-Supervised Learning: Application to Small Vessel Segmentation with Noisy Training Data")
st.markdown("""
This application can be used to perform vessel segmentation by uploading a `Magnetic Resonance Angiograph (MRA-ToF)`, ideally acquired at 7T, but it should also work on other field strengths, in NIFTI format (.nii or .nii.gz).

The list of weights includes those from the original [DS6 paper](https://doi.org/10.3390/jimaging8100259), from the [SMILE-UHURA challenge](https://doi.org/10.7303/syn47164761), from a follow-up research [SPOCKMIP](https://arxiv.org/abs/2407.08655), as well as a [fine-tuned version](https://www.medrxiv.org/content/10.1101/2024.10.03.24314845v1) of the DS6 model (trained on the SMILE-UHURA dataset) using the Cambridge 7T Cerebral Small Vessel Disease (CamSVD) dataset, which contains data from subjects with lacunar strokes with SVD, non-lacunar strokes without SVD, and healthy controls.

Segmentation can be performed in 2 different inference modes: full volume inference and patch-based inference. All these research works used patch-based inference. However, if the volume is not large enough, and only large vessels are to be segmented, full volume inference may be performed, which is significantly faster. Nevertheless, full volume inference might result in an out-of-memory error (if the volume is very large) and may not segment the small vessels properly.  
""")
with st.expander("List of available pretrained models"):
    st.markdown(
        """
        | Model Name | Description |
        |------------|-------------|
        | `DS6_UNet3D_woDeform` | UNet 3D, trained without deformation-aware learning |
        | `DS6_UNetMSS3D_woDeform` | UNet MSS 3D, trained without deformation-aware learning |
        | `DS6_UNetMSS3D_wDeform` | UNet MSS 3D, trained with deformation-aware learning [Proposed method, DS6] |
        | `SMILEUHURA_DS6_UNet3D_woDeform` | UNet 3D, trained without deformation-aware learning on the SMILE-UHURA dataset |
        | `SMILEUHURA_DS6_UNetMSS3D_woDeform` | UNet MSS 3D, trained without deformation-aware learning on the SMILE-UHURA dataset |
        | `SMILEUHURA_DS6_UNetMSS3D_wDeform` | UNet MSS 3D, trained with deformation-aware learning on the SMILE-UHURA dataset |
        | `SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_MIP` | UNet MSS 3D, trained with MIP (Maximum Intensity Projection) loss [Proposed method, SPOCKMIP] |
        | `SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_mMIP` | UNet MSS 3D, trained with Multi-axis MIP loss [Proposed method, SPOCKMIP] |
        | `SMILEUHURA_SPOCKMIP_UNet3D_MIP` | UNet 3D, trained with MIP (Maximum Intensity Projection) loss [Proposed method, SPOCKMIP] |
        | `SMILEUHURA_SPOCKMIP_UNet3D_mMIP` | UNet 3D, trained with Multi-axis MIP loss [Proposed method, SPOCKMIP] |
        | `SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_DS6MIP` | UNet MSS 3D, trained with deformation-aware learning, and then with MIP loss |
        | `SMILEUHURA_DS6_CamSVD_UNetMSS3D_wDeform` | UNet MSS 3D, initially trained with deformation-aware learning on the SMILE-UHURA dataset, and then fine-tuned on the CamSVD dataset with deformation-aware learning |
        """
    )
st.markdown("---")
    
# File uploader
uploaded_file = st.file_uploader(
    "Please upload a 3D NIfTI file (.nii or .nii.gz)",
    type=["nii", "nii.gz"]
)

# Model selection
model_options = [
    "DS6_UNet3D_woDeform",
    "DS6_UNetMSS3D_woDeform",
    "DS6_UNetMSS3D_wDeform",
    "SMILEUHURA_DS6_UNet3D_woDeform",
    "SMILEUHURA_DS6_UNetMSS3D_woDeform",
    "SMILEUHURA_DS6_UNetMSS3D_wDeform",
    "SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_MIP",
    "SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_mMIP",
    "SMILEUHURA_SPOCKMIP_UNet3D_MIP",
    "SMILEUHURA_SPOCKMIP_UNet3D_mMIP",
    "SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_DS6MIP",
    "SMILEUHURA_DS6_CamSVD_UNetMSS3D_wDeform"
]
selected_model = st.selectbox("Select a pretrained model:", model_options)

# Mode selection
mode_options = ["Full volume inference", "Patch-based inference [Default for all the published works]"]
selected_mode = st.selectbox("Select the inference mode:", mode_options)

# Parameters for patch-based inference
if selected_mode == "Patch-based inference [Default for all the published works]":
    col1, col2, col3 = st.columns(3)
    with col1:
        patch_size = st.number_input("Patch size:", min_value=1, value=64)
        stride_length = st.number_input("Stride length:", min_value=1, value=32)
    with col2:
        batch_size = st.number_input("Batch size:", min_value=1, value=14)
        stride_width = st.number_input("Stride width:", min_value=1, value=32)
    with col3:
        num_worker = st.number_input("Number of workers:", min_value=1, value=3)
        stride_depth = st.number_input("Stride depth:", min_value=1, value=16)

# Process button
process_button = st.button("Process")

if uploaded_file is not None and process_button:
    try:
        # Save the uploaded file to a temporary file
        file_extension = ''.join(Path(uploaded_file.name).suffixes)
        with tempfile.NamedTemporaryFile(suffix=file_extension) as tmp_file:
            tmp_file.write(uploaded_file.read())
            tmp_file.flush()

            # Load the NIfTI file from the temporary file
            nifti_img = nib.load(tmp_file.name)
            data = nifti_img.get_fdata()

        # Convert to PyTorch tensor
        tensor = torch.from_numpy(data).float()

        # Ensure it's 3D
        if tensor.ndim != 3:
            st.error("The uploaded NIfTI file is not a 3D volume. Please upload a valid 3D NIfTI file.")
        else:
            # Display input details
            st.success("File successfully uploaded and read.")
            st.write(f"Input tensor shape: `{tensor.shape}`")
            st.write(f"Selected pretrained model: `{selected_model}`")

            # Construct the model name based on the selected model
            model_name = f"soumickmj/{selected_model}"

            # Load the pre-trained model from Hugging Face
            @st.cache_resource
            def load_model(model_name):
                hf_token = os.environ.get('HF_API_TOKEN')
                if hf_token is None:
                    st.error("Hugging Face API token is not set. Please set the 'HF_API_TOKEN' environment variable.")
                    return None
                try:
                    model = AutoModel.from_pretrained(
                        model_name,
                        trust_remote_code=True,
                        use_auth_token=hf_token
                    )
                    model.eval()
                    return model
                except Exception as e:
                    st.error(f"Failed to load model: {e}")
                    return None

            with st.spinner('Loading the pre-trained model...'):
                model = load_model(model_name)
                if model is None:
                    st.stop()  # Stop the app if the model couldn't be loaded

            # Move model and tensor to CPU (ensure compatibility with Spaces)
            device = torch.device('cpu')
            model = model.to(device)
            tensor = tensor.to(device)

            # Process the tensor through the model
            with st.spinner('Processing the tensor through the model...'):
                if selected_mode == "Full volume inference":
                    st.info("Running full volume inference...")
                    output = infer_full_vol(tensor, model)
                else:
                    st.info("Running patch-based inference [Default for all the published works]...")
                    output = infer_patch_based(tensor, model, patch_size=patch_size, stride_length=stride_length, stride_width=stride_width, stride_depth=stride_depth, batch_size=batch_size, num_worker=num_worker)
            
            st.success("Processing complete.")
            st.write(f"Output tensor shape: `{output.shape}`")
            
            try:
                thresh = threshold_otsu(output)
                output = output > thresh
            except Exception as error:
                st.error(f"Otsu thresholding failed: {error}. Defaulting to a threshold of 0.5.")
                output = output > 0.5  # exception only if input image seems to have just one color 1.0.
            output = output.astype('uint16')
            
            # Save the output as a NIfTI file
            output_img = nib.Nifti1Image(output, affine=nifti_img.affine)
            output_path = tempfile.NamedTemporaryFile(suffix='.nii.gz', delete=False).name
            nib.save(output_img, output_path)

            # Read the saved file for download
            with open(output_path, "rb") as f:
                output_data = f.read()

            # Download button for NIfTI file
            st.download_button(
                label="Download Segmentation Output",
                data=output_data,
                file_name='segmentation_output.nii.gz',
                mime='application/gzip'
            )
            
    except Exception as e:
        st.error(f"An error occurred: {e}")
elif uploaded_file is None:
    st.info("Awaiting file upload...")
elif not process_button:
    st.info("Click the 'Process' button to start processing.")
    
# Footer
st.markdown(
    """
    ---
    ## Credits
    If you like this application, please click on **"Like"** on the top left!
    
    If you use this application and/or any of these models, please cite the following paper:

    ```
    @Article{chatterjee2022ds6,
            AUTHOR = {Chatterjee, Soumick and Prabhu, Kartik and Pattadkal, Mahantesh and Bortsova, Gerda and Sarasaen, Chompunuch and Dubost, Florian and Mattern, Hendrik and de Bruijne, Marleen and Speck, Oliver and Nürnberger, Andreas},
            TITLE = {DS6, Deformation-Aware Semi-Supervised Learning: Application to Small Vessel Segmentation with Noisy Training Data},
            JOURNAL = {Journal of Imaging},
            VOLUME = {8},
            YEAR = {2022},
            NUMBER = {10},
            ARTICLE-NUMBER = {259},
            URL = {https://www.mdpi.com/2313-433X/8/10/259},
            ISSN = {2313-433X},
            DOI = {10.3390/jimaging8100259}
    }
    ```
    
    If you use one of the models with the name starting with `SMILEUHURA`, please addiitonally cite the following paper:

    ```
    @article{chatterjee2023smile,
        title={SMILE-UHURA Challenge},
        author={Chatterjee, S and Mattern, H and Dubost, F and Schreiber, S and Nürnberger, A and Speck, O},
        year={2023},
            doi = {10.7303/syn47164761},
        URL = {https://doi.org/10.7303/syn47164761}
        }
    ```
    
    If you use one of the models that contains `SPOCKMIP` in its name, please addiitonally cite the following paper:

    ```
    @article{radhakrishna2024spockmip,
        title={SPOCKMIP: Segmentation of Vessels in MRAs with Enhanced Continuity using Maximum Intensity Projection as Loss},
        author={Radhakrishna, Chethan and Chintalapati, Karthikesh Varma and Kumar, Sri Chandana Hudukula Ram and Sutrave, Raviteja and Mattern, Hendrik and Speck, Oliver and N{\"u}rnberger, Andreas and Chatterjee, Soumick},
        journal={arXiv preprint arXiv:2407.08655},
        year={2024}
        }
    ```
    
    If you use the `SMILEUHURA_DS6_CamSVD_UNetMSS3D_wDeform` model (i.e. fine-tuned on the CamSVD dataset), please addiitonally cite the following paper:

    ```
    @article{ruiDS62024,
            author = {Li, Rui and Chatterjee, Soumick and Jiaerken, Yeerfan and Radhakrishna, Chethan and Benjamin, Philip and Nannoni, Stefania and Tozer, Daniel J. and Markus, Hugh and Rodgers, Christopher T.},
            title = {A Deep Learning Pipeline for Analysis of the 3D Morphology of the Cerebral Small Perforating Arteries from Time-of-Flight 7 Tesla MRI},
            year = {2024},
            doi = {10.1101/2024.10.03.24314845},
            publisher = {Cold Spring Harbor Laboratory Press},
            URL = {https://www.medrxiv.org/content/early/2024/10/04/2024.10.03.24314845},
            journal = {medRxiv}
        }

    ```
    
    """
)