Upload 6 files

.gitattributes

@@ -0,0 +1 @@
+unet3d_model.pth filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,89 @@
+---
+title: Brain Tumor Segmentation
+emoji: 🧠
+colorFrom: blue
+colorTo: pink
+sdk: streamlit
+sdk_version: 1.48.1
+app_file: app.py
+pinned: false
+license: mit
+---
+
+## Sample Data
+
+You can download sample NIfTI files for two patients to test the model from this Google Drive link:
+
+[Sample Data (Google Drive)](https://drive.google.com/drive/folders/19LzKOcoIrWQhwY91e_kn644AcQi4tl8z?usp=sharing)
+
+---
+title: Brain Tumor Segmentation
+emoji: 🧠
+colorFrom: blue
+colorTo: pink
+sdk: streamlit
+sdk_version: 1.48.1
+app_file: app.py
+pinned: false
+license: mit
+---
+
+# Brain Tumor Segmentation App
+
+<p align="center">
+  <img src="https://img.shields.io/badge/Streamlit-Online-brightgreen" alt="Streamlit">
+  <a href="https://huggingface.co/spaces/saketh-005/brain-tumor-segmentation"><img src="https://img.shields.io/badge/HuggingFace-Spaces-yellow" alt="Hugging Face Spaces"></a>
+</p>
+
+This project is a web application for brain tumor segmentation from 3D/4D NIfTI MRI scans using a 3D U-Net model, built with PyTorch and Streamlit. You can run it locally or deploy it on [Hugging Face Spaces](https://huggingface.co/spaces).
+
+## Features
+- Upload four 3D NIfTI brain scans (T1, T1ce, T2, FLAIR)
+- Automatic preprocessing and patch-based inference
+- Visualizes the predicted tumor mask overlayed on the MRI
+
+## Quick Start (Hugging Face Spaces)
+
+1. **Upload your trained model file** (`unet3d_model.pth`) to the Space's root directory or use a download link in the code.
+2. Click "Run" or "Duplicate Space" to use your own model.
+3. Use the web interface to upload your NIfTI files and view results.
+
+## Local Usage
+1. Clone this repository:
+   ```sh
+   git clone https://github.com/saketh-005/brain-tumor-segmentation.git
+   cd brain-tumor-segmentation
+   ```
+2. (Recommended) Create and activate a Python virtual environment:
+   ```sh
+   python3 -m venv .venv
+   source .venv/bin/activate
+   ```
+3. Install dependencies:
+   ```sh
+   pip install -r requirements.txt
+   ```
+4. Download the trained model file (`unet3d_model.pth`) and place it in this directory. (Due to file size, it is not included in the repo. Please contact the author or use your own trained model.)
+5. Run the app:
+   ```sh
+   streamlit run app.py
+   ```
+6. Open your browser to [http://localhost:8501](http://localhost:8501) and use the app.
+
+## File Structure
+- `app.py` - Main Streamlit app
+- `unet_model.py` - 3D U-Net model definition
+- `utils.py` - Preprocessing, postprocessing, and visualization utilities
+- `requirements.txt` - Python dependencies
+- `unet3d_model.pth` - Trained model weights (**not included**)
+
+## Notes
+- The model file (`unet3d_model.pth`) must be trained and exported separately.
+- For large files, use cloud storage and provide a download link in this README or in your Hugging Face Space.
+- For best results, ensure all input NIfTI files have the same dimensions and orientation.
+
+## License
+MIT License
+
+## Author
+[Saketh Jangala](https://github.com/saketh-005)

app.py

@@ -0,0 +1,245 @@
+# --- Patch-based Inference Helper ---
+def run_patch_inference(model, tensor, patch_depth=32):
+    """
+    Run model inference on 3D tensor in patches along the depth axis.
+    Args:
+        model: The 3D segmentation model.
+        tensor: Input tensor of shape [1, 4, D, H, W].
+        patch_depth: Depth of each patch.
+    Returns:
+        Output tensor stitched together.
+    """
+    device = next(model.parameters()).device if hasattr(model, 'parameters') else torch.device('cpu')
+    _, c, d, h, w = tensor.shape
+    output = []
+    for start in range(0, d, patch_depth):
+        end = min(start + patch_depth, d)
+        patch = tensor[:, :, start:end, :, :]
+        with torch.no_grad():
+            patch_out = model(patch.to(device))
+        output.append(patch_out.cpu())
+    # Concatenate along the depth axis
+    return torch.cat(output, dim=2)
+import streamlit as st
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import nibabel as nib
+import numpy as np
+import os
+import io
+import tempfile
+
+from utils import preprocess_nifti, postprocess_mask, visualize_prediction, combine_nifti_files
+
+# --- Page Configuration ---
+st.set_page_config(
+    page_title="Brain Tumor Segmentation App",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# --- App Title and Description ---
+st.title("Brain Tumor Segmentation")
+st.write("Upload the four 3D NIfTI brain scans (.nii or .nii.gz) for each modality to get a segmentation mask of the tumor.")
+st.markdown("---")
+
+# --- Model Architecture ---
+# A single block in the U-Net architecture.
+class DoubleConv(nn.Module):
+    """(convolution => GroupNorm => ReLU) * 2"""
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        # 3D convolutional layers, GroupNorm for stable training, and ReLU activation.
+        self.double_conv = nn.Sequential(
+            nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.GroupNorm(num_groups=out_channels // 2, num_channels=out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.GroupNorm(num_groups=out_channels // 2, num_channels=out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+# The downsampling part of the U-Net.
+class Down(nn.Module):
+    """Downscaling with maxpool then double conv"""
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.MaxPool3d(2),
+            DoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.encoder(x)
+
+# The upsampling part of the U-Net.
+class Up(nn.Module):
+    """Upscaling then double conv"""
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        # Use bilinear upsampling and then a convolution layer
+        self.up = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)
+        self.conv = DoubleConv(in_channels, out_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # Pad x1 to match the size of x2 for concatenation
+        diffX = x2.size()[2] - x1.size()[2]
+        diffY = x2.size()[3] - x1.size()[3]
+        diffZ = x2.size()[4] - x1.size()[4]
+        
+        x1 = F.pad(x1, [diffZ // 2, diffZ - diffZ // 2,
+                        diffY // 2, diffY - diffY // 2,
+                        diffX // 2, diffX - diffX // 2])
+        
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+# The final output convolutional layer.
+class Out(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(Out, self).__init__()
+        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+# The complete 3D U-Net model.
+class UNet3d(nn.Module):
+    def __init__(self, n_channels=4, n_classes=3):
+        super().__init__()
+        # The number of classes is 3 (tumor core, edema, enhancing tumor).
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+
+        # Contracting path
+        self.conv = DoubleConv(n_channels, 16)
+        self.enc1 = Down(16, 32)
+        self.enc2 = Down(32, 64)
+        self.enc3 = Down(64, 128)
+        self.enc4 = Down(128, 256)
+
+        # Expansive path
+        self.dec1 = Up(256 + 128, 128)
+        self.dec2 = Up(128 + 64, 64)
+        self.dec3 = Up(64 + 32, 32)
+        self.dec4 = Up(32 + 16, 16)
+        
+        self.out = Out(16, n_classes)
+
+    def forward(self, x):
+        x1 = self.conv(x)
+        x2 = self.enc1(x1)
+        x3 = self.enc2(x2)
+        x4 = self.enc3(x3)
+        x5 = self.enc4(x4)
+        
+        x = self.dec1(x5, x4)
+        x = self.dec2(x, x3)
+        x = self.dec3(x, x2)
+        x = self.dec4(x, x1)
+        
+        logits = self.out(x)
+        return logits
+        
+# --- Model Loading ---
+@st.cache_resource
+def load_model(model_path):
+    """Loads the trained PyTorch model from a .pth file."""
+    try:
+        # FIX: Directly load the model object, which is what was saved.
+        # The weights_only=False argument is needed for custom classes.
+        model = torch.load(model_path, map_location=torch.device('cpu'), weights_only=False)
+        
+        model.eval()
+        st.success("Model loaded successfully!")
+        return model
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        return None
+
+# --- Main App Logic ---
+model_file_path = "unet3d_model.pth"
+if not os.path.exists(model_file_path):
+    st.warning("Model file 'unet3d_model.pth' not found. Please ensure it is in the same directory.")
+    model = None
+else:
+    model = load_model(model_file_path)
+
+st.sidebar.header("Upload NIfTI Files")
+t1_file = st.sidebar.file_uploader("Choose a T1 scan (.nii or .nii.gz)", type=["nii", "gz"], key="t1")
+t1ce_file = st.sidebar.file_uploader("Choose a T1ce scan (.nii or .nii.gz)", type=["nii", "gz"], key="t1ce")
+t2_file = st.sidebar.file_uploader("Choose a T2 scan (.nii or .nii.gz)", type=["nii", "gz"], key="t2")
+flair_file = st.sidebar.file_uploader("Choose a FLAIR scan (.nii or .nii.gz)", type=["nii", "gz"], key="flair")
+
+if t1_file and t1ce_file and t2_file and flair_file and model is not None:
+    st.info("All files uploaded successfully. Processing...")
+    
+    # temp_combined_file_path is now defined at the start of the block
+    temp_combined_file_path = None
+    
+    with st.spinner("Combining NIfTI files and making prediction..."):
+        try:
+            # Create temporary files for each uploaded file
+            with tempfile.NamedTemporaryFile(suffix=f"_{t1_file.name}") as t1_temp, \
+                 tempfile.NamedTemporaryFile(suffix=f"_{t1ce_file.name}") as t1ce_temp, \
+                 tempfile.NamedTemporaryFile(suffix=f"_{t2_file.name}") as t2_temp, \
+                 tempfile.NamedTemporaryFile(suffix=f"_{flair_file.name}") as flair_temp:
+
+                t1_temp.write(t1_file.getvalue())
+                t1ce_temp.write(t1ce_file.getvalue())
+                t2_temp.write(t2_file.getvalue())
+                flair_temp.write(flair_file.getvalue())
+
+                # Pass the temporary file paths to the combine function
+                combined_nifti_img = combine_nifti_files(t1_temp.name, t1ce_temp.name, t2_temp.name, flair_temp.name)
+
+                original_data = combined_nifti_img.get_fdata()
+                
+                # Preprocess the combined image
+                # We need to save the combined NIfTI object to a file for nibabel to load it properly
+                temp_combined_file_path = "combined_4d.nii.gz"
+                nib.save(combined_nifti_img, temp_combined_file_path)
+
+                _, processed_tensor = preprocess_nifti(temp_combined_file_path)
+
+                if original_data is not None and processed_tensor is not None:
+                    st.success("Preprocessing complete!")
+
+                    # --- Patch-based Model Prediction ---
+                    st.info("Running patch-based model inference...")
+                    try:
+                        prediction_tensor = run_patch_inference(model, processed_tensor, patch_depth=32)
+                        st.success("Prediction complete!")
+                    except Exception as e:
+                        st.error(f"Error during patch-based inference: {e}")
+                        raise
+
+                    # Post-process the prediction to get a mask, resizing back to original size
+                    predicted_mask = postprocess_mask(prediction_tensor, original_data.shape)
+
+                    if predicted_mask is not None:
+                        st.header("Results")
+                        # Ensure mask is int and shape matches for visualization
+                        max_slices = original_data.shape[2]
+                        slice_index = st.slider("Select an axial slice to view", 0, max_slices - 1, max_slices // 2)
+                        fig = visualize_prediction(original_data, predicted_mask.astype(int), slice_index=slice_index)
+                        st.pyplot(fig)
+                    else:
+                        st.error("Could not post-process the model's prediction.")
+                
+        except Exception as e:
+            st.error(f"An error occurred during processing: {e}")
+            st.error("Please ensure the uploaded files are valid NIfTI files with the same dimensions.")
+        finally:
+            # Clean up temporary files
+            if os.path.exists(temp_combined_file_path):
+                os.remove(temp_combined_file_path)
+            
+# --- Footer ---
+st.markdown("---")
+st.markdown("Developed with PyTorch and Streamlit.")

requirements.txt

@@ -0,0 +1,9 @@
+streamlit
+torch
+torchvision
+numpy
+matplotlib
+pandas
+nibabel
+Pillow
+scikit-image

unet3d_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24b57e013f9f929e2a2c39af499ec86cf6dae2d42eb4481a0ccd4e2b94a984f7
+size 22655363

unet_model.py

@@ -0,0 +1,95 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class DoubleConv(nn.Module):
+    """(convolution => GroupNorm => ReLU) * 2"""
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.double_conv = nn.Sequential(
+            nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.GroupNorm(num_groups=out_channels // 2, num_channels=out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.GroupNorm(num_groups=out_channels // 2, num_channels=out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+class Down(nn.Module):
+    """Downscaling with maxpool then double conv"""
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.MaxPool3d(2),
+            DoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.encoder(x)
+
+class Up(nn.Module):
+    """Upscaling then double conv"""
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.up = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)
+        self.conv = DoubleConv(in_channels, out_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        diffX = x2.size()[2] - x1.size()[2]
+        diffY = x2.size()[3] - x1.size()[3]
+        diffZ = x2.size()[4] - x1.size()[4]
+        
+        x1 = F.pad(x1, [diffZ // 2, diffZ - diffZ // 2,
+                        diffY // 2, diffY - diffY // 2,
+                        diffX // 2, diffX - diffX // 2])
+        
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+class OutConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv, self).__init__()
+        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+class UNet3d(nn.Module):
+    def __init__(self, n_channels=4, n_classes=3):
+        super().__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+
+        # Contracting path
+        self.conv = DoubleConv(n_channels, 16)
+        self.enc1 = Down(16, 32)
+        self.enc2 = Down(32, 64)
+        self.enc3 = Down(64, 128)
+        self.enc4 = Down(128, 256)
+
+        # Expansive path
+        self.dec1 = Up(256 + 128, 128)
+        self.dec2 = Up(128 + 64, 64)
+        self.dec3 = Up(64 + 32, 32)
+        self.dec4 = Up(32 + 16, 16)
+        
+        self.out = OutConv(16, n_classes)
+
+    def forward(self, x):
+        x1 = self.conv(x)
+        x2 = self.enc1(x1)
+        x3 = self.enc2(x2)
+        x4 = self.enc3(x3)
+        x5 = self.enc4(x4)
+        
+        x = self.dec1(x5, x4)
+        x = self.dec2(x, x3)
+        x = self.dec3(x, x2)
+        x = self.dec4(x, x1)
+        
+        logits = self.out(x)
+        return logits

utils.py

@@ -0,0 +1,123 @@
+import nibabel as nib
+import numpy as np
+import torch
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
+import streamlit as st
+import io
+import tempfile
+from skimage.transform import resize
+
+def preprocess_nifti(nifti_file):
+    """
+    Loads a NIfTI file, preprocesses it, and returns a PyTorch tensor.
+    
+    Args:
+        nifti_file (str or io.BytesIO): Path to the NIfTI file or a file-like object.
+    
+    Returns:
+        tuple: A tuple containing the original image data and a preprocessed tensor.
+    """
+    try:
+        nifti_img = nib.load(nifti_file)
+        img_data = nifti_img.get_fdata()
+        
+        if len(img_data.shape) != 4 or img_data.shape[-1] != 4:
+            st.error("The uploaded NIfTI file must be a 4D image with 4 channels.")
+            return None, None
+            
+        for i in range(img_data.shape[-1]):
+            channel_data = img_data[..., i]
+            if np.max(channel_data) > 0:
+                img_data[..., i] = channel_data / np.max(channel_data)
+        
+        img_data = np.transpose(img_data, (3, 2, 0, 1))
+        
+        tensor_data = torch.from_numpy(img_data).float()
+        tensor_data = torch.unsqueeze(tensor_data, 0)
+        
+        return nifti_img.get_fdata(), tensor_data
+        
+    except Exception as e:
+        st.error(f"Error during NIfTI preprocessing: {e}")
+        return None, None
+
+def postprocess_mask(prediction_tensor, original_shape):
+    """
+    Converts model output (tensor) into a visualizable mask (numpy array)
+    and resizes it to the original image dimensions.
+    """
+    try:
+        probabilities = F.softmax(prediction_tensor, dim=1)
+        
+        mask = torch.argmax(probabilities, dim=1)
+        
+        mask = mask.detach().cpu().numpy()
+        mask = np.squeeze(mask)
+
+        mask = np.transpose(mask, (1, 2, 0))
+
+        resized_mask = resize(mask, original_shape[:3], order=0, preserve_range=True, anti_aliasing=False)
+        
+        return resized_mask
+    except Exception as e:
+        st.error(f"Error during mask post-processing: {e}")
+        return None
+
+def visualize_prediction(original_image, predicted_mask, slice_index=75):
+    """
+    Creates a 2-panel visualization of the original image and the predicted mask.
+    """
+    fig, axes = plt.subplots(1, 2, figsize=(15, 7))
+    # Show FLAIR channel for the original image
+    axes[0].imshow(np.rot90(original_image[:, :, slice_index, 3]), cmap='bone')
+    axes[0].set_title('Original Image (FLAIR)', fontsize=16)
+    axes[0].axis('off')
+
+    axes[1].imshow(np.rot90(original_image[:, :, slice_index, 3]), cmap='bone')
+    # Overlay the predicted mask directly
+    mask_slice = np.rot90(predicted_mask[:, :, slice_index])
+    axes[1].imshow(np.ma.masked_where(mask_slice == 0, mask_slice), cmap='jet', alpha=0.5)
+    axes[1].set_title('Predicted Tumor Mask', fontsize=16)
+    axes[1].axis('off')
+    return fig
+
+def combine_nifti_files(t1_file_path, t1ce_file_path, t2_file_path, flair_file_path):
+    """
+    Combines four 3D NIfTI files from given paths into a single 4D NIfTI file object.
+    
+    Args:
+        t1_file_path, t1ce_file_path, t2_file_path, flair_file_path (str): Paths to the temporary NIfTI files.
+        
+    Returns:
+        nib.Nifti1Image: A 4D NIfTI image object.
+    """
+    try:
+        # Load the four 3D NIfTI files from file paths
+        t1_img = nib.load(t1_file_path)
+        t1ce_img = nib.load(t1ce_file_path)
+        t2_img = nib.load(t2_file_path)
+        flair_img = nib.load(flair_file_path)
+
+        # Get the image data as NumPy arrays
+        t1_data = t1_img.get_fdata()
+        t1ce_data = t1ce_img.get_fdata()
+        t2_data = t2_img.get_fdata()
+        flair_data = flair_img.get_fdata()
+
+        # Ensure all files have the same shape
+        if not (t1_data.shape == t1ce_data.shape == t2_data.shape == flair_data.shape):
+            st.error("Error: Input NIfTI files do not have matching dimensions.")
+            return None
+
+        # Stack the 3D arrays along a new (4th) dimension to create a 4D array
+        combined_data = np.stack([t1_data, t1ce_data, t2_data, flair_data], axis=-1)
+
+        # Create a new 4D NIfTI image object
+        combined_img = nib.Nifti1Image(combined_data, t1_img.affine)
+
+        return combined_img
+    except Exception as e:
+        st.error(f"Error combining NIfTI files: {e}")
+        return None