patching added for trying

app.py

@@ -1,18 +1,15 @@
 import streamlit as st
-import json
 import math
 import numpy as np
 import nibabel as nib
 import torch
 import torch.nn.functional as F
-import scipy.io
-from io import BytesIO
 from transformers import AutoModel
 import os
 import tempfile
 from pathlib import Path
-import pandas as pd
 from skimage.filters import threshold_otsu
+import torchio as tio
 
 def infer_full_vol(tensor, model):
     tensor = torch.movedim(tensor, -1, -3)
@@ -46,6 +43,37 @@ def infer_full_vol(tensor, model):
     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))
+    overlap = np.subtract(patch_size, (stride_length, stride_width, stride_depth))
+    
+    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)
+        for _, patches_batch in enumerate(patch_loader):
+            local_batch = 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)
+
+        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)",
@@ -62,7 +90,7 @@ with st.sidebar:
     
     **Instructions**:
     - Upload your 3D NIfTI file (`.nii` or `.nii.gz`). It should be a single-slice cardiac long-axis dynamic CINE scan, where the first dimension represents time.
-    - Select a seed value from the dropdown menu.
+    - Select a model from the dropdown menu.
     - Click the "Process" button to generate the latent factors.
     """)
     st.markdown("---")
@@ -77,10 +105,14 @@ uploaded_file = st.file_uploader(
     type=["nii", "nii.gz"]
 )
 
-# Seed selection
+# Model selection
 model_options = ["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 DS6]"]
+selected_mode = st.selectbox("Select the running mode:", mode_options)
+
 # Process button
 process_button = st.button("Process")
 
@@ -111,7 +143,7 @@ if uploaded_file is not None and process_button:
             # Add batch and channel dimensions
             tensor = tensor.unsqueeze(0).unsqueeze(0)  # Shape: [1, 1, D, H, W]
 
-            # Construct the model name based on the selected seed
+            # Construct the model name based on the selected model
             model_name = f"soumickmj/{selected_model}"
 
             # Load the pre-trained model from Hugging Face
@@ -145,7 +177,12 @@ if uploaded_file is not None and process_button:
 
             # Process the tensor through the model
             with st.spinner('Processing the tensor through the model...'):
-                output = infer_full_vol(tensor, 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 DS6]...")
+                    output = infer_patch_based(tensor, model)
             
             st.success("Processing complete.")
             st.write(f"Output tensor shape: `{output.shape}`")

requirements.txt

@@ -1,7 +1,5 @@
+scikit-image
 nibabel
 torch
-pytorch_lightning
-scipy
 transformers
-torchvision
-scikit-image
+torchio