Update app.py

app.py

@@ -5,7 +5,6 @@ import shutil
 import uuid
 import zipfile
 import nibabel as nib
-import numpy as np
 import matplotlib.pyplot as plt
 
 def run_segmentation(uploaded_file, modality):
@@ -19,7 +18,7 @@ def run_segmentation(uploaded_file, modality):
         with open(input_filename, "wb") as f:
             f.write(uploaded_file.read())
     else:
-        return "Invalid file input", None
+        return None, None, "Invalid file input."
     
     command = ["TotalSegmentator", "-i", input_filename, "-o", output_folder]
     if modality == "MR":
@@ -28,8 +27,11 @@ def run_segmentation(uploaded_file, modality):
     try:
         subprocess.run(command, check=True)
     except subprocess.CalledProcessError as e:
-        return f"Error during segmentation: {e}", None
-
+        error_message = f"Error during segmentation: {e}"
+        if os.path.exists(input_filename): os.remove(input_filename)
+        if os.path.exists(output_folder): shutil.rmtree(output_folder)
+        return None, None, error_message
+    
     zip_filename = f"segmentations_{job_id}.zip"
     with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
         for root, dirs, files in os.walk(output_folder):
@@ -39,6 +41,7 @@ def run_segmentation(uploaded_file, modality):
                 zipf.write(file_path, arcname)
     
     seg_files = [os.path.join(output_folder, f) for f in os.listdir(output_folder) if f.endswith('.nii.gz')]
+    image_filename = None
     if seg_files:
         seg_file = seg_files[0]
         try:
@@ -55,13 +58,11 @@ def run_segmentation(uploaded_file, modality):
         except Exception as e:
             print(f"Error creating preview: {e}")
             image_filename = None
-    else:
-        image_filename = None
 
     os.remove(input_filename)
     shutil.rmtree(output_folder)
     
-    return zip_filename, image_filename
+    return zip_filename, image_filename, "Segmentation completed successfully."
 
 with gr.Blocks() as demo:
     gr.Markdown("# TotalSegmentator Gradio App")
@@ -78,7 +79,9 @@ with gr.Blocks() as demo:
         zip_output = gr.File(label="Download Segmentation Output (zip)")
         preview_output = gr.Image(label="Segmentation Preview")
     
+    status_output = gr.Textbox(label="Status", interactive=False)
+    
     run_btn = gr.Button("Run Segmentation")
-    run_btn.click(fn=run_segmentation, inputs=[uploaded_file, modality], outputs=[zip_output, preview_output])
+    run_btn.click(fn=run_segmentation, inputs=[uploaded_file, modality], outputs=[zip_output, preview_output, status_output])
 
 demo.launch()