Spaces:

SalmaHassan
/

MINDSETS-APP

Build error

File size: 49,565 Bytes

# 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
# import random
# import threading

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

# st.logo("resources/biomedia.png")

# 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. This feature is particularly useful for clinicians, researchers, and students who wish to examine specific regions of interest within the brain.
# """)

# 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 sliders or arrow keys.
#         - Real-time updates as users navigate through different slices.

#     - **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.

#     - **Segmentation Labels:**
#         - A legend or sidebar displays the labels corresponding to different segmented regions.
#         - Hovering over a segmented area in the image highlights the corresponding label in the legend.
#         - Provides detailed information about each structure, such as name and assigned label value.

#     - **Visualization Controls:**
#         - Zoom and pan functionalities for detailed examination of specific areas.
#     """)

# 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 sliders beneath each image to navigate through the slices.
#         - The slice number and anatomical location are displayed for reference.

#     4. **Interact with the Images:**
#         - Zoom in and out using the scroll wheel or zoom controls.
#         - Pan across the image by clicking and dragging.
#     """)

# 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)
# }

# def create_plotly_colorscale(segmentation_colors):
#     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

# # Create the Plotly-compatible colorscale
# segmentation_colorscale = create_plotly_colorscale(segmentation_colors)

# def get_label_name(value):
#     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',
#     }
#     return label_names.get(value, f'Unknown ({value})')

# def add_segmentation_legend_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):
#     if view == 'Axial':
#         data = np.flip(data, axis=0)
#         seg_data = np.flip(seg_data, axis=0)
#         aspect_ratio = slice_thickness / pixel_spacing[0]
#         height = 800
#         num_slices = data.shape[1]
#     elif view == 'Coronal':
#         aspect_ratio = slice_thickness / pixel_spacing[0]
#         height = 400
#         num_slices = data.shape[0]
#     elif view == 'Sagittal':
#         aspect_ratio = pixel_spacing[1] / pixel_spacing[0]
#         height = 400
#         num_slices = data.shape[2]

#     fig = go.Figure()

#     frames = []
#     for i in range(num_slices):
#         if view == 'Axial':
#             img = data[:, i, :]
#             seg_img = seg_data[:, i, :]
#         elif view == 'Coronal':
#             img = data[i, :, :]
#             seg_img = seg_data[i, :, :]
#         elif view == 'Sagittal':
#             img = np.rot90(data[:, :, i])
#             img = np.flip(img, axis=0)
#             seg_img = np.rot90(seg_data[:, :, i])
#             seg_img = np.flip(seg_img, axis=0)

#         # Create hover text matrix
#         hover_text = np.full(seg_img.shape, '', dtype='object')
#         for y in range(seg_img.shape[0]):
#             for x in range(seg_img.shape[1]):
#                 label_value = int(seg_img[y, x])
#                 if label_value != 0:  # Skip background
#                     hover_text[y, x] = get_label_name(label_value)

#         frame = go.Frame(
#             data=[
#                 go.Heatmap(z=img, colorscale='gray', showscale=False),
#                 go.Heatmap(
#                     z=seg_img,
#                     colorscale=segmentation_colorscale,
#                     showscale=False,
#                     opacity=0.5,
#                     hoverinfo='text',
#                     text=hover_text,
#                     hoverongaps=False
#                 )
#             ],
#             name=f'slice{i}',
#             layout=go.Layout()
#         )
#         frames.append(frame)

#     fig.frames = frames

#     # Add the middle slice as the initial visible layer
#     middle_slice = num_slices // 2
#     middle_hover_text = np.full(frames[middle_slice].data[1].z.shape, '', dtype='object')
#     for y in range(frames[middle_slice].data[1].z.shape[0]):
#         for x in range(frames[middle_slice].data[1].z.shape[1]):
#             label_value = int(frames[middle_slice].data[1].z[y, x])
#             if label_value != 0:
#                 middle_hover_text[y, x] = get_label_name(label_value)

#     fig.add_trace(go.Heatmap(z=frames[middle_slice].data[0].z, colorscale='gray', showscale=False))
#     fig.add_trace(go.Heatmap(
#         z=frames[middle_slice].data[1].z,
#         colorscale=segmentation_colorscale,
#         showscale=False,
#         opacity=0.5,
#         hoverinfo='text',
#         text=middle_hover_text,
#         hoverongaps=False
#     ))

