Update app.py

app.py

@@ -10,36 +10,22 @@ plt.switch_backend('Agg')
 
 def load_nrrd(file_obj):
     data, _ = nrrd.read(file_obj)
-    return data
+    num_slices = data.shape[2]
+    return data, num_slices
 
-def get_num_slices(data, view):
-    if view == 'Axial':
-        return data.shape[2]
-    elif view == 'Coronal':
-        return data.shape[1]
-    elif view == 'Sagittal':
-        return data.shape[0]
-
-def extract_slice(data, view, slice_index):
-    if view == 'Axial':
-        # Axial view: XY plane
-        slice_image = data[:, :, slice_index]
-    elif view == 'Coronal':
-        # Coronal view: XZ plane
-        slice_image = data[:, slice_index, :]
-        # Flip the image vertically for correct orientation
-        slice_image = np.flipud(slice_image).T
-    elif view == 'Sagittal':
-        # Sagittal view: YZ plane
-        slice_image = data[slice_index, :, :]
-        # Rotate the image for correct orientation
-        slice_image = np.rot90(slice_image, k=3)
+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)
     
-    return slice_image
+    slice_image = data[:, :, slice_index]
 
-def visualize_slice(file_obj, view, slice_index):
-    data = load_nrrd(file_obj)
-    slice_image = extract_slice(data, view, slice_index)
+    # Rotate the image 90 degrees clockwise
+    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
     fig, ax = plt.subplots()
@@ -55,24 +41,20 @@ def visualize_slice(file_obj, view, slice_index):
     
     return pil_img
 
-def update_slider(file_obj, view):
-    data = load_nrrd(file_obj)
-    num_slices = get_num_slices(data, view)
+def update_slider(file_obj):
+    _, num_slices = load_nrrd(file_obj)
     return gr.update(maximum=num_slices-1, value=0)
 
 with gr.Blocks() as app:
     gr.Markdown("## NRRD Slice Visualizer")
-    gr.Markdown("Upload an NRRD file, select a view, and use the slider to select and visualize slices.")
+    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.Radio(choices=["Axial", "Coronal", "Sagittal"], label="View Selector", value="Axial")
     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, 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, view_selector, slider], outputs=image_output)
-    slider.change(fn=visualize_slice, inputs=[file_input, view_selector, slider], outputs=image_output)
-    view_selector.change(fn=visualize_slice, inputs=[file_input, view_selector, slider], outputs=image_output)
+    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)
 
 app.launch()