Update app.py

app.py

@@ -10,21 +10,24 @@ 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)
-    
-    # Ensure the slice index is within the range of available slices
-    slice_index = min(max(0, slice_index), num_slices - 1)
+def get_slice(data, slice_index, view):
+    if view == "Axial":
+        return data[:, :, slice_index]
+    elif view == "Coronal":
+        return data[:, slice_index, :]
+    elif view == "Sagittal":
+        return data[slice_index, :, :]
+
+def visualize_slice(file_obj, slice_index, view):
+    data = load_nrrd(file_obj)
     
-    slice_image = data[:, :, slice_index]
+    # Get the appropriate slice based on the selected view
+    slice_image = get_slice(data, slice_index, view)
 
-    # Rotate the image 90 degrees clockwise
+    # Rotate the image 90 degrees clockwise and flip it from left to right
     slice_image = np.rot90(slice_image, k=-1)
-    
-    # Flip the image from left to right along the vertical axis
     slice_image = np.fliplr(slice_image)
 
     # Plot the slice
@@ -41,8 +44,14 @@ def visualize_slice(file_obj, slice_index):
     
     return pil_img
 
-def update_slider(file_obj):
-    _, num_slices = load_nrrd(file_obj)
+def update_slider(file_obj, view):
+    data = load_nrrd(file_obj)
+    if view == "Axial":
+        num_slices = data.shape[2]
+    elif view == "Coronal":
+        num_slices = data.shape[1]
+    elif view == "Sagittal":
+        num_slices = data.shape[0]
     return gr.update(maximum=num_slices-1, value=0)
 
 with gr.Blocks() as app:
@@ -50,11 +59,14 @@ with gr.Blocks() as app:
     gr.Markdown("Upload an NRRD file and use the slider to select and visualize slices.")
     
     file_input = gr.File(label="Upload NRRD File")
+    view_selector = gr.Dropdown(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 Slice")
     
-    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=update_slider, inputs=[file_input, view_selector], outputs=slider)
+    view_selector.change(fn=update_slider, inputs=[file_input, view_selector], outputs=slider)
+    file_input.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)
+    view_selector.change(fn=visualize_slice, inputs=[file_input, slider, view_selector], outputs=image_output)
 
 app.launch()