#     fig.update_layout(
#         title=f'{view} View',
#         xaxis=dict(title='X-axis', scaleanchor='y', scaleratio=aspect_ratio),
#         yaxis=dict(title='Y-axis'),
#         height=height,
#         margin=dict(l=0, r=0, t=30, b=0),
#         xaxis_visible=False,
#         yaxis_visible=False,
#         updatemenus=[{
#             'buttons': [
#                 {
#                     'args': [None, {'frame': {'duration': 100, 'redraw': True}, 'fromcurrent': True, 'mode': 'immediate',
#                               'transition': {'duration': 0}}],
#                     'label': '▶',
#                     'method': 'animate'
#                 },
#                 {
#                     'args': [[None], {'frame': {'duration': 0, 'redraw': True}, 'mode': 'immediate',
#                               'transition': {'duration': 0}}],
#                     'label': '◼',
#                     'method': 'animate'
#                 }
#             ],
#             'direction': 'left',
#             'pad': {'r': 10, 't': 87},
#             'showactive': False,
#             'type': 'buttons',
#             'x': 0.1,
#             'xanchor': 'right',
#             'y': 0,
#             'yanchor': 'top'
#         }],
#         sliders=[{
#             'active': middle_slice,
#             'yanchor': 'top',
#             'xanchor': 'left',
#             'currentvalue': {
#                 'font': {'size': 16},
#                 'prefix': 'Slice: ',
#                 'visible': True,
#                 'xanchor': 'right'
#             },
#             'transition': {'duration': 100, 'easing': 'cubic-in-out'},
#             'pad': {'b': 10, 't': 50},
#             'len': 0.9,
#             'x': 0.1,
#             'y': 0,
#             'steps': [{'args': [[f'slice{i}'],
#                                 {'frame': {'duration': 100, 'easing': 'cubic-in-out', 'redraw': True},
#                                  'transition': {'duration': 0}}],
#                        'label': str(i),
#                        'method': 'animate'} for i in range(num_slices)]
#         }]
#     )

#     return fig

# def get_random_fun_fact():
#     fun_facts = [
#         "Regular physical exercise can reduce the risk of developing dementia by up to 50%.",
#         "Eating a Mediterranean diet rich in fruits, vegetables, and omega-3 fatty acids may help prevent dementia.",
#         "Learning a new language or playing a musical instrument can help reduce the risk of developing dementia.",
#         "Socializing and maintaining strong social connections may lower the risk of cognitive decline.",
#         "Some studies suggest that drinking coffee in moderation may help protect against dementia.",
#         "Engaging in mentally stimulating activities like puzzles or reading can help maintain cognitive function.",
#         "Getting enough quality sleep is crucial for brain health and may help prevent dementia.",
#         "Controlling high blood pressure and cholesterol levels can reduce the risk of developing vascular dementia.",
#         "Non-smokers have a lower risk of developing dementia compared to smokers.",
#         "Early detection and treatment of hearing loss may help reduce the risk of cognitive decline."
#     ]
#     return random.choice(fun_facts)

# def prepare_figures(data, seg_data, pixel_spacing, slice_thickness, result):
#     axial_fig = plot_mri_slice(data, seg_data, 'Axial', pixel_spacing, slice_thickness)
#     coronal_fig = plot_mri_slice(data, seg_data, 'Coronal', pixel_spacing, slice_thickness)
#     sagittal_fig = plot_mri_slice(data, seg_data, 'Sagittal', pixel_spacing, slice_thickness)
#     result.extend([axial_fig, coronal_fig, sagittal_fig])

# if uploaded_file:
#     temp_dir = tempfile.mkdtemp()
#     path = os.path.join(temp_dir, uploaded_file.name)
#     with open(path, "wb") as f:
#         f.write(uploaded_file.getvalue())
    
#     img = nib.load(path)
#     data = img.get_fdata()
#     header = img.header

#     pixel_spacing = header['pixdim'][1:3]
#     slice_thickness = header['pixdim'][3]

#     seg_data = np.zeros_like(data)
#     if segments and segmentation_file:
#         seg_path = os.path.join(temp_dir, segmentation_file.name)
#         with open(seg_path, "wb") as f:
#             f.write(segmentation_file.getvalue())
        
#         seg_img = nib.load(seg_path)
#         seg_data = seg_img.get_fdata()

#         # Resample segmentation data to match MRI data dimensions
#         zoom_factors = np.array(data.shape) / np.array(seg_data.shape)
#         seg_data = zoom(seg_data, zoom_factors, order=0)  # Using nearest-neighbor interpolation

#     # Create a placeholder for the loading screen
#     loading_placeholder = st.empty()

#     # Prepare figures in a separate thread
#     result = []
#     thread = threading.Thread(target=prepare_figures, args=(data, seg_data, pixel_spacing, slice_thickness, result))
#     thread.start()

