From bfb5bb867598c69940f05b7b72b5af77fccbec2e Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 15:00:35 +0000 Subject: [PATCH] fix(ml): degrade to statistical features on any extractor failure, implement download-models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit extract_features only caught ImportError, but the ResNet18 path also downloads ImageNet weights from download.pytorch.org on first use — in an air-gapped environment that raises URLError/RuntimeError and crashed the whole ml_analysis step instead of falling back. Any extractor failure now degrades to the statistical features with a logged reason. The 'medcheck download-models' stub is now implemented: it pre-caches the weights so local analysis runs fully offline afterwards, and the README documents the one-time download (the previous wording implied no network access at all). Fixes #137 Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01W4svt5QTs4WUMSy4HwiVt9 --- README.md | 2 +- src/medcheck/main.py | 20 +++- src/medcheck/pipeline/ml_analysis.py | 103 +++++++++++-------- tests/unit/test_cli.py | 22 ++++ tests/unit/test_pipeline/test_ml_analysis.py | 15 +++ 5 files changed, 115 insertions(+), 47 deletions(-) diff --git a/README.md b/README.md index f79cd59..a24b2ec 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ and generate structured, radiology-style reports — from the CLI, a web UI, or - **Plug & Play Docker** — single `docker run` command, no local setup required - **Multiple data sources** — local DICOM folders/ZIPs, easyRadiology portal links, and custom plugins -- **Local ML analysis** — on-device anomaly detection and feature extraction; no API key required +- **Local ML analysis** — on-device anomaly detection and feature extraction; no API key required (one-time model download on first use, or pre-fetch with `medcheck download-models` for offline environments) - **Vision-LLM analysis** — Claude Opus 4.8, GPT-5.5, and Gemini 3.5 Flash (opt-in, consent-gated) - **Privacy by default** — nothing leaves your machine without explicit consent; `--deidentify` pseudonymizes reports - **Clinical context input** — attach symptoms, trauma history, and suspected diagnosis to guide the analysis diff --git a/src/medcheck/main.py b/src/medcheck/main.py index 1f7c5dc..2d25d27 100644 --- a/src/medcheck/main.py +++ b/src/medcheck/main.py @@ -291,6 +291,22 @@ def models() -> None: @app.command() def download_models() -> None: - """Download pretrained ML models for local analysis.""" + """Download and cache the pretrained ML models for local analysis. + + Run this once while online (e.g. before moving to an air-gapped + environment); afterwards local ML analysis needs no network access. + """ console.print("[bold blue]MedCheck - Model Download[/bold blue]") - console.print("[yellow]Model download coming in next release.[/yellow]") + from medcheck.pipeline.ml_analysis import _build_feature_extractor + + try: + console.print("Downloading ResNet18 feature-extractor weights (one-time, ~45 MB)...") + _build_feature_extractor() + except ImportError: + console.print("[red]PyTorch is not installed.[/red] Install the local-models extra first:") + console.print(" uv sync --extra local-models") + raise typer.Exit(code=1) from None + except Exception as exc: + console.print(f"[red]Download failed:[/red] {exc}") + raise typer.Exit(code=1) from None + console.print("[green]Done — model weights are cached; local ML analysis now runs fully offline.[/green]") diff --git a/src/medcheck/pipeline/ml_analysis.py b/src/medcheck/pipeline/ml_analysis.py index d964199..9f929f3 100644 --- a/src/medcheck/pipeline/ml_analysis.py +++ b/src/medcheck/pipeline/ml_analysis.py @@ -77,56 +77,71 @@ def analyze_signal_intensity(volume: np.ndarray) -> SignalStats: ) -def extract_features(volume: np.ndarray) -> np.ndarray: - """Extract features per slice using ResNet18 or fallback to simple stats.""" - try: - import torch - from torchvision import transforms +def _resnet_features(volume: np.ndarray) -> np.ndarray: + """Extract per-slice features with the ResNet18 extractor (needs torch).""" + import torch + from torchvision import transforms - feature_extractor = _get_feature_extractor() + feature_extractor = _get_feature_extractor() - transform = transforms.Compose( - [ - transforms.Resize((224, 224)), - transforms.ToTensor(), - transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), - ] - ) + transform = transforms.Compose( + [ + transforms.Resize((224, 224)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ] + ) - features = [] - with torch.no_grad(): - for i in range(volume.shape[0]): - sl = volume[i] - sl_uint8 = ((sl - sl.min()) / (sl.max() - sl.min() + 1e-8) * 255).astype(np.uint8) - img = Image.fromarray(sl_uint8).convert("RGB") - tensor = transform(img).unsqueeze(0) - feat = feature_extractor(tensor).squeeze().numpy() - features.append(feat) - return np.array(features) - - except ImportError: - # Fallback: simple statistical features when PyTorch not installed - console.print("[yellow]PyTorch not available, using simple statistical features[/yellow]") - features = [] + features = [] + with torch.no_grad(): for i in range(volume.shape[0]): sl = volume[i] - feat = np.array( - [ - sl.mean(), - sl.std(), - sl.min(), - sl.max(), - np.percentile(sl, 25), - np.percentile(sl, 50), - np.percentile(sl, 75), - np.percentile(sl, 95), - np.percentile(sl, 99), - (sl > sl.mean()).mean(), - ], - dtype=np.float32, - ) + sl_uint8 = ((sl - sl.min()) / (sl.max() - sl.min() + 1e-8) * 255).astype(np.uint8) + img = Image.fromarray(sl_uint8).convert("RGB") + tensor = transform(img).unsqueeze(0) + feat = feature_extractor(tensor).squeeze().numpy() features.append(feat) - return np.array(features) + return np.array(features) + + +def _statistical_features(volume: np.ndarray) -> np.ndarray: + """Per-slice statistical features — the no-torch / offline fallback.""" + features = [] + for i in range(volume.shape[0]): + sl = volume[i] + feat = np.array( + [ + sl.mean(), + sl.std(), + sl.min(), + sl.max(), + np.percentile(sl, 25), + np.percentile(sl, 50), + np.percentile(sl, 75), + np.percentile(sl, 95), + np.percentile(sl, 99), + (sl > sl.mean()).mean(), + ], + dtype=np.float32, + ) + features.append(feat) + return np.array(features) + + +def extract_features(volume: np.ndarray) -> np.ndarray: + """Extract features per slice using ResNet18 or fallback to simple stats.""" + try: + return _resnet_features(volume) + except Exception as exc: + # The ResNet path is best-effort: torch may be missing (ImportError), + # or the one-time ImageNet weight download may fail in an air-gapped + # environment (URLError/OSError/RuntimeError). Any of these must + # degrade to the statistical features, not crash the analysis step. + reason = f"{exc.__class__.__name__}: {exc}" if str(exc) else exc.__class__.__name__ + console.print( + f"[yellow]Local ML feature extractor unavailable ({reason}); using simple statistical features[/yellow]" + ) + return _statistical_features(volume) class MLAnalysisStep(PipelineStep): diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index 593d8ab..eec7073 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -97,3 +97,25 @@ def test_cli_models(): assert "claude" in result.stdout assert "openai" in result.stdout assert "gemini" in result.stdout + + +def test_download_models_reports_missing_torch(monkeypatch): + from unittest.mock import patch as mock_patch + + with mock_patch( + "medcheck.pipeline.ml_analysis._build_feature_extractor", + side_effect=ImportError("No module named 'torch'"), + ): + result = runner.invoke(app, ["download-models"]) + assert result.exit_code == 1 + assert "PyTorch is not installed" in result.output + + +def test_download_models_caches_weights(): + from unittest.mock import patch as mock_patch + + with mock_patch("medcheck.pipeline.ml_analysis._build_feature_extractor", return_value=object()) as build: + result = runner.invoke(app, ["download-models"]) + assert result.exit_code == 0 + build.assert_called_once() + assert "cached" in result.output diff --git a/tests/unit/test_pipeline/test_ml_analysis.py b/tests/unit/test_pipeline/test_ml_analysis.py index 3fbb3ee..49d9500 100644 --- a/tests/unit/test_pipeline/test_ml_analysis.py +++ b/tests/unit/test_pipeline/test_ml_analysis.py @@ -90,3 +90,18 @@ def test_ml_step_validate_with_volumes(): def test_ml_step_name(): assert MLAnalysisStep().name == "ml_analysis" + + +def test_extract_features_falls_back_when_resnet_path_fails(): + """A weight-download failure (not just missing torch) must degrade gracefully.""" + volume = np.random.rand(3, 16, 16).astype(np.float32) + with patch.object(ml_analysis, "_resnet_features", side_effect=RuntimeError("download.pytorch.org unreachable")): + features = ml_analysis.extract_features(volume) + assert features.shape == (3, 10) # statistical fallback: 10 features per slice + + +def test_extract_features_falls_back_on_import_error(): + volume = np.random.rand(2, 16, 16).astype(np.float32) + with patch.object(ml_analysis, "_resnet_features", side_effect=ImportError("No module named 'torch'")): + features = ml_analysis.extract_features(volume) + assert features.shape == (2, 10)