vis.py

import streamlit as st
import tempfile
import os
import nibabel as nib
import numpy as np
import plotly.graph_objects as go
from scipy.ndimage import zoom
import time

# Set the page layout to wide
st.set_page_config(layout="wide")

# Optional: logo (uncomment if you have the resource)
# st.image("resources/biomedia.png", width=200)

st.markdown(
    """
    <style>
    /* Change the font size of the sidebar */
    .sidebar .sidebar-content {
        font-size: 2px;  
    }
    /* Change the line spacing of the sidebar */
    .sidebar .sidebar-content p {
        line-height: 0.2;  /* Adjust the line height as needed */
    }
    </style>
    """,
    unsafe_allow_html=True
)

st.title('MRI Scan Visualization :brain:')

st.write("""
         The **MRI Visualization** page offers an interactive platform for users to explore 3D MRI scans in detail. It is designed to enhance the understanding of anatomical structures and any associated segmentation data.
""")

with st.expander("### **Features:**"):
    st.write("""
    - **Upload MRI Scans:**
        - Users can upload their own 3D MRI scans in NIfTI format (`.nii` or `.nii.gz`).
        - An optional checkbox allows users to upload a corresponding segmentation mask if available.
    
    - **Interactive Slice Navigation:**
        - Provides axial, coronal, and sagittal views of the MRI scan.
        - Users can scroll through slices in each plane using the interactive slider.

    - **Segmentation Overlay:**
        - When a segmentation mask is provided, it overlays onto the MRI images.
        - Segmented regions are highlighted with transparent colors to maintain visibility of underlying anatomy.
    """)

with st.expander("## **Instructions:**"):
    st.write("""
    1. **Upload Your MRI Scan:**
        - Use the file uploader to select your MRI scan file.
        - Supported formats: NIfTI (`.nii`, `.nii.gz`).

    2. **Upload Segmentation Mask (Optional):**
        - If you have a segmentation mask, check the **"Include Segmentation Mask"** checkbox.
        - Upload your segmentation mask file in NIfTI format.

    3. **View MRI Slices:**
        - After uploading, the application will display the MRI scan in three views: axial, coronal, and sagittal.
        - Use the slider beneath each visualization to navigate through the slices.
    """)

uploaded_file = st.file_uploader("Upload MRI file", type=["gz", "nii", "nii.gz"])

segments = st.checkbox("Show Segmentation")
segmentation_file = None

if segments:
    segmentation_file = st.file_uploader("Upload segmentation file", type=["gz", "nii", "nii.gz"])

# Define segmentation colors
segmentation_colors = {
    0: '#000000',  # Background (black)
    2: '#759479',  # Left cerebral white matter (purple)
    3: '#f7d587',  # Left cerebral cortex (orange)
    4: '#bb7761',  # Left lateral ventricle (brown)
    5: '#53bad5',  # Left inferior lateral ventricle (cyan)
    7: '#e95c47',  # Left cerebellum white matter (red)
    8: '#ec7c5e',  # Left cerebellum cortex (light salmon)
    10: '#6cf185', # Left thalamus (green)
    11: '#ce6253', # Left caudate (brown)
    12: '#d6f600', # Left putamen (yellow)
    13: '#513e00', # Left pallidum (saddle brown)
    14: '#fffad9', # 3rd ventricle (white)
    15: '#e8dc00', # 4th ventricle (yellow)
    16: '#c8c8ee', # Brain-stem (lavender)
    17: '#fafacd', # Left hippocampus (lemon chiffon)
    18: '#fad500', # Left amygdala (yellow)
    24: '#009ad3', # CSF (cyan)
    26: '#e95c47', # Left accumbens area (red)
    28: '#bc9be1', # Left ventral DC (lavender)
    41: '#afd7d3', # Right cerebral white matter (light cyan)
    42: '#8ebed1', # Right cerebral cortex (light blue)
    43: '#53bad5', # Right lateral ventricle (cyan)
    44: '#a9d5f5', # Right inferior lateral ventricle (cyan)
    46: '#00af5c', # Right cerebellum white matter (light green)
    47: '#4ac87c', # Right cerebellum cortex (pale green)
    49: '#00bfff', # Right thalamus (cyan)
    50: '#009400', # Right caudate (green)
    51: '#d2e772', # Right putamen (green yellow)
    52: '#513e00', # Right pallidum (saddle brown)
    53: '#d0ffff', # Right hippocampus (light cyan)
    54: '#90fdfb', # Right amygdala (cyan)
    58: '#6de3e6', # Right accumbens area (cyan)
    60: '#cc3400', # Right ventral DC (red)
}

