Update app.py

app.py

@@ -10,24 +10,22 @@ from lcmv_class import LCMVSourceEstimator
 BUILT_IN_GPSC = Path("ghw280_from_egig.gpsc")
 BUILT_IN_SRC = Path("fsaverage-vol-5mm-src.fif")
 
-# Validate files exist
 if not BUILT_IN_GPSC.is_file():
     raise FileNotFoundError(f"Required montage file not found: {BUILT_IN_GPSC}")
 if not BUILT_IN_SRC.is_file():
     raise FileNotFoundError(f"Required source space file not found: {BUILT_IN_SRC}")
 
-# Prepare fsaverage structure to avoid MNE download
+# Prepare fsaverage structure
 FS_DIR = Path("derivatives/lcmv/fsaverage")
 BEM_DIR = FS_DIR / "bem"
 BEM_DIR.mkdir(parents=True, exist_ok=True)
 
-# Copy source space to expected location
 EXPECTED_SRC = Path("derivatives/lcmv/fsaverage-vol-5mm-src.fif")
 EXPECTED_SRC.parent.mkdir(parents=True, exist_ok=True)
 if not EXPECTED_SRC.exists():
     shutil.copy(BUILT_IN_SRC, EXPECTED_SRC)
 
-# Fixed subject/task (no user input)
+# Fixed subject/task
 SUBJECT_ID = "sub"
 TASK = ""
 
@@ -35,7 +33,8 @@ def run_lcmv(
     cleaned_fif,
     run_difumo: bool = True,
     reg: float = 0.01,
-    n_jobs: int = 1,  # HF Spaces: use 1 to avoid OOM
+    n_jobs: int = 1,
+    download_option: str = "full_zip",  # "full_zip" or "difumo_only"
 ):
     abs_base = Path(".").resolve()
     subject_output = abs_base / "derivatives" / "lcmv" / f"{SUBJECT_ID}_{TASK}"
@@ -59,9 +58,7 @@ def run_lcmv(
     }
 
     try:
-        # Set SUBJECTS_DIR so MNE finds fsaverage
         os.environ['SUBJECTS_DIR'] = str(abs_base / "derivatives" / "lcmv")
-
         estimator = LCMVSourceEstimator(config)
         metadata = estimator.run_enhanced_computation()
 
@@ -69,25 +66,29 @@ def run_lcmv(
             difumo_config = {'dimension': 512, 'resolution_mm': 2}
             estimator.run_difumo_extraction(difumo_config=difumo_config)
 
-        # Bundle outputs
-        zip_path = abs_base / f"{SUBJECT_ID}_{TASK}_lcmv_output.zip"
-        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
-            for root, _, files in os.walk(subject_output):
-                for file in files:
-                    file_path = Path(root) / file
-                    arcname = file_path.relative_to(abs_base)
-                    zf.write(file_path, arcname)
-
-        return str(zip_path)
+        # ✅ Dual output logic
+        if download_option == "difumo_only":
+            difumo_file = subject_output / "difumo_time_courses.npy"
+            if difumo_file.exists():
+                return str(difumo_file)
+            else:
+                raise FileNotFoundError("DiFuMo file not found. Ensure 'Extract DiFuMo' is enabled.")
+
+        else:  # full_zip
+            zip_path = abs_base / f"{SUBJECT_ID}_{TASK}_lcmv_output.zip"
+            with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+                for root, _, files in os.walk(subject_output):
+                    for file in files:
+                        file_path = Path(root) / file
+                        arcname = file_path.relative_to(abs_base)
+                        zf.write(file_path, arcname)
+            return str(zip_path)
 
     except Exception as e:
         error_log = subject_output / "lcmv_error.log"
         with open(error_log, "w") as f:
             f.write(f"LCMV failed:\n{str(e)}")
-        zip_path = abs_base / f"{SUBJECT_ID}_{TASK}_lcmv_error.zip"
-        with zipfile.ZipFile(zip_path, "w") as zf:
-            zf.write(error_log, error_log.name)
-        return str(zip_path)
+        return str(error_log)
 
 # Gradio UI
 with gr.Blocks(theme=gr.themes.Base(), title="LCMV Source Estimation") as demo:
@@ -100,15 +101,21 @@ with gr.Blocks(theme=gr.themes.Base(), title="LCMV Source Estimation") as demo:
             run_difumo = gr.Checkbox(True, label="Extract DiFuMo 512 time courses")
             reg = gr.Number(0.01, label="LCMV Regularization (reg)")
             n_jobs = gr.Number(1, label="Parallel Jobs (n_jobs)", precision=0)
+            download_option = gr.Radio(
+                choices=["full_zip", "difumo_only"],
+                value="full_zip",
+                label="Download Option",
+                info="• full_zip: All outputs (STC, metadata, plots, DiFuMo)\n• difumo_only: Only difumo_time_courses.npy (for connectivity)"
+            )
             run_btn = gr.Button("Run LCMV", variant="primary")
 
         with gr.Column():
-            output_file = gr.File(label="Download LCMV Results (.zip)")
+            output_file = gr.File(label="Download Output")
 
     run_btn.click(
         fn=run_lcmv,
-        inputs=[fif_input, run_difumo, reg, n_jobs],
+        inputs=[fif_input, run_difumo, reg, n_jobs, download_option],
         outputs=output_file,
     )
 
-demo.launch()
+demo.launch()