ADD 3 VIEWS

app.py

@@ -10,24 +10,43 @@ plt.switch_backend('Agg')
 
 def load_nrrd(file_obj):
     data, _ = nrrd.read(file_obj)
-    num_slices = data.shape[2]
-    return data, num_slices
+    return data
 
-def visualize_slice(file_obj, slice_index):
-    data, num_slices = load_nrrd(file_obj)
+def visualize_slices(file_obj, slice_index):
+    data = load_nrrd(file_obj)
     
-    # Ensure the slice index is within the range of available slices
-    slice_index = min(max(0, slice_index), num_slices - 1)
-    
-    slice_image = data[:, :, slice_index]
+    num_slices_axial = data.shape[2]
+    num_slices_coronal = data.shape[1]
+    num_slices_sagittal = data.shape[0]
 
-    # Plot the slice
-    fig, ax = plt.subplots()
-    ax.imshow(slice_image, cmap='gray')
-    plt.axis('off')
+    # Ensure the slice index is within the range of available slices for each view
+    slice_index_axial = min(max(0, slice_index), num_slices_axial - 1)
+    slice_index_coronal = min(max(0, slice_index), num_slices_coronal - 1)
+    slice_index_sagittal = min(max(0, slice_index), num_slices_sagittal - 1)
     
+    # Extract slices
+    slice_axial = data[:, :, slice_index_axial]
+    slice_coronal = data[:, slice_index_coronal, :]
+    slice_sagittal = data[slice_index_sagittal, :, :]
+
+    # Plot the slices
+    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+
+    axes[0].imshow(slice_axial, cmap='gray')
+    axes[0].set_title("Axial View")
+    axes[0].axis('off')
+
+    axes[1].imshow(slice_coronal, cmap='gray')
+    axes[1].set_title("Coronal View")
+    axes[1].axis('off')
+
+    axes[2].imshow(slice_sagittal, cmap='gray')
+    axes[2].set_title("Sagittal View")
+    axes[2].axis('off')
+
     # Convert matplotlib figure to PIL Image
     buf = io.BytesIO()
+    plt.tight_layout()
     fig.savefig(buf, format='png')
     plt.close(fig)
     buf.seek(0)
@@ -36,19 +55,21 @@ def visualize_slice(file_obj, slice_index):
     return pil_img
 
 def update_slider(file_obj):
-    _, num_slices = load_nrrd(file_obj)
-    return gr.update(maximum=num_slices-1, value=0)
+    data = load_nrrd(file_obj)
+    # Set maximum based on the axis with the smallest number of slices
+    max_slices = min(data.shape)
+    return gr.update(maximum=max_slices-1, value=0)
 
 with gr.Blocks() as app:
-    gr.Markdown("## NRRD Slice Visualizer")
-    gr.Markdown("Upload an NRRD file and use the slider to select and visualize slices.")
+    gr.Markdown("## NRRD Multi-View Slice Visualizer")
+    gr.Markdown("Upload an NRRD file and use the slider to select and visualize axial, coronal, and sagittal slices.")
     
     file_input = gr.File(label="Upload NRRD File")
     slider = gr.Slider(minimum=0, maximum=1, step=1, value=0, label="Slice Selector")
-    image_output = gr.Image(type="pil", label="Selected Slice")
+    image_output = gr.Image(type="pil", label="Selected Slices (Axial, Coronal, Sagittal)")
     
     file_input.change(fn=update_slider, inputs=file_input, outputs=slider)
-    file_input.change(fn=visualize_slice, inputs=[file_input, slider], outputs=image_output)
-    slider.change(fn=visualize_slice, inputs=[file_input, slider], outputs=image_output)
+    file_input.change(fn=visualize_slices, inputs=[file_input, slider], outputs=image_output)
+    slider.change(fn=visualize_slices, inputs=[file_input, slider], outputs=image_output)
 
 app.launch()