Update app.py

app.py

@@ -12,64 +12,59 @@ def load_nrrd(file_obj):
     data, _ = nrrd.read(file_obj)
     return data
 
-def visualize_slices(file_obj, slice_index):
+def visualize_slice(file_obj, slice_index, view):
     data = load_nrrd(file_obj)
     
-    num_slices_axial = data.shape[2]
-    num_slices_coronal = data.shape[1]
-    num_slices_sagittal = data.shape[0]
-
-    # 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)
+    # Choose the appropriate slice based on the selected view
+    if view == "Axial":
+        num_slices = data.shape[2]
+        slice_index = min(max(0, slice_index), num_slices - 1)
+        slice_image = data[:, :, slice_index]
+    elif view == "Coronal":
+        num_slices = data.shape[1]
+        slice_index = min(max(0, slice_index), num_slices - 1)
+        slice_image = data[:, slice_index, :]
+    elif view == "Sagittal":
+        num_slices = data.shape[0]
+        slice_index = min(max(0, slice_index), num_slices - 1)
+        slice_image = data[slice_index, :, :]
+    
+    # Plot the slice
+    fig, ax = plt.subplots(figsize=(5, 5 * slice_image.shape[0] / slice_image.shape[1]))  # Keep aspect ratio
+    ax.imshow(slice_image, cmap='gray')
+    plt.axis('off')
     
-    # Extract slices
-    slice_axial = data[:, :, slice_index_axial]
-    slice_coronal = data[:, slice_index_coronal, :]
-    slice_sagittal = data[slice_index_sagittal, :, :]
-
-    # Plot the slices with the correct aspect ratios
-    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
-
-    axes[0].imshow(slice_axial, cmap='gray', aspect='auto')
-    axes[0].set_title("Axial View")
-    axes[0].axis('off')
-
-    axes[1].imshow(slice_coronal.T, cmap='gray', aspect='auto', extent=[0, data.shape[0], 0, data.shape[2]])
-    axes[1].set_title("Coronal View")
-    axes[1].axis('off')
-
-    axes[2].imshow(slice_sagittal.T, cmap='gray', aspect='auto', extent=[0, data.shape[1], 0, data.shape[2]])
-    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')
+    fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
     plt.close(fig)
     buf.seek(0)
     pil_img = Image.open(buf)
     
     return pil_img
 
-def update_slider(file_obj):
+def update_slider(file_obj, view):
     data = load_nrrd(file_obj)
-    # Set maximum based on the axis with the smallest number of slices
-    max_slices = min(data.shape)
+    if view == "Axial":
+        max_slices = data.shape[2]
+    elif view == "Coronal":
+        max_slices = data.shape[1]
+    elif view == "Sagittal":
+        max_slices = data.shape[0]
     return gr.update(maximum=max_slices-1, value=0)
 
 with gr.Blocks() as app:
-    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.")
+    gr.Markdown("## NRRD Slice Visualizer with View Selector")
+    gr.Markdown("Upload an NRRD file, select a view (Axial, Coronal, Sagittal), and use the slider to select and visualize slices.")
     
     file_input = gr.File(label="Upload NRRD File")
+    view_selector = gr.Radio(choices=["Axial", "Coronal", "Sagittal"], value="Axial", label="View Selector")
     slider = gr.Slider(minimum=0, maximum=1, step=1, value=0, label="Slice Selector")
-    image_output = gr.Image(type="pil", label="Selected Slices (Axial, Coronal, Sagittal)")
+    image_output = gr.Image(type="pil", label="Selected Slice")
     
-    file_input.change(fn=update_slider, inputs=file_input, outputs=slider)
-    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)
+    file_input.change(fn=update_slider, inputs=[file_input, view_selector], outputs=slider)
+    view_selector.change(fn=update_slider, inputs=[file_input, view_selector], outputs=slider)
+    view_selector.change(fn=visualize_slice, inputs=[file_input, slider, view_selector], outputs=image_output)
+    slider.change(fn=visualize_slice, inputs=[file_input, slider, view_selector], outputs=image_output)
 
 app.launch()