Updated app.py

app.py

@@ -31,6 +31,10 @@ def run_nnunet_predict(nifti_file):
     os.makedirs(INPUT_DIR, exist_ok=True)
     os.makedirs(OUTPUT_DIR, exist_ok=True)
 
+    # Extract the original filename without the extension
+    original_filename = os.path.basename(nifti_file.name)
+    base_filename = original_filename.replace(".nii.gz", "").replace("_0000", "")
+    
     # 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
@@ -50,19 +54,23 @@ def run_nnunet_predict(nifti_file):
         "-d", "004",                  # Dataset ID
         "-c", "3d_fullres",           # Configuration
         "-tr", "nnUNetTrainer_8000epochs",
-        "-device", "cuda"  # Explicitly use GPU 0
+        "-device", "cuda"  # Explicitly use GPU
     ]
     try:
         subprocess.run(command, check=True)
-        # Get the output file
+
+        # Rename the output file to match the original input filename
         output_file = os.path.join(OUTPUT_DIR, "image.nii.gz")
+        new_output_file = os.path.join(OUTPUT_DIR, f"{base_filename}_LesionMask.nii.gz")
         if os.path.exists(output_file):
-            return output_file
+            os.rename(output_file, new_output_file)
+            return new_output_file
         else:
             return "Error: Output file not found."
     except subprocess.CalledProcessError as e:
         return f"Error: {e}"
 
+
 # Gradio Interface
 interface = gr.Interface(
     fn=run_nnunet_predict,