#     # Display loading screen with fun facts
#     start_time = time.time()
#     while thread.is_alive():
#         with loading_placeholder.container():
#             st.markdown("## Loading MRI Visualization")
#             fun_fact = get_random_fun_fact()
#             st.markdown(f"**Fun Fact:** {fun_fact}")
#             progress = min(100, int((time.time() - start_time) / 0.5))  # Assume max 50 seconds loading time
#             st.progress(progress)
#         time.sleep(2)  # Update fun fact every 2 seconds

#     thread.join()  # Ensure the thread has completed

#     # Clear the loading screen
#     loading_placeholder.empty()

#     # Unpack the prepared figures
#     axial_fig, coronal_fig, sagittal_fig = result

#     # Display all figures at once
#     col1, col2 = st.columns([2, 1], gap="small")

#     with col1:
#         axial_chart = st.plotly_chart(axial_fig, use_container_width=True)

#     with col2:
#         coronal_chart = st.plotly_chart(coronal_fig, use_container_width=True)
#         sagittal_chart = st.plotly_chart(sagittal_fig, use_container_width=True)

#     # Add segmentation legend to the sidebar
#     add_segmentation_legend_sidebar()

# # if uploaded_file:
# #     temp_dir = tempfile.mkdtemp()
# #     path = os.path.join(temp_dir, uploaded_file.name)
# #     with open(path, "wb") as f:
# #         f.write(uploaded_file.getvalue())
    
# #     img = nib.load(path)
# #     data = img.get_fdata()
# #     header = img.header

# #     pixel_spacing = header['pixdim'][1:3]
# #     slice_thickness = header['pixdim'][3]

# #     seg_data = np.zeros_like(data)
# #     if segments and segmentation_file:
# #         seg_path = os.path.join(temp_dir, segmentation_file.name)
# #         with open(seg_path, "wb") as f:
# #             f.write(segmentation_file.getvalue())
        
# #         seg_img = nib.load(seg_path)
# #         seg_data = seg_img.get_fdata()

# #         # Resample segmentation data to match MRI data dimensions
# #         zoom_factors = np.array(data.shape) / np.array(seg_data.shape)
# #         seg_data = zoom(seg_data, zoom_factors, order=0)  # Using nearest-neighbor interpolation

# #     # Create a placeholder for the loading screen
# #     loading_placeholder = st.empty()

# #     # Display loading screen
# #     with loading_placeholder.container():
# #         st.markdown("## Loading MRI Visualization")

# #     # Prepare the figures
# #     axial_fig = plot_mri_slice(data, seg_data, 'Axial', pixel_spacing, slice_thickness)
# #     coronal_fig = plot_mri_slice(data, seg_data, 'Coronal', pixel_spacing, slice_thickness)
# #     sagittal_fig = plot_mri_slice(data, seg_data, 'Sagittal', pixel_spacing, slice_thickness)

# #     # Clear the loading screen
# #     loading_placeholder.empty()

# #     # Display all figures at once
# #     col1, col2 = st.columns([2, 1], gap="small")

# #     with col1:
# #         axial_chart = st.plotly_chart(axial_fig, use_container_width=True)

# #     with col2:
# #         coronal_chart = st.plotly_chart(coronal_fig, use_container_width=True)
# #         sagittal_chart = st.plotly_chart(sagittal_fig, use_container_width=True)

# #     # Add segmentation legend to the sidebar
# #     add_segmentation_legend_sidebar()

# 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.")


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")

st.logo("resources/biomedia.png")

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.
        - Note: For performance reasons, only the middle slice and 10 slices before and after are rendered.
    """)

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 to only render middle slice and 10 slices before and after
    """
    # 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)
    
    # Calculate middle slice and range
    middle_idx = num_slices // 2
    slice_range = 10  # Number of slices before and after middle to render
    start_idx = max(0, middle_idx - slice_range)
    end_idx = min(num_slices, middle_idx + slice_range + 1)
    
    # Create frames only for the middle slice and neighboring slices
    frames = []
    slice_labels = []
    
    for i in range(start_idx, end_idx):
        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}'))
        slice_labels.append(str(i))
    
    # Add the initial (middle) slice to the figure
    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
    slider_steps = [
        {
            'args': [
                [f'slice{i}'],
                {'frame': {'duration': 0, 'redraw': True}}
            ],
            'label': str(i),
            'method': 'animate'
        }
        for i in range(start_idx, end_idx)
    ]
    
    # Calculate middle position for slider
    slider_midpoint = (middle_idx - start_idx) / (end_idx - start_idx - 1) if end_idx > start_idx + 1 else 0.5
    
    sliders = [{
        'active': middle_idx - start_idx,  # Set to middle slice
        'currentvalue': {'prefix': 'Slice: ', 'visible': True},
        'pad': {'t': 50},
        'len': 0.9,
        'x': 0.1,
        'y': 0,
        'steps': slider_steps
    }]
    
    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': 200, '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.")