Updated app.py

app.py

@@ -3,11 +3,29 @@ import subprocess
 import os
 import shutil
 
-# Paths
+export nnUNet_results="./model"
+
+# Define paths
+MODEL_DIR = "./model"  # Local directory to store the downloaded model
+DATASET_DIR = os.path.join(MODEL_DIR, "Dataset004_WML")  # Directory for Dataset004_WML
 INPUT_DIR = "/tmp/input"
 OUTPUT_DIR = "/tmp/output"
-MODEL_DIR = "./model"
 
+# Hugging Face Model Repository
+REPO_ID = "FrancescoLR/FLAMeS-model"  # Replace with your actual model repository ID
+
+# Function to download the model files from Hugging Face Model Hub
+def download_model():
+    if not os.path.exists(DATASET_DIR):
+        os.makedirs(DATASET_DIR, exist_ok=True)
+        print("Downloading Dataset004_WML...")
+        hf_hub_download(repo_id=REPO_ID, filename="Dataset004_WML.zip", cache_dir=MODEL_DIR)
+        # Unzip the dataset into the correct location
+        subprocess.run(["unzip", "-o", os.path.join(MODEL_DIR, "Dataset004_WML.zip"), "-d", DATASET_DIR])
+        os.remove(os.path.join(MODEL_DIR, "Dataset004_WML.zip"))
+        print("Dataset004_WML downloaded and extracted.")
+
+# Function to run nnUNet inference
 def run_nnunet_predict(nifti_file):
     # Prepare directories
     os.makedirs(INPUT_DIR, exist_ok=True)
@@ -15,14 +33,15 @@ def run_nnunet_predict(nifti_file):
 
     # Save the uploaded file to the input directory
     input_path = os.path.join(INPUT_DIR, "image.nii.gz")
-    shutil.copy(nifti_file.name, input_path)
+    with open(input_path, "wb") as f:
+        f.write(nifti_file.read())
 
     # Set environment variables for nnUNet
-    os.environ["nnUNet_raw"] = MODEL_DIR
-    os.environ["nnUNet_preprocessed"] = MODEL_DIR
-    os.environ["nnUNet_results"] = MODEL_DIR
+    os.environ["nnUNet_raw"] = DATASET_DIR
+    os.environ["nnUNet_preprocessed"] = DATASET_DIR
+    os.environ["nnUNet_results"] = DATASET_DIR
 
-    # Construct the nnUNetv2_predict command
+    # Construct and run the nnUNetv2_predict command
     command = [
         "nnUNetv2_predict",
         "-i", INPUT_DIR,
@@ -31,8 +50,6 @@ def run_nnunet_predict(nifti_file):
         "-c", "3d_fullres",           # Configuration
         "-tr", "nnUNetTrainer_8000epochs",  # Trainer name
     ]
-
-    # Run the command
     try:
         subprocess.run(command, check=True)
         # Get the output file
@@ -46,14 +63,13 @@ interface = gr.Interface(
     fn=run_nnunet_predict,
     inputs=gr.File(label="Upload FLAIR Image (.nii.gz)"),
     outputs=gr.File(label="Download Segmentation Mask"),
-    title="FLAMeS: FLAIR Lesion Analysis in Multiple Sclerosis",
-    description=(
-        "Upload a skull-stripped FLAIR image (.nii.gz) to generate a binary segmentation of MS lesions. "
-        "This model uses nnUNetv2 for inference with ensemble predictions."
-    ),
+    title="FLAMeS: Multiple Sclerosis Lesion Segmentation",
+    description="Upload a skull-stripped FLAIR image (.nii.gz) to generate a binary segmentation of MS lesions."
 )
 
+# Download model files before launching the app
+download_model()
+
 # Launch the app
 if __name__ == "__main__":
     interface.launch()
-