# Label names for legend
label_names = {
    2: 'Left cerebral white matter',
    3: 'Left cerebral cortex',
    4: 'Left lateral ventricle',
    5: 'Left inferior lateral ventricle',
    7: 'Left cerebellum white matter',
    8: 'Left cerebellum cortex',
    10: 'Left thalamus',
    11: 'Left caudate',
    12: 'Left putamen',
    13: 'Left pallidum',
    14: '3rd ventricle',
    15: '4th ventricle',
    16: 'Brain-stem',
    17: 'Left hippocampus',
    18: 'Left amygdala',
    24: 'CSF',
    26: 'Left accumbens area',
    28: 'Left ventral DC',
    41: 'Right cerebral white matter',
    42: 'Right cerebral cortex',
    43: 'Right lateral ventricle',
    44: 'Right inferior lateral ventricle',
    46: 'Right cerebellum white matter',
    47: 'Right cerebellum cortex',
    49: 'Right thalamus',
    50: 'Right caudate',
    51: 'Right putamen',
    52: 'Right pallidum',
    53: 'Right hippocampus',
    54: 'Right amygdala',
    58: 'Right accumbens area',
    60: 'Right ventral DC',
}

@st.cache_data
def create_plotly_colorscale(segmentation_colors):
    """Create a Plotly-compatible colorscale for segmentation"""
    max_value = max(segmentation_colors.keys())
    colorscale = []
    for i in range(max_value + 1):
        color = segmentation_colors.get(i, '#000000')
        colorscale.append([i / max_value, color])
    return colorscale

def get_label_name(value):
    """Get the label name for a segmentation value"""
    return label_names.get(value, f'Unknown ({value})')

def add_segmentation_legend_sidebar():
    """Add segmentation legend to sidebar"""
    st.sidebar.title("Segmentation Labels")
    for value, color in segmentation_colors.items():
        if value != 0:  # Skip background
            st.sidebar.markdown(f"<span style='color:{color};'>⬤</span> {get_label_name(value)}", unsafe_allow_html=True)

