Create app.py

app.py

@@ -0,0 +1,54 @@
+import gradio as gr
+import subprocess
+import os
+import shutil
+import uuid
+import zipfile
+
+def run_segmentation(uploaded_file, modality):
+    job_id = str(uuid.uuid4())
+    input_filename = f"input_{job_id}.nii.gz"
+    output_folder = f"segmentations_{job_id}"
+    
+    with open(input_filename, "wb") as f:
+        f.write(uploaded_file.read())
+
+    command = ["TotalSegmentator", "-i", input_filename, "-o", output_folder]
+    if modality == "MR":
+        command.extend(["--task", "total_mr"])
+    
+    try:
+        subprocess.run(command, check=True)
+    except subprocess.CalledProcessError as e:
+        return f"Error during segmentation: {e}"
+    
+    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):
+            for file in files:
+                file_path = os.path.join(root, file)
+                arcname = os.path.relpath(file_path, output_folder)
+                zipf.write(file_path, arcname)
+    
+    os.remove(input_filename)
+    shutil.rmtree(output_folder)
+    
+    return zip_filename
+
+with gr.Blocks() as demo:
+    gr.Markdown("# TotalSegmentator Gradio App")
+    gr.Markdown(
+        "Upload a CT or MR image (in NIfTI format) and run segmentation using TotalSegmentator. "
+        "For MR images, the task flag is set accordingly."
+    )
+    
+    with gr.Row():
+        uploaded_file = gr.File(label="Upload NIfTI Image (.nii.gz)")
+        modality = gr.Radio(choices=["CT", "MR"], label="Select Image Modality", value="CT")
+    
+    output_file = gr.File(label="Download Segmentation Output (zip)")
+    
+    run_btn = gr.Button("Run Segmentation")
+    run_btn.click(fn=run_segmentation, inputs=[uploaded_file, modality], outputs=output_file)
+
+demo.launch()