Update app.py

app.py

@@ -172,7 +172,7 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180, center=
     
 # Function to run nnUNet inference
 @spaces.GPU(duration=90)  # Decorate the function to allocate GPU for its execution
-def run_nnunet_predict(nifti_file):
+def run_nnunet_predict(nifti_file,hd_bet=False):
     # Prepare directories
     os.makedirs(INPUT_DIR, exist_ok=True)
     os.makedirs(OUTPUT_DIR, exist_ok=True)
@@ -184,21 +184,21 @@ def run_nnunet_predict(nifti_file):
     # Save the uploaded file to the input directory
     input_path = os.path.join(INPUT_DIR, "image_0000.nii.gz")
     os.rename(nifti_file.name, input_path)  # Move the uploaded file to the expected input location
-    
-    # Apply skull-stripping with HD-BET
-    hd_bet_output_path = os.path.join(INPUT_DIR, "image_0000.nii.gz")
-    try:
-       subprocess.run([
-        "hd-bet",
-        "-i", input_path,
-        "-o", hd_bet_output_path,
-        "-device", "cuda",  # or "cpu"
-        "--disable_tta"   
-        ], check=True)
-       print("Skull-stripping completed.")
-       input_path = hd_bet_output_path
-    except subprocess.CalledProcessError as e:
-       return f"HD-BET Error: {e}"
+
+    if hd_bet:
+       # Apply skull-stripping with HD-BET
+       hd_bet_output_path = os.path.join(INPUT_DIR, "image_0000.nii.gz")
+       try:
+          subprocess.run([
+          "hd-bet",
+          "-i", input_path,
+          "-o", hd_bet_output_path,
+          "-device", "cuda",  # or "cpu"
+          "--disable_tta" ], check=True)
+          print("Skull-stripping completed.")
+          input_path = hd_bet_output_path
+       except subprocess.CalledProcessError as e:
+          return f"HD-BET Error: {e}"
     
     # Debugging: List files in the /tmp/input directory
     print("Files in /tmp/input:")
@@ -268,6 +268,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column(scale=1):
             flair_input = gr.File(label="Upload a FLAIR Image (.nii.gz)")
+            use_hd_bet = gr.Checkbox(label="Apply skull-stripping", value=False)
             submit_button = gr.Button("Submit")
         with gr.Column(scale=2):
             seg_output = gr.File(label="Download the Lesion Segmentation Mask")