def plot_mri_slice(data, seg_data, view, pixel_spacing, slice_thickness):
    """
    Plot MRI slice with segmentation overlay - optimized version
    """
    # Create an empty frames list to store all slices
    frames = []
    
    # Set appropriate dimensions and orientation based on view
    if view == 'Axial':
        data = np.flip(data, axis=0)
        seg_data = np.flip(seg_data, axis=0) if seg_data is not None else None
        aspect_ratio = slice_thickness / pixel_spacing[0]
        height = 800
        num_slices = data.shape[1]
        
        # Function to extract slice for this view
        def extract_slice(i):
            img = data[:, i, :]
            seg_img = None if seg_data is None else seg_data[:, i, :]
            return img, seg_img
            
    elif view == 'Coronal':
        aspect_ratio = slice_thickness / pixel_spacing[0]
        height = 400
        num_slices = data.shape[0]
        
        # Function to extract slice for this view
        def extract_slice(i):
            img = data[i, :, :]
            seg_img = None if seg_data is None else seg_data[i, :, :]
            return img, seg_img
            
    elif view == 'Sagittal':
        aspect_ratio = pixel_spacing[1] / pixel_spacing[0]
        height = 400
        num_slices = data.shape[2]
        
        # Function to extract slice for this view
        def extract_slice(i):
            img = np.rot90(data[:, :, i])
            img = np.flip(img, axis=0)
            seg_img = None
            if seg_data is not None:
                seg_img = np.rot90(seg_data[:, :, i])
                seg_img = np.flip(seg_img, axis=0)
            return img, seg_img
    
    # Initialize figure
    fig = go.Figure()
    
    # Get colorscale for segmentation
    segmentation_colorscale = create_plotly_colorscale(segmentation_colors)
    
    # Create frames for each slice - this is the key optimization
    # We only create a limited number of frames for performance but keep smooth animation
    # For very large volumes, we sample every nth slice
    skip_factor = max(1, num_slices // 100)  # Don't create more than ~100 frames
    
    for i in range(0, num_slices, skip_factor):
        img, seg_img = extract_slice(i)
        
        frame_data = [go.Heatmap(z=img, colorscale='gray', showscale=False)]
        
        if seg_img is not None:
            # Only create hover text for visible segmentation values (non-zero)
            hover_text = np.full(seg_img.shape, '', dtype='object')
            for unique_val in np.unique(seg_img):
                if unique_val > 0 and unique_val in label_names:
                    mask = seg_img == unique_val
                    hover_text[mask] = get_label_name(unique_val)
            
            frame_data.append(go.Heatmap(
                z=seg_img,
                colorscale=segmentation_colorscale,
                showscale=False,
                opacity=0.5,
                hoverinfo='text',
                text=hover_text,
                hoverongaps=False
            ))
        
        frames.append(go.Frame(data=frame_data, name=f'slice{i}'))
    
    # Add the initial (middle) slice to the figure
    middle_idx = num_slices // 2
    middle_img, middle_seg = extract_slice(middle_idx)
    
    fig.add_trace(go.Heatmap(z=middle_img, colorscale='gray', showscale=False))
    
    if middle_seg is not None:
        # Create hover text for middle slice
        middle_hover = np.full(middle_seg.shape, '', dtype='object')
        for unique_val in np.unique(middle_seg):
            if unique_val > 0 and unique_val in label_names:
                mask = middle_seg == unique_val
                middle_hover[mask] = get_label_name(unique_val)
        
        fig.add_trace(go.Heatmap(
            z=middle_seg,
            colorscale=segmentation_colorscale,
            showscale=False,
            opacity=0.5,
            hoverinfo='text',
            text=middle_hover,
            hoverongaps=False
        ))
    
    # Add frames to the figure
    fig.frames = frames
    
    # Configure layout with slider
    sliders = [{
        'active': num_slices // (2 * skip_factor),  # Set to middle slice
        'currentvalue': {'prefix': 'Slice: ', 'visible': True},
        'pad': {'t': 50},
        'len': 0.9,
        'x': 0.1,
        'y': 0,
        'steps': [
            {
                'args': [
                    [f'slice{i}'],
                    {'frame': {'duration': 0, 'redraw': True}}
                ],
                'label': str(i),
                'method': 'animate'
            }
            for i in range(0, num_slices, skip_factor)
        ]
    }]
    
    fig.update_layout(
        title=f'{view} View',
        height=height,
        xaxis=dict(scaleanchor='y', scaleratio=aspect_ratio),
        margin=dict(l=0, r=0, t=30, b=0),
        xaxis_visible=False,
        yaxis_visible=False,
        sliders=sliders,
        updatemenus=[{
            'buttons': [
                {
                    'args': [None, {'frame': {'duration': 500, 'redraw': True}}],
                    'label': '▶',
                    'method': 'animate'
                },
                {
                    'args': [[None], {'frame': {'duration': 0, 'redraw': True}}],
                    'label': '◼',
                    'method': 'animate'
                }
            ],
            'type': 'buttons',
            'direction': 'left',
            'showactive': False,
            'x': 0.1,
            'y': 0,
            'pad': {'r': 10, 't': 60}
        }]
    )
    
    return fig

# Main execution
if uploaded_file:
    # Create progress indicator
    progress_bar = st.progress(0)
    status_text = st.empty()
    
    # Step 1: Load MRI file
    status_text.text("Loading MRI file...")
    temp_dir = tempfile.mkdtemp()
    mri_path = os.path.join(temp_dir, uploaded_file.name)
    with open(mri_path, "wb") as f:
        f.write(uploaded_file.getvalue())
    
    # Load MRI image
    progress_bar.progress(25)
    img = nib.load(mri_path)
    data = img.get_fdata()
    header = img.header
    
    # Get pixel spacing for aspect ratio
    pixel_spacing = header['pixdim'][1:3]
    slice_thickness = header['pixdim'][3]
    
    # Step 2: Load segmentation if provided
    progress_bar.progress(40)
    seg_data = None
    if segments and segmentation_file:
        status_text.text("Loading segmentation file...")
        seg_path = os.path.join(temp_dir, segmentation_file.name)
        with open(seg_path, "wb") as f:
            f.write(segmentation_file.getvalue())
        
        # Load segmentation
        seg_img = nib.load(seg_path)
        seg_data = seg_img.get_fdata()
        
        # Resample segmentation if dimensions don't match
        if seg_data.shape != data.shape:
            status_text.text("Resampling segmentation to match MRI dimensions...")
            zoom_factors = np.array(data.shape) / np.array(seg_data.shape)
            seg_data = zoom(seg_data, zoom_factors, order=0)  # Use nearest-neighbor interpolation
    
    # Step 3: Create visualizations
    progress_bar.progress(60)
    status_text.text("Creating visualizations...")
    
    # Create columns for layout
    col1, col2 = st.columns([2, 1], gap="small")
    
    # Create axial view (top-down)
    with col1:
        progress_bar.progress(70)
        status_text.text("Rendering axial view...")
        axial_fig = plot_mri_slice(data, seg_data, 'Axial', pixel_spacing, slice_thickness)
        axial_chart = st.plotly_chart(axial_fig, use_container_width=True)
    
    with col2:
        # Create coronal view (front-back)
        progress_bar.progress(85)
        status_text.text("Rendering coronal view...")
        coronal_fig = plot_mri_slice(data, seg_data, 'Coronal', pixel_spacing, slice_thickness)
        coronal_chart = st.plotly_chart(coronal_fig, use_container_width=True)
        
        # Create sagittal view (side)
        progress_bar.progress(95)
        status_text.text("Rendering sagittal view...")
        sagittal_fig = plot_mri_slice(data, seg_data, 'Sagittal', pixel_spacing, slice_thickness)
        sagittal_chart = st.plotly_chart(sagittal_fig, use_container_width=True)
    
    # Add segmentation legend to sidebar if using segmentation
    if segments and segmentation_file and seg_data is not None:
        add_segmentation_legend_sidebar()
    
    # Cleanup
    progress_bar.progress(100)
    status_text.text("Visualization complete!")
    time.sleep(1)
    status_text.empty()
    progress_bar.empty()
    
    # Clean up temporary files
    try:
        os.remove(mri_path)
        if segments and segmentation_file:
            os.remove(seg_path)
        os.rmdir(temp_dir)
    except:
        pass  # Ignore cleanup errors
else:
    st.info("Please upload an MRI file to begin visualization.")