from __future__ import annotations import contextlib import io import os import re import shutil import stat import tempfile import time from pathlib import Path import nibabel as nib import numpy as np import streamlit as st from data_prep.prep_utils import ( DatasetConfig, combine_standardized, combine_standardized_images_only, run_prep, run_prep_images_only, ) from inference.streamlit_panel import render_inference_panel PROJECT_ROOT = Path(__file__).resolve().parents[1] WEBAPP_UPLOAD_ROOT = PROJECT_ROOT / "data" / "raw" / "webapp_upload" PREP_OUT_ROOT = PROJECT_ROOT / "data" / "prep_outputs" MODEL_INPUT_DEST = PROJECT_ROOT / "data" / "processed" / "test_input" MODEL_OUTPUT_DEST = PROJECT_ROOT / "data" / "processed" / "test_output" HF_RUNTIME_CACHE_ROOT = PROJECT_ROOT / "data" / "runtime_bundle_from_hf" HF_RUNTIME_DEFAULT_REPO = os.environ.get("HF_RUNTIME_REPO_ID", "rmbielski/Stroke_Dependencies") HF_RUNTIME_DEFAULT_REVISION = os.environ.get("HF_RUNTIME_REVISION", "main") HF_MODEL_CACHE_ROOT = PROJECT_ROOT / "data" / "model_bundle_from_hf" HF_MODEL_DEFAULT_REPO = os.environ.get("HF_MODEL_REPO_ID", "rmbielski/ARC_ATLAS_v3.1") HF_MODEL_DEFAULT_REVISION = os.environ.get("HF_MODEL_REVISION", HF_RUNTIME_DEFAULT_REVISION) HF_MODEL_DEFAULT_SUBDIR = os.environ.get("HF_MODEL_SUBDIR", "") HF_DATA_CACHE_ROOT = PROJECT_ROOT / "data" / "hf_data_cache" PUBLIC_HF_DATASETS = ["rmbielski/Atlas_2", "rmbielski/ARC", "rmbielski/Approximate_Numeracy"] NATIVE_PREVIEW_ROOT = PROJECT_ROOT / "data" / "raw" / "webapp_native_preview" ANTS_PROFILE_CHOICES = ["quick", "fast", "balanced", "accurate"] HF_RUNTIME_ALLOW_PATTERNS = [ "tools/ants/**", "data/templateflow/**", "**/tools/ants/**", "**/data/templateflow/**", ] @st.cache_data(show_spinner=False) def _load_nifti(path: str): img = nib.load(path) data = np.asarray(img.dataobj) zooms = img.header.get_zooms() return data, tuple(float(z) for z in zooms[:3]) @st.cache_data(show_spinner=False) def _load_nifti_from_bytes(file_name: str, payload: bytes): suffix = ".nii.gz" if file_name.lower().endswith(".nii.gz") else ".nii" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: tmp.write(payload) tmp_path = tmp.name try: return _load_nifti(tmp_path) finally: try: os.remove(tmp_path) except OSError: pass def _is_nifti_filename(name: str) -> bool: lower = name.lower() return lower.endswith(".nii") or lower.endswith(".nii.gz") def _clean_name(name: str) -> str: cleaned = re.sub(r"[^a-zA-Z0-9._-]+", "_", name.strip()) return cleaned or "input.nii.gz" def _save_upload(uploaded_file, directory: Path) -> Path: directory.mkdir(parents=True, exist_ok=True) out_path = directory / _clean_name(uploaded_file.name) out_path.write_bytes(uploaded_file.getbuffer()) return out_path def _clear_dir(path: Path) -> None: if path.exists(): shutil.rmtree(path) path.mkdir(parents=True, exist_ok=True) def _ensure_executable_binary(path: Path) -> Path: if not path.exists(): raise FileNotFoundError(f"Missing binary: {path}") # HF dataset downloads can drop execute bits; restore them before use. try: mode = path.stat().st_mode path.chmod(mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) except Exception: pass if os.access(path, os.X_OK): return path # Fallback: copy to a local executable cache and retry. exec_dir = Path(tempfile.gettempdir()) / "stroke_runtime_exec" exec_dir.mkdir(parents=True, exist_ok=True) fallback = exec_dir / path.name shutil.copy2(path, fallback) fallback.chmod(fallback.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) if os.access(fallback, os.X_OK): return fallback raise PermissionError( f"Binary is not executable: {path}. " "This may be due to file permissions or a non-executable mount." ) def _persist_native_preview(uploaded_t1, uploaded_mask=None): t1_dir = NATIVE_PREVIEW_ROOT / "t1" mask_dir = NATIVE_PREVIEW_ROOT / "mask" _clear_dir(t1_dir) _save_upload(uploaded_t1, t1_dir) if uploaded_mask is not None: _clear_dir(mask_dir) _save_upload(uploaded_mask, mask_dir) elif mask_dir.exists(): shutil.rmtree(mask_dir, ignore_errors=True) def _load_persisted_native_preview(): t1_dir = NATIVE_PREVIEW_ROOT / "t1" mask_dir = NATIVE_PREVIEW_ROOT / "mask" if not t1_dir.exists(): return None, None, None t1_paths = sorted(t1_dir.glob("*.nii*")) if not t1_paths: return None, None, None t1_path = t1_paths[-1] t1_data, t1_zooms = _load_nifti(str(t1_path)) mask_data = None mask_paths = sorted(mask_dir.glob("*.nii*")) if mask_dir.exists() else [] if mask_paths: mask_data, _ = _load_nifti(str(mask_paths[-1])) return t1_data, t1_zooms, mask_data def _discover_runtime_bundle_root(download_root: Path) -> Path: direct_reg = download_root / "tools" / "ants" / "bin" / "antsRegistration" if direct_reg.exists(): return download_root candidates = sorted(download_root.rglob("tools/ants/bin/antsRegistration")) for reg in candidates: bundle_root = reg.parents[3] if (bundle_root / "tools" / "ants" / "bin" / "antsApplyTransforms").exists(): return bundle_root raise FileNotFoundError( "Could not find runtime bundle layout under downloaded HF dataset. " "Expected `tools/ants/bin/antsRegistration` and `tools/ants/bin/antsApplyTransforms`." ) def _set_runtime_env(bundle_root: Path) -> dict[str, str]: reg = bundle_root / "tools" / "ants" / "bin" / "antsRegistration" app = bundle_root / "tools" / "ants" / "bin" / "antsApplyTransforms" tf_home = bundle_root / "data" / "templateflow" missing = [str(p) for p in (reg, app, tf_home) if not p.exists()] if missing: raise FileNotFoundError( "Runtime bundle is missing required paths:\n" + "\n".join(missing) ) reg_exec = _ensure_executable_binary(reg) app_exec = _ensure_executable_binary(app) os.environ["ANTS_REG"] = str(reg_exec) os.environ["ANTS_APPLY"] = str(app_exec) os.environ["TEMPLATEFLOW_HOME"] = str(tf_home) # Disable fallback auto-install once HF runtime is configured. os.environ["ANTS_AUTO_INSTALL"] = "0" return { "runtime_bundle_root": str(bundle_root), "ants_registration": str(reg_exec), "ants_apply_transforms": str(app_exec), "templateflow_home": str(tf_home), } def _hf_local_dir_for_repo(repo_id: str, revision: str | None = None) -> Path: safe = re.sub(r"[^a-zA-Z0-9._-]+", "__", repo_id.strip()) safe_rev = re.sub(r"[^a-zA-Z0-9._-]+", "__", (revision or "main").strip() or "main") return HF_RUNTIME_CACHE_ROOT / safe / safe_rev class _HFFileWrapper: """Minimal shim that satisfies the UploadedFile interface (.name + .getbuffer()).""" def __init__(self, name: str, data: bytes) -> None: self.name = name self._data = data def getbuffer(self) -> bytes: return self._data def _list_nifti_files_from_hf( repo_id: str, revision: str = "main", token: str | None = None, ) -> tuple[list[str], list[str]]: """Return (t1_files, mask_files) lists of relative NIfTI paths in the dataset.""" try: from huggingface_hub import list_repo_files except Exception as exc: raise RuntimeError( "huggingface_hub is required. Install with `pip install huggingface_hub`." ) from exc all_files = list( list_repo_files( repo_id=repo_id.strip(), repo_type="dataset", revision=revision or "main", token=(token.strip() if token else None), ) ) nifti = sorted( f for f in all_files if f.lower().endswith(".nii") or f.lower().endswith(".nii.gz") ) _MASK_KEYWORDS = {"mask", "lesion", "seg", "label"} mask_files = [f for f in nifti if any(k in f.lower() for k in _MASK_KEYWORDS)] t1_files = [f for f in nifti if f not in set(mask_files)] return t1_files, mask_files def _download_hf_nifti( repo_id: str, filename: str, revision: str = "main", token: str | None = None, ) -> bytes: """Download a single NIfTI file from an HF dataset and return its bytes.""" try: from huggingface_hub import hf_hub_download except Exception as exc: raise RuntimeError( "huggingface_hub is required. Install with `pip install huggingface_hub`." ) from exc safe_repo = re.sub(r"[^a-zA-Z0-9._-]+", "__", repo_id.strip()) cache_dir = HF_DATA_CACHE_ROOT / safe_repo / (revision or "main") cache_dir.mkdir(parents=True, exist_ok=True) local_path = hf_hub_download( repo_id=repo_id.strip(), filename=filename, repo_type="dataset", revision=revision or "main", token=(token.strip() if token else None), local_dir=str(cache_dir), ) return Path(local_path).read_bytes() def _sync_runtime_bundle_from_hf( repo_id: str, revision: str, token: str | None = None, force_download: bool = False, ) -> dict[str, str]: try: from huggingface_hub import snapshot_download except Exception as exc: raise RuntimeError( "huggingface_hub is required to download runtime dependencies. " "Install it with `pip install huggingface_hub`." ) from exc repo_id = repo_id.strip() if not repo_id: raise ValueError("Runtime dataset repo_id is empty.") revision = (revision or "main").strip() or "main" local_dir = _hf_local_dir_for_repo(repo_id, revision) if force_download and local_dir.exists(): shutil.rmtree(local_dir, ignore_errors=True) local_dir.mkdir(parents=True, exist_ok=True) if not force_download: try: bundle_root = _discover_runtime_bundle_root(local_dir) except FileNotFoundError: pass else: info = _set_runtime_env(bundle_root) info["hf_repo_id"] = repo_id info["hf_revision"] = revision info["hf_local_dir"] = str(local_dir) info["hf_cache_hit"] = "true" return info snapshot_download( repo_id=repo_id, repo_type="dataset", local_dir=str(local_dir), allow_patterns=HF_RUNTIME_ALLOW_PATTERNS, revision=revision, token=(token.strip() if token else None), ) bundle_root = _discover_runtime_bundle_root(local_dir) info = _set_runtime_env(bundle_root) info["hf_repo_id"] = repo_id info["hf_revision"] = revision info["hf_local_dir"] = str(local_dir) info["hf_cache_hit"] = "false" return info def _paired_mask_path(t1_path: Path, mask_dir: Path) -> Path | None: if not mask_dir.exists(): return None key = t1_path.name.replace("_T1w_MNI_norm", "") for mask_path in sorted(mask_dir.glob("*.nii*")): mask_key = mask_path.name.replace("_lesion_mask_MNI_clean", "") if mask_key == key: return mask_path return None def _normalize_slice(slice_data: np.ndarray) -> np.ndarray: low, high = np.percentile(slice_data, [1, 99]) if low == high: low, high = float(slice_data.min()), float(slice_data.max()) if low == high: low, high = 0.0, 1.0 normalized = (slice_data - low) / (high - low) return np.clip(normalized, 0.0, 1.0) def _select_volume(data: np.ndarray, key_prefix: str, label: str) -> np.ndarray: if data.ndim <= 3: return data if data.shape[-1] <= 1: return np.squeeze(data, axis=-1) vol_idx = st.slider( f"{label} volume index", 0, data.shape[-1] - 1, 0, key=f"{key_prefix}_vol_idx", ) return data[..., vol_idx] def _overlay_mask(normalized_slice: np.ndarray, mask_slice: np.ndarray, opacity: float) -> np.ndarray: rgb = np.stack([normalized_slice, normalized_slice, normalized_slice], axis=-1) mask_bool = np.asarray(mask_slice) > 0 overlay_color = np.array([1.0, 0.0, 0.0], dtype=rgb.dtype) rgb[mask_bool] = (1.0 - opacity) * rgb[mask_bool] + opacity * overlay_color return rgb def _render_inline_image(container, preview: np.ndarray, caption: str) -> None: arr = np.asarray(preview, dtype=np.float32) arr = np.clip(np.nan_to_num(arr, nan=0.0, posinf=1.0, neginf=0.0), 0.0, 1.0) container.image((arr * 255.0).astype(np.uint8), caption=caption, use_container_width=True) def _render_scaled_image(preview: np.ndarray, caption: str, key_prefix: str): left, _ = st.columns([1, 1]) _render_inline_image(left, preview, caption) def _render_overlay_from_arrays( image_data: np.ndarray, image_zooms: tuple[float, float, float], mask_data: np.ndarray | None, key_prefix: str, title: str, ): st.markdown(f"**{title}**") image_vol = _select_volume(image_data, key_prefix, title) if image_vol.ndim < 2: st.error(f"Unsupported image dimensions: {image_data.shape}") return col_a, col_b, col_c = st.columns(3) col_a.metric("Image shape", str(image_data.shape)) col_b.metric("Image dtype", str(image_data.dtype)) col_c.metric("Voxel size", " x ".join(f"{z:.3g}" for z in image_zooms)) if image_vol.ndim == 2: slice_data = np.asarray(image_vol) caption = "2D image" axis = None idx = None else: axis_options = {"Axial (Z)": 2, "Coronal (Y)": 1, "Sagittal (X)": 0} axis_label = st.selectbox( f"{title} slice axis", list(axis_options.keys()), key=f"{key_prefix}_axis", ) axis = axis_options[axis_label] idx = st.slider( f"{title} slice index", 0, image_vol.shape[axis] - 1, image_vol.shape[axis] // 2, key=f"{key_prefix}_slice", ) slice_data = np.take(image_vol, idx, axis=axis) caption = f"{axis_label} slice {idx}" preview = _normalize_slice(np.asarray(slice_data)) if mask_data is not None: mask_vol = _select_volume(mask_data, f"{key_prefix}_mask", f"{title} mask") mask_slice = None if axis is None: if mask_vol.shape == image_vol.shape: mask_slice = mask_vol else: st.warning("Mask shape does not match image shape. Showing image only.") else: if mask_vol.shape == image_vol.shape: mask_slice = np.take(mask_vol, idx, axis=axis) else: st.warning( f"Mask shape {mask_vol.shape} does not match image shape {image_vol.shape}. " "Showing image only." ) if mask_slice is not None: opacity = st.slider("Mask opacity", 0.0, 1.0, 0.4, 0.05, key=f"{key_prefix}_mask_alpha") preview = _overlay_mask(preview, np.asarray(mask_slice), opacity) _render_scaled_image(preview, caption=caption, key_prefix=key_prefix) def _coerce_3d_for_compare(volume: np.ndarray) -> np.ndarray | None: if volume.ndim == 3: return volume if volume.ndim == 4 and volume.shape[-1] == 1: return np.squeeze(volume, axis=-1) if volume.ndim >= 4: return volume[..., 0] return None def _render_side_by_side_comparison( native_data: np.ndarray, native_mask: np.ndarray | None, prepared_data: np.ndarray, prepared_mask: np.ndarray | None, ): native_vol = _coerce_3d_for_compare(native_data) prepared_vol = _coerce_3d_for_compare(prepared_data) if native_vol is None or prepared_vol is None: st.info( "Side-by-side comparison currently supports 3D volumes (or 4D with a singleton/first volume)." ) return st.markdown("**Native vs Registered (side-by-side)**") axis_options = {"Axial (Z)": 2, "Coronal (Y)": 1, "Sagittal (X)": 0} axis_label = st.selectbox("Comparison slice axis", list(axis_options.keys()), key="cmp_axis") axis = axis_options[axis_label] slice_pct = st.slider("Comparison slice position (%)", 0, 100, 50, key="cmp_slice_pct") idx_native = int(round((native_vol.shape[axis] - 1) * (slice_pct / 100.0))) idx_prepared = int(round((prepared_vol.shape[axis] - 1) * (slice_pct / 100.0))) show_masks = st.checkbox("Overlay masks in comparison", value=True, key="cmp_show_masks") opacity = st.slider("Comparison mask opacity", 0.0, 1.0, 0.4, 0.05, key="cmp_mask_alpha") native_slice = np.take(native_vol, idx_native, axis=axis) prepared_slice = np.take(prepared_vol, idx_prepared, axis=axis) native_preview = _normalize_slice(np.asarray(native_slice)) prepared_preview = _normalize_slice(np.asarray(prepared_slice)) if show_masks and native_mask is not None: native_mask_vol = _coerce_3d_for_compare(native_mask) if native_mask_vol is not None and native_mask_vol.shape == native_vol.shape: native_mask_slice = np.take(native_mask_vol, idx_native, axis=axis) native_preview = _overlay_mask(native_preview, np.asarray(native_mask_slice), opacity) if show_masks and prepared_mask is not None: prepared_mask_vol = _coerce_3d_for_compare(prepared_mask) if prepared_mask_vol is not None and prepared_mask_vol.shape == prepared_vol.shape: prepared_mask_slice = np.take(prepared_mask_vol, idx_prepared, axis=axis) prepared_preview = _overlay_mask(prepared_preview, np.asarray(prepared_mask_slice), opacity) left, right = st.columns(2) _render_inline_image( left, native_preview, caption=f"Native {axis_label} slice {idx_native} / {native_vol.shape[axis] - 1}", ) _render_inline_image( right, prepared_preview, caption=f"Registered {axis_label} slice {idx_prepared} / {prepared_vol.shape[axis] - 1}", ) def _run_registration(uploaded_t1, uploaded_mask, already_mni: bool, prep_profile: str) -> tuple[bool, str]: t0 = time.monotonic() _clear_dir(WEBAPP_UPLOAD_ROOT) image_dir = WEBAPP_UPLOAD_ROOT / "Images" mask_dir = WEBAPP_UPLOAD_ROOT / "Masks" t1_saved = _save_upload(uploaded_t1, image_dir) has_mask = uploaded_mask is not None if has_mask: _save_upload(uploaded_mask, mask_dir) _clear_dir(MODEL_INPUT_DEST) log_buffer = io.StringIO() with contextlib.redirect_stdout(log_buffer): os.environ["PREP_ANTS_PROFILE"] = (prep_profile or "quick").strip().lower() print("[prep] staging complete") print(f"[prep] uploaded_t1={t1_saved.name}") print(f"[prep] uploaded_mask={'yes' if has_mask else 'no'}") print(f"[prep] already_mni={already_mni}") print(f"[prep] ants_profile={os.environ.get('PREP_ANTS_PROFILE')}") print(f"[prep] command_timeout_sec={os.environ.get('PREP_CMD_TIMEOUT_SEC', os.environ.get('ANTS_CMD_TIMEOUT_SEC', '900'))}") if has_mask: datasets = [ DatasetConfig( name="WEBAPP", images_dir=image_dir, masks_dir=mask_dir, t1_glob="**/*.nii*", mask_glob="**/*.nii*", already_mni=already_mni, overwrite=True, ) ] print("[prep] running image+mask prep pipeline...") outputs = run_prep(datasets, PREP_OUT_ROOT, force_overwrite=True) if not outputs: raise RuntimeError("No outputs were generated from image+mask prep.") print("[prep] combining standardized image+mask outputs...") combine_standardized(outputs, MODEL_INPUT_DEST) else: print("[prep] running image-only prep pipeline...") out_ds = run_prep_images_only( images_dir=image_dir, out_root=PREP_OUT_ROOT, name="WEBAPP", t1_glob="**/*.nii*", already_mni=already_mni, overwrite=True, ) print("[prep] combining standardized image-only outputs...") combine_standardized_images_only([out_ds], MODEL_INPUT_DEST) prepared_t1 = sorted((MODEL_INPUT_DEST / "t1").glob("*.nii*")) if not prepared_t1: raise RuntimeError( "Prep finished but no registered/normalized images were found in " f"{MODEL_INPUT_DEST / 't1'}." ) message = ( f"Prepared {len(prepared_t1)} image(s). " f"Model input root: {MODEL_INPUT_DEST}. " f"Uploaded source: {t1_saved.name}" ) elapsed = time.monotonic() - t0 return has_mask, message + f"\nPrep wall time: {elapsed:.2f}s\n\n" + log_buffer.getvalue() def _render_hf_dataset_browser() -> None: """Render the HF dataset browser inside the 'Browse HF Dataset' tab.""" token = (st.session_state.get("hf_runtime_token") or "").strip() or None dataset_options = PUBLIC_HF_DATASETS + ["Custom..."] selected_preset = st.selectbox("Dataset", dataset_options, key="hf_browser_preset") if selected_preset == "Custom...": repo_id = st.text_input( "Repository ID", key="hf_browser_custom_repo", placeholder="org/dataset-name", ) else: repo_id = selected_preset revision = st.text_input( "Revision (branch / tag / commit)", value="main", key="hf_browser_revision", ) if token: st.caption("HF token from Runtime Setup will be used (supports private repositories).") else: st.caption("For private repositories, enter your HF token in the Runtime Setup section above.") if st.button("List files", key="hf_browser_list_btn") and repo_id: with st.spinner(f"Listing NIfTI files in {repo_id}..."): try: t1_files, mask_files = _list_nifti_files_from_hf(repo_id, revision, token) st.session_state["hf_browser_t1_files"] = t1_files st.session_state["hf_browser_mask_files"] = mask_files st.session_state["hf_browser_list_error"] = "" st.session_state["hf_browser_repo_id"] = repo_id st.session_state["hf_browser_revision_used"] = revision except Exception as exc: st.session_state["hf_browser_list_error"] = str(exc) st.session_state["hf_browser_t1_files"] = [] st.session_state["hf_browser_mask_files"] = [] if st.session_state.get("hf_browser_list_error"): st.error(st.session_state["hf_browser_list_error"]) t1_files: list[str] = st.session_state.get("hf_browser_t1_files", []) mask_files: list[str] = st.session_state.get("hf_browser_mask_files", []) if t1_files or mask_files: if not t1_files: st.warning("No T1-like NIfTI files found. All NIfTI files were classified as masks.") selected_t1 = None else: selected_t1 = st.selectbox("T1 image file", t1_files, key="hf_browser_selected_t1") mask_options = ["— none —"] + mask_files selected_mask_label = st.selectbox( "Mask file (optional)", mask_options, key="hf_browser_selected_mask" ) selected_mask = None if selected_mask_label == "— none —" else selected_mask_label if st.button("Load selected", key="hf_browser_load_btn", disabled=selected_t1 is None): cached_repo = st.session_state.get("hf_browser_repo_id", repo_id) cached_rev = st.session_state.get("hf_browser_revision_used", revision) try: with st.spinner(f"Downloading {Path(selected_t1).name}..."): t1_bytes = _download_hf_nifti(cached_repo, selected_t1, cached_rev, token) st.session_state["hf_loaded_t1"] = {"name": Path(selected_t1).name, "bytes": t1_bytes} if selected_mask: with st.spinner(f"Downloading {Path(selected_mask).name}..."): mask_bytes = _download_hf_nifti(cached_repo, selected_mask, cached_rev, token) st.session_state["hf_loaded_mask"] = { "name": Path(selected_mask).name, "bytes": mask_bytes, } else: st.session_state["hf_loaded_mask"] = None st.session_state["hf_load_error"] = "" except Exception as exc: st.session_state["hf_load_error"] = str(exc) if st.session_state.get("hf_load_error"): st.error(st.session_state["hf_load_error"]) loaded_t1 = st.session_state.get("hf_loaded_t1") loaded_mask = st.session_state.get("hf_loaded_mask") if loaded_t1: mask_label = f" | Mask: {loaded_mask['name']}" if loaded_mask else "" st.success(f"Loaded: {loaded_t1['name']}{mask_label}") if st.button("Clear HF selection", key="hf_browser_clear_btn"): for k in ("hf_loaded_t1", "hf_loaded_mask", "hf_browser_t1_files", "hf_browser_mask_files", "hf_load_error"): st.session_state.pop(k, None) st.rerun() def render_app(): st.set_page_config(page_title="Chronic Stroke Segmentation Prep", layout="wide") st.title("Input Registration + Prep") st.write( "Upload a T1 NIfTI and optionally a lesion mask, preview native-space data, " "run registration/normalization via `src/data_prep`, and compare outputs side-by-side." ) with st.expander("0) Runtime setup (Hugging Face dataset)", expanded=False): st.write( "Runtime dependencies are pulled from your HF dataset and cached locally. " "These paths are used for every registration run." ) runtime_repo_id = st.text_input( "Runtime dataset repo_id", value=HF_RUNTIME_DEFAULT_REPO, key="hf_runtime_repo_id", help="Example: rmbielski/Stroke_Dependencies", ) runtime_revision = st.text_input( "Dataset revision", value=HF_RUNTIME_DEFAULT_REVISION, key="hf_runtime_revision", help="Branch, tag, or commit; usually `main`.", ) runtime_token = st.text_input( "HF token (optional)", value=os.environ.get("HF_TOKEN", ""), type="password", key="hf_runtime_token", help="Required only for private datasets.", ) runtime_force = st.checkbox( "Force re-download runtime bundle", value=False, key="hf_runtime_force", ) setup_btn = st.button("Sync runtime bundle from HF", key="setup_runtime_btn") if setup_btn: with st.spinner("Syncing runtime bundle from Hugging Face..."): try: runtime_info = _sync_runtime_bundle_from_hf( repo_id=runtime_repo_id, revision=runtime_revision, token=runtime_token, force_download=runtime_force, ) except Exception as exc: st.session_state["runtime_setup_error"] = str(exc) st.session_state["runtime_setup_info"] = None else: st.session_state["runtime_setup_error"] = "" st.session_state["runtime_setup_info"] = runtime_info if st.session_state.get("runtime_setup_error"): st.error(st.session_state["runtime_setup_error"]) if st.session_state.get("runtime_setup_info"): st.success("Runtime bundle is available and active.") st.code("\n".join(f"{k}: {v}" for k, v in st.session_state["runtime_setup_info"].items())) if "native_upload_key_version" not in st.session_state: st.session_state["native_upload_key_version"] = 0 upload_key_version = int(st.session_state["native_upload_key_version"]) with st.expander("1) Upload inputs", expanded=True): upload_tab, hf_tab = st.tabs(["Upload file", "Browse HF Dataset"]) with upload_tab: uploaded_t1 = st.file_uploader( "T1 MRI (.nii/.nii.gz)", type=None, key=f"prep_t1_{upload_key_version}", ) include_mask = st.checkbox("I also want to upload a mask", value=False, key="prep_has_mask") uploaded_mask = st.file_uploader( "Lesion mask (.nii/.nii.gz)", type=None, key=f"prep_mask_{upload_key_version}", disabled=not include_mask, ) with hf_tab: _render_hf_dataset_browser() already_mni = st.checkbox( "Input is already in MNI space (skip ANTs registration)", value=False, key="prep_already_mni", ) prep_profile = st.selectbox( "Registration profile", ANTS_PROFILE_CHOICES, index=ANTS_PROFILE_CHOICES.index( (os.environ.get("PREP_ANTS_PROFILE") or "quick").strip().lower() if (os.environ.get("PREP_ANTS_PROFILE") or "quick").strip().lower() in ANTS_PROFILE_CHOICES else "quick" ), key="prep_ants_profile", ) run_prep_btn = st.button("Register + normalize", type="primary", key="prep_run_btn") # Resolve effective files: local upload takes priority; fall back to HF-loaded. effective_t1 = uploaded_t1 effective_mask = uploaded_mask if include_mask else None if effective_t1 is None and st.session_state.get("hf_loaded_t1"): d = st.session_state["hf_loaded_t1"] effective_t1 = _HFFileWrapper(d["name"], d["bytes"]) if effective_mask is None and st.session_state.get("hf_loaded_mask"): d = st.session_state["hf_loaded_mask"] effective_mask = _HFFileWrapper(d["name"], d["bytes"]) include_mask = True native_data, native_zooms, native_mask_data = _load_persisted_native_preview() if effective_t1 is not None: if not _is_nifti_filename(effective_t1.name): st.error("T1 file must be `.nii` or `.nii.gz`.") elif include_mask and effective_mask is not None and not _is_nifti_filename(effective_mask.name): st.error("Mask file must be `.nii` or `.nii.gz`.") else: try: _persist_native_preview( uploaded_t1=effective_t1, uploaded_mask=effective_mask, ) native_data, native_zooms, native_mask_data = _load_persisted_native_preview() except Exception as exc: st.error(f"Failed to persist native preview files: {exc}") st.subheader("2) Native-space viewer (before registration)") if st.button("Clear native viewer results", key="clear_native_results_btn"): if NATIVE_PREVIEW_ROOT.exists(): shutil.rmtree(NATIVE_PREVIEW_ROOT, ignore_errors=True) st.session_state["native_upload_key_version"] = upload_key_version + 1 st.rerun() if native_data is None: st.info("Upload a T1 file or load one from the HF Dataset browser above to preview the native-space MRI.") else: _render_overlay_from_arrays( image_data=native_data, image_zooms=native_zooms, mask_data=native_mask_data, key_prefix="native_view", title="Native uploaded image", ) if run_prep_btn: if effective_t1 is None: st.error("Upload a T1 NIfTI or load one from the HF Dataset browser before running registration.") elif not _is_nifti_filename(effective_t1.name): st.error("T1 file must be `.nii` or `.nii.gz`.") elif include_mask and effective_mask is None: st.error("Mask is enabled but no mask file was provided.") elif include_mask and effective_mask is not None and not _is_nifti_filename(effective_mask.name): st.error("Mask file must be `.nii` or `.nii.gz`.") else: stage_status = st.empty() t0_total = time.monotonic() with st.spinner("Syncing runtime bundle + running registration pipeline..."): try: stage_status.info("Step 1/2: Resolving runtime bundle...") t0_sync = time.monotonic() runtime_info = _sync_runtime_bundle_from_hf( repo_id=runtime_repo_id, revision=runtime_revision, token=runtime_token, force_download=runtime_force, ) sync_elapsed = time.monotonic() - t0_sync st.session_state["runtime_setup_error"] = "" st.session_state["runtime_setup_info"] = runtime_info cache_hit = str(runtime_info.get("hf_cache_hit", "false")).lower() == "true" cache_text = "cache hit" if cache_hit else "downloaded" stage_status.info( "Step 2/2: Running registration + normalization " f"(runtime {cache_text}, sync {sync_elapsed:.1f}s)..." ) t0_prep = time.monotonic() has_mask, prep_log = _run_registration( uploaded_t1=effective_t1, uploaded_mask=effective_mask, already_mni=already_mni, prep_profile=prep_profile, ) prep_elapsed = time.monotonic() - t0_prep except Exception as exc: st.session_state["prep_error"] = str(exc) st.session_state["prep_log"] = "" st.session_state["prep_ready"] = False st.session_state["prep_message"] = "" stage_status.error("Registration pipeline failed. Check 'Prep logs' and traceback details.") else: total_elapsed = time.monotonic() - t0_total st.session_state["prep_error"] = "" runtime_summary = ( "[runtime]\n" f"hf_repo_id={runtime_info.get('hf_repo_id')}\n" f"hf_revision={runtime_info.get('hf_revision')}\n" f"hf_local_dir={runtime_info.get('hf_local_dir')}\n" f"hf_cache_hit={runtime_info.get('hf_cache_hit')}\n" f"already_mni={already_mni}\n" f"ants_profile={os.environ.get('PREP_ANTS_PROFILE', 'balanced')}\n" f"ants_registration={runtime_info.get('ants_registration')}\n" f"ants_apply_transforms={runtime_info.get('ants_apply_transforms')}\n" f"templateflow_home={runtime_info.get('templateflow_home')}\n" f"prep_cmd_timeout_sec={os.environ.get('PREP_CMD_TIMEOUT_SEC', os.environ.get('ANTS_CMD_TIMEOUT_SEC', '900'))}\n" "\n[timing]\n" f"sync_seconds={sync_elapsed:.2f}\n" f"prep_seconds={prep_elapsed:.2f}\n" f"total_seconds={total_elapsed:.2f}\n" ) st.session_state["prep_log"] = runtime_summary + "\n" + prep_log st.session_state["prep_ready"] = True st.session_state["prep_has_mask_out"] = has_mask st.session_state["prep_message"] = ( "Registration/prep completed. Viewer and model-ready files were updated. " f"Total time: {total_elapsed / 60.0:.1f} min." ) stage_status.success( f"Registration pipeline completed in {total_elapsed / 60.0:.1f} min." ) if st.session_state.get("prep_error"): st.error(st.session_state["prep_error"]) if st.session_state.get("prep_message") and not st.session_state.get("prep_error"): st.success(st.session_state["prep_message"]) if st.session_state.get("prep_log"): with st.expander("Prep logs", expanded=False): st.text(st.session_state["prep_log"]) st.subheader("3) Registered/normalized viewer") if st.button("Clear registered results", key="clear_registered_results_btn"): _clear_dir(MODEL_INPUT_DEST) for k in ("prep_log", "prep_message", "prep_error", "prep_ready", "prep_has_mask_out"): st.session_state.pop(k, None) st.rerun() t1_dir = MODEL_INPUT_DEST / "t1" mask_dir = MODEL_INPUT_DEST / "masks" prepared_t1 = sorted(t1_dir.glob("*.nii*")) if t1_dir.exists() else [] if not prepared_t1: st.info( "No prepared files found yet. Run the registration button above to generate " "`data/processed/test_input/t1` and optional masks." ) return case_names = [p.name for p in prepared_t1] selected_name = st.selectbox("Prepared case", case_names, key="prepared_case") selected_t1 = t1_dir / selected_name selected_mask = _paired_mask_path(selected_t1, mask_dir) prepared_data, prepared_zooms = _load_nifti(str(selected_t1)) prepared_mask_data = None if selected_mask is not None and selected_mask.exists(): prepared_mask_data, _ = _load_nifti(str(selected_mask)) _render_overlay_from_arrays( image_data=prepared_data, image_zooms=prepared_zooms, mask_data=prepared_mask_data, key_prefix="prepared_view", title="Prepared registered image", ) st.subheader("4) Side-by-side comparison") if native_data is None: st.info("Upload a native T1 above to enable side-by-side native vs registered comparison.") else: _render_side_by_side_comparison( native_data=native_data, native_mask=native_mask_data, prepared_data=prepared_data, prepared_mask=prepared_mask_data, ) st.subheader("5) Lesion segmentation") render_inference_panel( model_input_root=MODEL_INPUT_DEST, project_root=PROJECT_ROOT, model_dir_default=PROJECT_ROOT / "ARC_ATLAS_Model_V3", output_root_default=MODEL_OUTPUT_DEST, model_cache_root_default=HF_MODEL_CACHE_ROOT, hf_model_repo_default=HF_MODEL_DEFAULT_REPO, hf_model_revision_default=HF_MODEL_DEFAULT_REVISION, hf_model_subdir_default=HF_MODEL_DEFAULT_SUBDIR, ) st.subheader("Prepared paths") st.code( "\n".join( [ f"T1 directory: {t1_dir}", f"Mask directory: {mask_dir} (optional)", f"Manifest: {MODEL_INPUT_DEST / 'manifest.csv'}", f"Selected T1: {selected_t1}", f"Selected mask: {selected_mask if selected_mask else 'None'}", ] ) ) if __name__ == "__main__": render_app()