Update app.py

app.py

@@ -1,4 +1,4 @@
-import gradio as gr
+import os
 import numpy as np
 import matplotlib.pyplot as plt
 import plotly.graph_objects as go
@@ -7,7 +7,7 @@ from skimage import measure
 from mpl_toolkits.mplot3d.art3d import Poly3DCollection
 from lungmask import LMInferer
 import SimpleITK as sitk
-import os
+import gradio as gr
 
 # --- Lung Segmentation Functions ---
 
@@ -56,15 +56,21 @@ def plot_3d_lungs(lungs_volume, threshold=0.5):
     fig.update_layout(scene_aspectmode='data')  # Maintain aspect ratio
     return fig
 
+# --- Gradio Interface ---
 
+def process_and_visualize(selected_folder):
+    if selected_folder not in ["tumor", "lungs"]:
+        return "Invalid folder selection."  # Handle invalid input
 
-# --- Create Gradio Interface ---
+    volume = segment_lungs_from_dicom(selected_folder)
+    visualization = plot_3d_lungs(volume)
+    return visualization
 
-inputs = gr.Textbox(label="DICOM Folder Path")
+inputs = gr.Dropdown(choices=["tumor", "lungs"], label="Select DICOM Folder")
 output = gr.Plot(label="3D Segmented Lungs")
 
 iface = gr.Interface(
-    fn=lambda dcm_folder: plot_3d_lungs(segment_lungs_from_dicom(dcm_folder)),
+    fn=process_and_visualize,
     inputs=inputs,
     outputs=output,
     title="3D Lung Segmentation Visualization",