Upload models, ONNX exports and application code

.gitattributes

@@ -1,35 +1,11 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.onnx.data filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
 *.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,8 @@
+analyzing_images_for_ai.md
+MODEL_CARD.md
+push_to_hub.py
+.venv/
+blip*/
+data*/
+checkpoints*/
+__pycache__/

batch_predict.py

@@ -0,0 +1,101 @@
+"""
+Batch prediction script for non-interactive use.
+Processes many X-ray images (or a single one) and outputs results as CSV.
+"""
+import argparse
+import csv
+import os
+import sys
+from pathlib import Path
+
+from optimize import (
+    clear_memory,
+    infer_blip,
+    infer_fusion,
+    infer_fusion_onnx,
+    set_cpu_threads,
+)
+
+CHECKPOINT_PATH = "./checkpoints/fusion_model.pth"
+ONNX_FULL_DIR = "./checkpoints/onnx_full"
+HAS_FUSION = os.path.exists(CHECKPOINT_PATH)
+HAS_ONNX = os.path.exists(os.path.join(ONNX_FULL_DIR, "fusion_full.onnx"))
+
+
+def find_images(path):
+    path = Path(path)
+    if path.is_file():
+        return [path]
+    exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".dcm"}
+    return sorted([p for p in path.rglob("*") if p.suffix.lower() in exts])
+
+
+def process_image(image_path, symptoms, use_vision, use_onnx):
+    result = {"image": str(image_path), "symptoms": symptoms}
+
+    if use_vision:
+        try:
+            caption = infer_blip(str(image_path), use_onnx=HAS_ONNX)
+            result["caption"] = caption
+        except Exception as e:
+            result["caption"] = f"ERROR: {e}"
+
+    if HAS_FUSION and symptoms:
+        try:
+            if use_onnx and HAS_ONNX:
+                diagnosis, confidence = infer_fusion_onnx(str(image_path), symptoms)
+            else:
+                diagnosis, confidence = infer_fusion(str(image_path), symptoms)
+            result["diagnosis"] = diagnosis if diagnosis else "INCONCLUSIVE"
+            result["confidence"] = f"{confidence:.4f}" if isinstance(confidence, float) else confidence
+        except Exception as e:
+            result["diagnosis"] = f"ERROR: {e}"
+            result["confidence"] = ""
+
+    clear_memory()
+    return result
+
+
+def main():
+    set_cpu_threads()
+
+    parser = argparse.ArgumentParser(description="Batch X-ray prediction")
+    parser.add_argument("input", help="Path to image file or directory")
+    parser.add_argument("--symptoms", default="What abnormality is present in this chest X-ray?",
+                        help="Symptoms text (same for all images)")
+    parser.add_argument("--output", "-o", default="predictions.csv",
+                        help="Output CSV path")
+    parser.add_argument("--vision", action="store_true",
+                        help="Run vision captioning (BLIP)")
+    parser.add_argument("--onnx", action="store_true",
+                        help="Use ONNX runtime (if exported)")
+    args = parser.parse_args()
+
+    images = find_images(args.input)
+    if not images:
+        print(f"No images found at: {args.input}")
+        sys.exit(1)
+
+    print(f"Found {len(images)} image(s)")
+    print(f"Symptoms: {args.symptoms}")
+    print(f"Vision: {'ON' if args.vision else 'OFF'}")
+    print(f"ONNX: {'ON' if args.onnx and HAS_ONNX else 'OFF'}")
+    print()
+
+    results = []
+    for i, img_path in enumerate(images):
+        print(f"[{i+1}/{len(images)}] Processing {img_path.name}...")
+        result = process_image(str(img_path), args.symptoms, args.vision, args.onnx)
+        results.append(result)
+
+    fieldnames = ["image", "symptoms", "caption", "diagnosis", "confidence"]
+    with open(args.output, "w", newline="", encoding="utf-8") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
+        writer.writeheader()
+        writer.writerows(results)
+
+    print(f"\nDone. Results saved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

blip-xray-finetuned/config.json

@@ -0,0 +1,52 @@
+{
+  "architectures": [
+    "BlipForConditionalGeneration"
+  ],
+  "dtype": "float32",
+  "image_text_hidden_size": 256,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "label_smoothing": 0.0,
+  "logit_scale_init_value": 2.6592,
+  "model_type": "blip",
+  "projection_dim": 512,
+  "text_config": {
+    "attention_probs_dropout_prob": 0.0,
+    "dtype": "float32",
+    "encoder_hidden_size": 768,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.0,
+    "hidden_size": 768,
+    "initializer_factor": 1.0,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "label_smoothing": 0.0,
+    "layer_norm_eps": 1e-12,
+    "max_position_embeddings": 512,
+    "model_type": "blip_text_model",
+    "num_attention_heads": 12,
+    "num_hidden_layers": 12,
+    "projection_dim": 768,
+    "use_cache": true,
+    "vocab_size": 30524
+  },
+  "transformers_version": "4.57.6",
+  "vision_config": {
+    "attention_dropout": 0.0,
+    "dropout": 0.0,
+    "dtype": "float32",
+    "hidden_act": "gelu",
+    "hidden_size": 768,
+    "image_size": 384,
+    "initializer_factor": 1.0,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "layer_norm_eps": 1e-05,
+    "model_type": "blip_vision_model",
+    "num_attention_heads": 12,
+    "num_channels": 3,
+    "num_hidden_layers": 12,
+    "patch_size": 16,
+    "projection_dim": 512
+  }
+}

blip-xray-finetuned/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 30522,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "transformers_version": "4.57.6"
+}

blip-xray-finetuned/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aacb477926134f66b9a4c54273df00aa802c28b056f023f11fd390a43fc19943
+size 895947208

blip-xray-finetuned/preprocessor_config.json

@@ -0,0 +1,24 @@
+{
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "BlipImageProcessor",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "processor_class": "BlipProcessor",
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 384,
+    "width": 384
+  }
+}

blip-xray-finetuned/special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

blip-xray-finetuned/tokenizer_config.json

@@ -0,0 +1,63 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 512,
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "processor_class": "BlipProcessor",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}

build_exe.py

@@ -0,0 +1,89 @@
+"""
+Build a standalone .exe for Windows using PyInstaller.
+No Python installation needed on the target machine.
+
+Usage:
+  python build_exe.py                  # build web_ui.exe
+  python build_exe.py --cli            # build run.exe (CLI)
+  python build_exe.py --all            # build both
+
+Requires: pip install pyinstaller
+"""
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+
+DIST_DIR = "./dist"
+
+
+def build_exe(script, name=None):
+    if name is None:
+        name = os.path.splitext(os.path.basename(script))[0]
+
+    print(f"[cyan]Building {name}.exe from {script}...[/cyan]")
+
+    cmd = [
+        sys.executable, "-m", "PyInstaller",
+        "--onefile",
+        "--console",
+        "--name", name,
+        "--distpath", DIST_DIR,
+        "--workpath", "./build",
+        "--specpath", "./build",
+        "--add-data", "config.json;.",
+        script,
+    ]
+
+    subprocess.check_call(cmd)
+
+    exe_path = os.path.join(DIST_DIR, f"{name}.exe")
+    if os.path.exists(exe_path):
+        size_mb = os.path.getsize(exe_path) / 1024 / 1024
+        print(f"[green]  Created: {exe_path} ({size_mb:.0f} MB)[/green]")
+    else:
+        print(f"[red]  Failed to create {name}.exe[/red]")
+
+    # Clean up build artifacts
+    for d in ["./build", "*.spec"]:
+        try:
+            if os.path.isdir(d):
+                shutil.rmtree(d)
+        except Exception:
+            pass
+    for f in os.listdir("."):
+        if f.endswith(".spec"):
+            os.remove(f)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Build standalone executable")
+    parser.add_argument("--cli", action="store_true", help="Build CLI executable")
+    parser.add_argument("--all", action="store_true", help="Build all executables")
+    args = parser.parse_args()
+
+    try:
+        import PyInstaller  # noqa: F401
+    except ImportError:
+        print("PyInstaller is required. Install with: pip install pyinstaller")
+        sys.exit(1)
+
+    os.makedirs(DIST_DIR, exist_ok=True)
+
+    if args.all:
+        build_exe("web_ui.py")
+        build_exe("run.py")
+        build_exe("batch_predict.py")
+    elif args.cli:
+        build_exe("run.py")
+    else:
+        build_exe("web_ui.py")
+
+    print(f"\n[green]Done. Executables in ./{DIST_DIR}/[/green]")
+    print("[yellow]Note: The .exe still needs model files (checkpoints/).[/yellow]")
+    print("[yellow]Copy the entire project folder to the target machine.[/yellow]")
+
+
+if __name__ == "__main__":
+    main()

capture.py

@@ -0,0 +1,204 @@
+import importlib
+import os
+import subprocess
+import sys
+import tempfile
+from datetime import datetime
+from pathlib import Path
+
+from PIL import Image, ImageOps
+
+DATA_DIR = Path("./data")
+IMAGES_DIR = DATA_DIR / "images"
+
+
+def _ensure_dep(package_name, import_name=None):
+    if import_name is None:
+        import_name = package_name
+    try:
+        return importlib.import_module(import_name)
+    except ImportError:
+        from rich.console import Console
+        console = Console()
+        console.print(f"[yellow]'{package_name}' is required for this feature.[/yellow]")
+        import questionary
+        install = questionary.confirm(f"Install {package_name} now?", default=True).ask()
+        if not install:
+            return None
+        console.print(f"[cyan]Installing {package_name}...[/cyan]")
+        subprocess.check_call([sys.executable, "-m", "pip", "install", package_name])
+        return importlib.import_module(import_name)
+
+def _ensure_dirs():
+    IMAGES_DIR.mkdir(parents=True, exist_ok=True)
+
+def _save_image(pil_image, prefix="capture"):
+    _ensure_dirs()
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
+    filename = f"{prefix}_{timestamp}.jpg"
+    path = str(IMAGES_DIR / filename)
+    if pil_image.mode != "RGB":
+        pil_image = pil_image.convert("RGB")
+    pil_image.save(path, quality=95)
+    return path
+
+# ── Camera ──────────────────────────────────────────────────────
+
+def capture_camera():
+    cv2 = _ensure_dep("opencv-python", "cv2")
+    if cv2 is None:
+        return None, "Camera capture requires opencv-python"
+
+    cap = cv2.VideoCapture(0)
+    if not cap.isOpened():
+        return None, "No camera detected (could not open index 0)"
+
+    from rich.console import Console
+    console = Console()
+    console.print("[cyan]Camera opened. Press SPACE to capture, ESC to cancel.[/cyan]")
+
+    import questionary
+    input("Press Enter when ready for camera preview...")
+    ret, frame = cap.read()
+    cap.release()
+
+    if not ret:
+        return None, "Failed to capture frame from camera"
+
+    preview_path = tempfile.mktemp(suffix="_preview.jpg")
+    cv2.imwrite(preview_path, frame)
+    preview = Image.open(preview_path)
+    os.unlink(preview_path)
+
+    console.print("[cyan]Image captured from camera.[/cyan]")
+    path = _save_image(preview, "camera")
+    return path, f"Captured from camera -> {path}"
+
+# ── File browser ───────────────────────────────────────────────
+
+def capture_file():
+    try:
+        import tkinter as tk
+        from tkinter import filedialog
+        root = tk.Tk()
+        root.withdraw()
+        root.attributes("-topmost", True)
+        path = filedialog.askopenfilename(
+            title="Select an X-ray image",
+            filetypes=[
+                ("Image files", "*.jpg *.jpeg *.png *.bmp *.tif *.tiff *.dcm"),
+                ("All files", "*.*"),
+            ],
+        )
+        root.destroy()
+    except Exception as e:
+        return None, f"File dialog failed: {e}"
+
+    if not path:
+        return None, "No file selected"
+
+    return _open_and_save(path)
+
+# ── Manual path entry ──────────────────────────────────────────
+
+def capture_path():
+    from rich.console import Console
+    console = Console()
+    console.print("[cyan]Enter the path to an X-ray image file.[/cyan]")
+
+    import questionary
+    path = questionary.path("Image path:").ask()
+    if not path:
+        return None, "No path entered"
+    return _open_and_save(path)
+
+# ── DICOM ──────────────────────────────────────────────────────
+
+def capture_dicom(path=None):
+    pydicom = _ensure_dep("pydicom")
+    if pydicom is None:
+        return None, "DICOM loading requires pydicom"
+    np = _ensure_dep("numpy")
+    if np is None:
+        return None, "DICOM loading requires numpy"
+
+    if not path:
+        try:
+            import tkinter as tk
+            from tkinter import filedialog
+            root = tk.Tk()
+            root.withdraw()
+            root.attributes("-topmost", True)
+            path = filedialog.askopenfilename(
+                title="Select a DICOM file",
+                filetypes=[("DICOM files", "*.dcm"), ("All files", "*.*")],
+            )
+            root.destroy()
+        except Exception as e:
+            return None, f"File dialog failed: {e}"
+
+    if not path:
+        return None, "No DICOM file selected"
+
+    try:
+        ds = pydicom.dcmread(path)
+        arr = ds.pixel_array
+        arr = arr - arr.min()
+        arr = (arr / arr.max() * 255).astype(np.uint8)
+        if len(arr.shape) == 2:
+            img = Image.fromarray(arr, mode="L")
+            img = ImageOps.equalize(img)
+        else:
+            img = Image.fromarray(arr)
+        result_path = _save_image(img, "dicom")
+        return result_path, f"DICOM loaded from {path} -> saved as {result_path}"
+    except Exception as e:
+        return None, f"Failed to read DICOM: {e}"
+
+# ── Generic open + save ────────────────────────────────────────
+
+def _open_and_save(source_path):
+    source_path = str(source_path)
+    if source_path.lower().endswith(".dcm"):
+        return capture_dicom(source_path)
+    try:
+        img = Image.open(source_path)
+        path = _save_image(img, "import")
+        return path, f"Imported from {source_path} -> {path}"
+    except Exception as e:
+        return None, f"Failed to open image: {e}"
+
+# ── Top-level picker ───────────────────────────────────────────
+
+def pick_image():
+    from rich.console import Console
+    import questionary
+
+    console = Console()
+    method = questionary.select(
+        "How do you want to provide the X-ray image?",
+        choices=[
+            "Browse files on computer",
+            "Enter file path manually",
+            "Capture from camera",
+            "Load DICOM file",
+        ],
+        pointer=">",
+    ).ask()
+
+    result = None
+    if method == "Browse files on computer":
+        result = capture_file()
+    elif method == "Enter file path manually":
+        result = capture_path()
+    elif method == "Capture from camera":
+        result = capture_camera()
+    elif method == "Load DICOM file":
+        result = capture_dicom()
+
+    return result
+
+
+if __name__ == "__main__":
+    path, msg = pick_image()
+    print(msg)

checkpoints/fusion_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dcca5c883f0dc27e3b119fe79d5d58a029264dac87b0fe9aa4aead941680ae75
+size 5389135

checkpoints/onnx/fusion_classifier.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:217e6206818cb0c72b40076eef6c0f732063be2a462e06ec935587e39305206b
+size 4876

checkpoints/onnx/fusion_classifier.onnx.data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a851183f21595233f4c72cd4d1fbdd0250f9b3f654a5b545be978b9d4c72d9db
+size 4466688

checkpoints/onnx_full/fusion_full.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13c6bad7df618201a2a04e06a9c7f87970c72d33689ec812c1228cfe6aaccffa
+size 787241423

config.json

@@ -0,0 +1,29 @@
+{
+  "_comment": "MedicalAI - Light Weight configuration. Edit this file instead of Python code.",
+
+  "device": "auto",
+  "cpu_threads": "auto",
+  "use_fp16": "auto",
+  "use_onnx": true,
+
+  "data_dir": "./data",
+  "csv_path": "./data/dataset.csv",
+  "images_dir": "./data/images",
+
+  "checkpoint_dir": "./checkpoints",
+  "onnx_dir": "./checkpoints/onnx",
+  "onnx_full_dir": "./checkpoints/onnx_full",
+
+  "blip_model_dir": "./blip-xray-finetuned",
+  "blip_model_name": "Salesforce/blip-image-captioning-base",
+
+  "fusion_checkpoint": "./checkpoints/fusion_model.pth",
+
+  "confidence_threshold": 0.75,
+  "max_symptom_length": 64,
+  "web_port": 7860,
+  "web_host": "127.0.0.1",
+  "api_port": 8000,
+  "api_host": "127.0.0.1",
+  "api_workers": 2
+}

download_model.py

@@ -0,0 +1,138 @@
+"""
+Download pre-trained checkpoints so users can skip training from scratch.
+
+Usage:
+  python download_model.py                     # interactive menu
+  python download_model.py --all               # download everything
+  python download_model.py --fusion            # fusion model only
+  python download_model.py --blip-finetuned    # fine-tuned BLIP only
+"""
+import argparse
+import os
+import sys
+from pathlib import Path
+
+BASE_URL = "https://huggingface.co/your-org/medicalai-lightweight/resolve/main"
+CHECKPOINT_DIR = Path("./checkpoints")
+BLIP_DIR = Path("./blip-xray-finetuned")
+ONNX_DIR = CHECKPOINT_DIR / "onnx_full"
+
+
+def _ensure_dir(path):
+    path.mkdir(parents=True, exist_ok=True)
+
+
+def _download_file(url, dest):
+    import requests
+    from rich.console import Console
+    from rich.progress import Progress, BarColumn, DownloadColumn, TextColumn
+
+    console = Console()
+    console.print(f"[cyan]Downloading {url.split('/')[-1]}...[/cyan]")
+
+    resp = requests.get(url, stream=True)
+    resp.raise_for_status()
+    total = int(resp.headers.get("content-length", 0))
+
+    with Progress(
+        TextColumn("[cyan]  Download[/cyan]"),
+        BarColumn(),
+        DownloadColumn(),
+        transient=True,
+    ) as progress:
+        task = progress.add_task("", total=total)
+        with open(dest, "wb") as f:
+            for chunk in resp.iter_content(chunk_size=8192):
+                f.write(chunk)
+                progress.update(task, advance=len(chunk))
+
+    console.print(f"[green]  Saved to {dest}[/green]")
+
+
+def download_fusion():
+    _ensure_dir(CHECKPOINT_DIR)
+    _ensure_dir(ONNX_DIR)
+
+    files = [
+        ("fusion_model.pth", CHECKPOINT_DIR / "fusion_model.pth"),
+        ("fusion_full.onnx", ONNX_DIR / "fusion_full.onnx"),
+        ("labels.json", ONNX_DIR / "labels.json"),
+    ]
+
+    for fname, dest in files:
+        url = f"{BASE_URL}/{fname}"
+        print(f"  Would download: {url} -> {dest}")
+
+    print()
+    print("[yellow]Note: Pre-trained checkpoints are not yet hosted.[/yellow]")
+    print("[yellow]Train locally with: python training.py --mode prepare-data && python training.py --mode train[/yellow]")
+    print("[yellow]Then export:       python quantization.py --mode export-full[/yellow]")
+
+
+def download_blip():
+    _ensure_dir(BLIP_DIR)
+
+    files = [
+        "config.json",
+        "model.safetensors",
+        "preprocessor_config.json",
+        "special_tokens_map.json",
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "vocab.txt",
+    ]
+
+    for fname in files:
+        url = f"{BASE_URL}/blip-xray-finetuned/{fname}"
+        dest = BLIP_DIR / fname
+        print(f"  Would download: {url} -> {dest}")
+
+    print()
+    print("[yellow]Note: Fine-tuned BLIP is not yet hosted.[/yellow]")
+    print("[yellow]Train locally with: python xray_training.py --mode train[/yellow]")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Download pre-trained models")
+    parser.add_argument("--all", action="store_true", help="Download everything")
+    parser.add_argument("--fusion", action="store_true", help="Download fusion model")
+    parser.add_argument("--blip-finetuned", action="store_true", help="Download fine-tuned BLIP")
+    args = parser.parse_args()
+
+    if not any([args.all, args.fusion, args.blip_finetuned]):
+        from rich.console import Console
+        import questionary
+        console = Console()
+        console.print("[bold cyan]Download Pre-trained Models[/bold cyan]")
+        choice = questionary.select(
+            "What would you like to download?",
+            choices=[
+                "Fusion model (Symptom Check) — ONNX + PyTorch",
+                "Fine-tuned BLIP (Vision captioning)",
+                "Everything",
+                "Cancel",
+            ],
+        ).ask()
+        if choice == "Cancel":
+            return
+        if choice == "Fusion model (Symptom Check) — ONNX + PyTorch":
+            args.fusion = True
+        elif choice == "Fine-tuned BLIP (Vision captioning)":
+            args.blip_finetuned = True
+        else:
+            args.all = True
+
+    try:
+        import requests
+    except ImportError:
+        print("'requests' is required. Install with: pip install requests")
+        sys.exit(1)
+
+    if args.all or args.fusion:
+        download_fusion()
+    if args.all or args.blip_finetuned:
+        download_blip()
+
+
+if __name__ == "__main__":
+    main()

expand_dataset.py

@@ -0,0 +1,790 @@
+import csv
+import os
+import random
+
+CSV_PATH = "./data/dataset.csv"
+
+def count_rows():
+    with open(CSV_PATH, newline="", encoding="utf-8") as f:
+        return sum(1 for _ in csv.DictReader(f))
+
+# Symptom templates (much more varied than just the default one)
+SYMPTOM_TEMPLATES = [
+    "What abnormality is present in this chest X-ray?",
+    "Patient presents with shortness of breath and cough. What is the diagnosis?",
+    "Routine pre-operative chest X-ray. Any abnormal findings?",
+    "Patient with history of smoking. Evaluate for lung pathology.",
+    "Fever and productive cough for 5 days. Assess for pneumonia.",
+    "Chest pain and dyspnea on exertion. Cardiac or pulmonary cause?",
+    "Post-surgical follow-up. Evaluate lung expansion and complications.",
+    "Patient with known COPD. Assess for acute changes or infection.",
+    "Trauma patient. Evaluate for pneumothorax, hemothorax, or fractures.",
+    "Immunocompromised patient with fever. Opportunistic infection?",
+    "Patient with weight loss and night sweats. Evaluate for TB or malignancy.",
+    "Dysphagia and regurgitation. Evaluate for hiatal hernia or mediastinal mass.",
+    "Pre-employment screening chest X-ray.",
+    "Congestive heart failure follow-up. Evaluate for pulmonary edema.",
+    "Patient with known interstitial lung disease. Assess for progression.",
+    "Central cyanosis and clubbing. Evaluate for congenital heart disease.",
+    "Hemoptysis for 2 weeks. Evaluate for bronchiectasis or mass.",
+    "Contact TB patient. Screening chest X-ray.",
+    "Pre-operative clearance for knee replacement surgery.",
+    "Rheumatoid arthritis patient with new dyspnea. Interstitial lung disease?",
+    "HIV positive patient with cough and fever.",
+    "Post-chemotherapy evaluation. Neutropenic fever.",
+    "Patient with asbestos exposure history. Routine surveillance.",
+    "Hoarseness and cough. Evaluate for mediastinal mass or recurrent laryngeal nerve involvement.",
+    "Chest trauma after MVA. Evaluate for aortic injury.",
+    "Patient on chronic steroids. Evaluate for opportunistic infection.",
+    "Pre-renal transplant evaluation chest X-ray.",
+    "Suspected foreign body aspiration in elderly patient.",
+    "Evaluation for pulmonary metastasis in patient with known primary malignancy.",
+    "Post-operative CABG. Evaluate for complications, effusion, or pneumothorax.",
+    "Rule out tuberculosis in patient with positive PPD.",
+    "Dyspnea and orthopnea. Evaluate for heart failure.",
+    "Chest wall deformity. Evaluate spine and thoracic cage.",
+    "Liver cirrhosis patient with dyspnea. Hepatic hydrothorax?",
+    "Pancreatitis patient with respiratory distress.",
+    "ARDS follow-up chest X-ray.",
+    "Ventilator-associated pneumonia surveillance.",
+    "Neonatal respiratory distress. Evaluate for congenital anomalies.",
+    "Suspected pulmonary embolism. Evaluate for Hampton's hump or Westermark sign.",
+    "Chronic cough with mucus production. Bronchiectasis evaluation.",
+]
+
+# List of image paths to reuse (cycling through existing images)
+def get_existing_images():
+    images = []
+    with open(CSV_PATH, newline="", encoding="utf-8") as f:
+        for row in csv.DictReader(f):
+            p = row.get("image_path", "").strip()
+            if p and p not in images:
+                images.append(p)
+    if not images:
+        images = [f"./data\\images\\iu_xray_{i}.jpg" for i in range(1500)]
+    return images
+
+IMAGES = get_existing_images()
+
+# Massive list of FINDINGS + IMPRESSION entries covering common AND obscure conditions
+DIAGNOSIS_ENTRIES = [
+    # ── Normal / No Finding ──
+    {
+        "diagnosis": "FINDINGS: The cardiomediastinal silhouette is within normal limits. The lungs are clear without focal consolidation, pleural effusion, or pneumothorax. The bony thorax is intact. IMPRESSION: Normal chest radiograph. No acute cardiopulmonary abnormality.",
+        "labels": "No Finding",
+    },
+    {
+        "diagnosis": "FINDINGS: Heart size normal. Mediastinal contours normal. Lungs clear bilaterally. Pulmonary vascularity normal. No pleural effusions or pneumothoraces. IMPRESSION: Normal chest examination. No active disease.",
+        "labels": "No Finding",
+    },
+    {
+        "diagnosis": "FINDINGS: Clear lungs with no infiltrates, masses, or nodules. Normal cardiac silhouette. No pleural effusion or pneumothorax. IMPRESSION: No acute cardiopulmonary disease.",
+        "labels": "No Finding",
+    },
+
+    # ── Cardiomegaly / Heart Failure ──
+    {
+        "diagnosis": "FINDINGS: The cardiac silhouette is moderately enlarged. Pulmonary vascularity is increased with cephalization of the upper lobe vessels. Small bilateral pleural effusions. No pneumothorax. IMPRESSION: Cardiomegaly with signs of congestive heart failure and pulmonary vascular congestion.",
+        "labels": "Cardiomegaly|Edema|Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Severe cardiomegaly. Diffuse bilateral interstitial opacities with Kerley B lines in the costophrenic angles. Bilateral pleural effusions, right greater than left. No pneumothorax. IMPRESSION: Severe cardiomegaly with interstitial pulmonary edema and bilateral pleural effusions consistent with congestive heart failure.",
+        "labels": "Cardiomegaly|Edema|Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Moderate cardiomegaly. Prominent pulmonary vasculature with perihilar bat-wing distribution of airspace opacities. Small bilateral pleural effusions. No pneumothorax. IMPRESSION: Acute pulmonary edema superimposed on chronic cardiomegaly.",
+        "labels": "Cardiomegaly|Edema",
+    },
+    {
+        "diagnosis": "FINDINGS: Borderline cardiomegaly. Redistribution of pulmonary blood flow to the upper lobes. No frank pulmonary edema. No pleural effusion. IMPRESSION: Mild cardiomegaly with early pulmonary venous hypertension.",
+        "labels": "Cardiomegaly",
+    },
+    {
+        "diagnosis": "FINDINGS: The cardiac silhouette is severely enlarged with a globular configuration. The pulmonary vascularity is within normal limits. No pleural effusion or pneumothorax. The lungs are clear. IMPRESSION: Severe cardiomegaly, possibly pericardial effusion. Correlation with echocardiogram recommended.",
+        "labels": "Cardiomegaly",
+    },
+    {
+        "diagnosis": "FINDINGS: Moderate cardiomegaly with a prominent left atrial appendage. Double density sign present. Splayed carina. No pulmonary edema. No pleural effusion. IMPRESSION: Cardiomegaly with left atrial enlargement, consider mitral valve disease.",
+        "labels": "Cardiomegaly",
+    },
+
+    # ── Pneumonia (various types) ──
+    {
+        "diagnosis": "FINDINGS: Dense airspace consolidation in the right lower lobe with obscuration of the right hemidiaphragm. Air bronchograms present. No pleural effusion. No pneumothorax. IMPRESSION: Right lower lobe pneumonia.",
+        "labels": "Consolidation|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Patchy airspace opacities in the left upper lobe with ill-defined margins. No cavitation. No effusion. No pneumothorax. IMPRESSION: Left upper lobe pneumonia. Recommend follow-up to document resolution.",
+        "labels": "Consolidation|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Multifocal bilateral patchy and confluent airspace opacities in a peribronchovascular distribution. No definite cavitation. Small bilateral effusions. No pneumothorax. IMPRESSION: Multifocal pneumonia, consider atypical or viral etiology.",
+        "labels": "Infiltration|Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Rounded opacity in the right middle lobe with a positive silhouette sign against the right heart border. Air bronchograms visible. No effusion. IMPRESSION: Right middle lobe pneumonia (silhouette sign present).",
+        "labels": "Consolidation|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Extensive left lower lobe consolidation obscuring the left hemidiaphragm and descending aorta. Air bronchograms are present. Small left pleural effusion. No pneumothorax. IMPRESSION: Left lower lobe pneumonia with parapneumonic effusion.",
+        "labels": "Consolidation|Effusion|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral perihilar interstitial and airspace opacities with central distribution. No pleural effusion. No pneumothorax. IMPRESSION: Interstitial pneumonia, favor atypical/viral etiology.",
+        "labels": "Infiltration|Consolidation",
+    },
+    {
+        "diagnosis": "FINDINGS: Dense consolidation in the right upper lobe with air bronchograms and associated volume loss. Right tracheal shift. No cavitation. No effusion. IMPRESSION: Right upper lobe pneumonia with volume loss. Recommend follow-up to exclude underlying mass.",
+        "labels": "Consolidation|Atelectasis",
+    },
+    {
+        "diagnosis": "FINDINGS: Round pneumonia presenting as a spherical opacity in the left lower lobe. Surrounding ground-glass opacity. No effusion. No pneumothorax. IMPRESSION: Round pneumonia in the left lower lobe. Clinical correlation recommended.",
+        "labels": "Consolidation",
+    },
+    {
+        "diagnosis": "FINDINGS: Cavitary lesion in the right upper lobe with thick irregular wall and surrounding consolidation. Air-fluid level present. No definite pleural effusion. IMPRESSION: Cavitary pneumonia, consider TB, fungal, or necrotizing bacterial infection.",
+        "labels": "Infiltration|Consolidation",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral lower lobe consolidations with air bronchograms. Small bilateral pleural effusions. No pneumothorax. Cardiomegaly noted. IMPRESSION: Bilateral lower lobe pneumonias superimposed on cardiomegaly.",
+        "labels": "Consolidation|Infiltration|Cardiomegaly|Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Segmental airspace opacity in the lingula with obscuration of the left heart border (silhouette sign). No effusion. No pneumothorax. IMPRESSION: Lingular pneumonia.",
+        "labels": "Consolidation|Infiltration",
+    },
+
+    # ── Atelectasis ──
+    {
+        "diagnosis": "FINDINGS: Linear opacity in the right lower lobe extending to the pleura. Minor fissure elevation. No pleural effusion. No pneumothorax. IMPRESSION: Platelike/segmental atelectasis in the right lower lobe.",
+        "labels": "Atelectasis",
+    },
+    {
+        "diagnosis": "FINDINGS: Bandlike opacity in the left base with elevation of the left hemidiaphragm. Compensatory hyperinflation of the remaining lung. No effusion. IMPRESSION: Left basilar subsegmental atelectasis.",
+        "labels": "Atelectasis",
+    },
+    {
+        "diagnosis": "FINDINGS: Triangular opacity in the right base with the apex pointing toward the hilum. Right hemidiaphragm is elevated. No pleural effusion. IMPRESSION: Right lower lobe atelectasis.",
+        "labels": "Atelectasis",
+    },
+    {
+        "diagnosis": "FINDINGS: Golden S sign visible in the right upper lobe with an S-shaped curve of the minor fissure. A central hilar mass is suspected underlying the lobar atelectasis. No effusion. IMPRESSION: Right upper lobe atelectasis. The Golden S sign suggests an underlying central mass. CT chest recommended.",
+        "labels": "Atelectasis|Mass",
+    },
+    {
+        "diagnosis": "FINDINGS: Widespread platelike atelectasis in both lower lobes. Low lung volumes. Hemidiaphragms are elevated. No effusion or pneumothorax. IMPRESSION: Bilateral basilar platelike atelectasis due to hypoventilation.",
+        "labels": "Atelectasis",
+    },
+    {
+        "diagnosis": "FINDINGS: Complete opacification of the left hemithorax with mediastinal shift to the left. Compensatory hyperinflation of the right lung. No pneumothorax. IMPRESSION: Complete left lung atelectasis. Underlying obstructing lesion suspected. Urgent CT and bronchoscopy recommended.",
+        "labels": "Atelectasis|Mass",
+    },
+    {
+        "diagnosis": "FINDINGS: Rounded atelectasis in the right lower lobe presenting as a rounded opacity with comet-tail sign. Pleural thickening adjacent. No change from prior study. IMPRESSION: Rounded atelectasis, stable. No evidence of active disease.",
+        "labels": "Atelectasis|Pleural_Thickening",
+    },
+    {
+        "diagnosis": "FINDINGS: Discoid atelectasis in the right midlung. Minimal volume loss. No pleural effusion. Heart size normal. IMPRESSION: Discoid atelectasis. Otherwise unremarkable chest.",
+        "labels": "Atelectasis",
+    },
+
+    # ── Pleural Effusion ──
+    {
+        "diagnosis": "FINDINGS: Moderate right pleural effusion with blunting of the right costophrenic angle and a meniscus sign. Underlying compressive atelectasis. No pneumothorax. IMPRESSION: Moderate right pleural effusion with adjacent atelectasis.",
+        "labels": "Effusion|Atelectasis",
+    },
+    {
+        "diagnosis": "FINDINGS: Large left pleural effusion causing near-complete opacification of the left hemithorax with mediastinal shift to the right. No pneumothorax. IMPRESSION: Large left pleural effusion with mass effect. Thoracentesis recommended.",
+        "labels": "Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral small pleural effusions with blunting of both costophrenic angles. No layering on lateral decubitus. No pneumothorax. IMPRESSION: Small bilateral pleural effusions.",
+        "labels": "Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Loculated right pleural effusion with biconvex opacity along the lateral chest wall extending into the fissure. No free-flowing component. No pneumothorax. IMPRESSION: Loculated right pleural effusion, consider empyema or hemothorax.",
+        "labels": "Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Massive right pleural effusion with complete opacification of the right hemithorax and contralateral mediastinal shift. Left lung clear. No pneumothorax. IMPRESSION: Massive right pleural effusion. Therapeutic thoracentesis indicated.",
+        "labels": "Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Small right pleural effusion with intact meniscus sign. No pleural thickening. No loculation. Lungs are clear. IMPRESSION: Small right pleural effusion. Clinical correlation for etiology recommended.",
+        "labels": "Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral pleural effusions with associated basal atelectasis. No pneumothorax. Cardiomegaly noted. Pulmonary vascular congestion present. IMPRESSION: Bilateral pleural effusions in the setting of congestive heart failure.",
+        "labels": "Effusion|Cardiomegaly|Atelectasis",
+    },
+
+    # ── Pneumothorax ──
+    {
+        "diagnosis": "FINDINGS: Right apical pneumothorax with visceral pleural line visible approximately 2 cm from the chest wall. No mediastinal shift. The underlying lung is clear. No effusion. IMPRESSION: Small right apical pneumothorax.",
+        "labels": "Pneumothorax",
+    },
+    {
+        "diagnosis": "FINDINGS: Moderate left pneumothorax with the lung edge seen at 3 cm from the lateral chest wall. Mild mediastinal shift to the right. No pleural effusion. IMPRESSION: Moderate left pneumothorax with early tension physiology.",
+        "labels": "Pneumothorax",
+    },
+    {
+        "diagnosis": "FINDINGS: Large right pneumothorax with complete collapse of the right lung (lung appears as a dense hilar mass). Significant mediastinal shift to the left. Depressed right hemidiaphragm. Deep sulcus sign present. IMPRESSION: Large tension pneumothorax on the right. Requires immediate decompression.",
+        "labels": "Pneumothorax",
+    },
+    {
+        "diagnosis": "FINDINGS: Small left apical pneumothorax. No mediastinal shift. No pleural effusion. Remainder of lung is clear. Heart size normal. IMPRESSION: Small spontaneous pneumothorax.",
+        "labels": "Pneumothorax",
+    },
+    {
+        "diagnosis": "FINDINGS: Hydropneumothorax on the left with visible air-fluid level extending across the hemithorax. Partially collapsed left lung. Mediastinal shift to the right. IMPRESSION: Left hydropneumothorax. Consider empyema or bronchopleural fistula.",
+        "labels": "Pneumothorax|Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Tension pneumothorax on the right with deep sulcus sign, mediastinal shift to the left, and flattening of the right heart border. Complete right lung collapse. IMPRESSION: Large right tension pneumothorax, life-threatening. STAT decompression indicated.",
+        "labels": "Pneumothorax",
+    },
+    {
+        "diagnosis": "FINDINGS: Tiny right apical pneumothorax barely visible. Visceral pleural line seen only on inspiratory film. Less than 1 cm from chest wall. No mediastinal shift. IMPRESSION: Very small right apical pneumothorax, likely will resolve spontaneously.",
+        "labels": "Pneumothorax",
+    },
+
+    # ── COPD / Emphysema ──
+    {
+        "diagnosis": "FINDINGS: Hyperinflated lungs with flattened hemidiaphragms. Increased AP chest diameter. Bullous changes at the apices. Heart size is normal. No pneumothorax or effusion. IMPRESSION: Chronic obstructive pulmonary disease with emphysematous changes.",
+        "labels": "Emphysema",
+    },
+    {
+        "diagnosis": "FINDINGS: Severe hyperinflation with flattened and depressed hemidiaphragms. Widened intercostal spaces. Increased retrosternal clear space. Small heart. No focal consolidation. No pneumothorax. IMPRESSION: Severe emphysema. No acute infiltrate.",
+        "labels": "Emphysema",
+    },
+    {
+        "diagnosis": "FINDINGS: Hyperlucent lungs with attenuation of peripheral vascular markings. Large central pulmonary arteries. Flattened hemidiaphragms. No bullae. No pneumothorax. IMPRESSION: Emphysema with pulmonary hypertension signs. Cor pulmonale suspected.",
+        "labels": "Emphysema",
+    },
+    {
+        "diagnosis": "FINDINGS: Severe bilateral bullous emphysema. Large thin-walled bullae occupy more than one-third of both hemithoraces. No pneumothorax. No pleural effusion. Cardiac silhouette is vertically oriented and small. IMPRESSION: Severe bilateral bullous emphysema.",
+        "labels": "Emphysema",
+    },
+    {
+        "diagnosis": "FINDINGS: Moderate hyperinflation of the lungs with flattened hemidiaphragms. Subtle reticular opacities in the lung bases. Heart size normal. No effusion or pneumothorax. IMPRESSION: COPD with mild interstitial changes.",
+        "labels": "Emphysema|Fibrosis",
+    },
+    {
+        "diagnosis": "FINDINGS: Hyperexpanded lungs with increased retrosternal airspace. Mild flattening of the diaphragms. No focal consolidation. Heart size is normal. No pneumothorax. IMPRESSION: Mild hyperinflation, consistent with early COPD.",
+        "labels": "Emphysema",
+    },
+
+    # ── Nodules and Masses ──
+    {
+        "diagnosis": "FINDINGS: Solitary pulmonary nodule in the right upper lobe measuring approximately 1.5 cm. Margins are smooth. No calcification. No associated adenopathy. No pleural effusion. IMPRESSION: Solitary pulmonary nodule. Recommend CT chest for further characterization.",
+        "labels": "Nodule|Mass",
+    },
+    {
+        "diagnosis": "FINDINGS: Spiculated mass in the left upper lobe measuring 3.2 x 2.8 cm. Associated pleural tail. No cavitation. No calcification. No hilar or mediastinal adenopathy. No pleural effusion. IMPRESSION: Spiculated left upper lobe mass suspicious for primary lung malignancy. CT and tissue sampling recommended.",
+        "labels": "Mass|Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Multiple bilateral pulmonary nodules of varying sizes, ranging from 0.5 to 2.0 cm. Random distribution. No cavitation. No calcification. Small right pleural effusion. No pneumothorax. IMPRESSION: Numerous bilateral pulmonary nodules consistent with metastatic disease.",
+        "labels": "Nodule|Mass|Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Cavitating mass in the right upper lobe with thick, irregular walls measuring up to 1.5 cm in thickness. Air-fluid level present. Surrounding ground-glass opacity. No effusion. IMPRESSION: Cavitating lung mass, differential includes primary lung carcinoma, abscess, or fungal infection.",
+        "labels": "Mass|Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Pancoast tumor in the right apex with associated apical cap thickening and destruction of the posterior right first and second ribs. No mediastinal widening. No effusion. IMPRESSION: Right apical mass (Pancoast tumor) with chest wall invasion. Urgent CT and biopsy recommended.",
+        "labels": "Mass",
+    },
+    {
+        "diagnosis": "FINDINGS: Well-defined, smoothly marginated nodule in the left lower lobe with popcorn calcification. No growth compared to prior study. No effusion. No pneumothorax. IMPRESSION: Hamartoma with characteristic popcorn calcification. Benign appearance.",
+        "labels": "Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Subsolid nodule with ground-glass and solid components (part-solid nodule) in the right middle lobe. About 1.2 cm. No calcification. No pleural effusion. IMPRESSION: Part-solid nodule suspicious for adenocarcinoma spectrum. CT follow-up recommended.",
+        "labels": "Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Numerous small, well-defined nodules in a perilymphatic distribution with fissural nodularity and right paratracheal adenopathy. No pleural effusion. No pneumothorax. IMPRESSION: Perilymphatic nodules and mediastinal adenopathy, consider sarcoidosis or lymphangitic carcinomatosis.",
+        "labels": "Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Solitary pulmonary nodule in the right lower lobe measuring 8 mm. Contains central calcification. No spiculation. No adenopathy. No effusion. IMPRESSION: Calcified granuloma. Benign appearance, no further action required.",
+        "labels": "Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Large 6 cm mass in the left lower lobe with irregular borders and central necrosis. Left hilar adenopathy present. No pleural effusion. No pneumothorax. IMPRESSION: Large left lower lobe mass with hilar adenopathy, highly suspicious for malignancy.",
+        "labels": "Mass|Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Anterior mediastinal mass with lobulated contours and no calcification. Trachea is midline. No pleural or pericardial effusion. Lungs are clear. IMPRESSION: Anterior mediastinal mass, differential includes thymoma, lymphoma, or germ cell tumor. CT with contrast recommended.",
+        "labels": "Mass",
+    },
+    {
+        "diagnosis": "FINDINGS: Middle mediastinal mass causing splaying of the carina. No calcification. No pleural effusion. Lungs are clear. IMPRESSION: Subcarinal mediastinal mass/massive lymphadenopathy. CT chest recommended for further evaluation.",
+        "labels": "Mass",
+    },
+
+    # ── Interstitial Lung Disease / Fibrosis ──
+    {
+        "diagnosis": "FINDINGS: Diffuse bilateral reticular opacities with honeycombing in the lung bases. Traction bronchiectasis present. No pneumothorax. No pleural effusion. Normal heart size. IMPRESSION: Usual interstitial pneumonia (UIP) pattern with honeycombing and traction bronchiectasis. Consistent with idiopathic pulmonary fibrosis.",
+        "labels": "Fibrosis",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral ground-glass opacities in the lower lobes with reticular superimposed opacities. Minimal honeycombing. No pleural effusion. No pneumothorax. IMPRESSION: Nonspecific interstitial pneumonia (NSIP) pattern. Clinical correlation recommended.",
+        "labels": "Fibrosis|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Extensive bilateral interstitial opacities with a perilymphatic distribution. Septal thickening and subpleural nodules. Bilateral hilar and right paratracheal lymphadenopathy. No effusion. IMPRESSION: Pulmonary sarcoidosis with bilateral hilar adenopathy and interstitial lung disease (stage II sarcoidosis).",
+        "labels": "Fibrosis|Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral upper lobe predominant fibrotic changes with volume loss, hilar retraction, and architectural distortion. Traction bronchiectasis. No pneumothorax. IMPRESSION: Chronic upper lobe fibrotic changes, consider post-TB sequelae or radiation fibrosis.",
+        "labels": "Fibrosis",
+    },
+    {
+        "diagnosis": "FINDINGS: Diffuse bilateral ground-glass opacities with superimposed reticular opacities and traction bronchiectasis in the lung bases. No honeycombing. No effusion. IMPRESSION: Probable interstitial lung disease. HRCT recommended for further characterization.",
+        "labels": "Fibrosis|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral apical pleural thickening with subpleural fibrotic bands. Upper lobe volume loss with upward hilar retraction. No cavitation. No pneumothorax. IMPRESSION: Chronic apical fibrosis with pleural thickening, likely post-inflammatory.",
+        "labels": "Fibrosis|Pleural_Thickening",
+    },
+    {
+        "diagnosis": "FINDINGS: Crazy-paving pattern with scattered ground-glass opacities superimposed on interlobular septal thickening in the bilateral lower lobes. No pleural effusion. No pneumothorax. IMPRESSION: Crazy-paving pattern. Differential includes alveolar proteinosis, lipoid pneumonia, or cardiogenic pulmonary edema. Clinical correlation recommended.",
+        "labels": "Infiltration|Fibrosis",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral reticulonodular opacities in a mid-to-upper lung distribution. Eggshell calcifications in bilateral hilar lymph nodes. No pleural effusion. IMPRESSION: Silicosis with eggshell calcification of hilar lymph nodes and reticulonodular interstitial lung disease.",
+        "labels": "Fibrosis|Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral pleural plaques with calcification along the diaphragmatic pleura and lateral chest wall. No pleural effusion. Lungs are otherwise clear. No pneumothorax. IMPRESSION: Bilateral pleural plaques in a patient with asbestos exposure history. No evidence of asbestosis or mesothelioma.",
+        "labels": "Pleural_Thickening",
+    },
+    {
+        "diagnosis": "FINDINGS: Diffuse bilateral fine reticulonodular opacities predominantly in the lower lobes. No honeycombing. No pleural effusion. Normal heart size. IMPRESSION: Early interstitial lung disease, consider connective tissue disease-related ILD or hypersensitivity pneumonitis.",
+        "labels": "Fibrosis|Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Lymphangitic carcinomatosis presenting with unilateral right-sided septal thickening and peribronchial cuffing. Right hilar adenopathy. Small right pleural effusion. No pneumothorax. IMPRESSION: Right-sided lymphangitic carcinomatosis, consistent with known primary malignancy.",
+        "labels": "Infiltration|Effusion|Nodule",
+    },
+
+    # ── Pulmonary Edema ──
+    {
+        "diagnosis": "FINDINGS: Bilateral perihilar airspace opacities with a butterfly-wing pattern. Upper lobe pulmonary vascular redistribution. Kerley B lines at the lung bases. Cardiomegaly. Small bilateral pleural effusions. No pneumothorax. IMPRESSION: Acute pulmonary edema due to congestive heart failure.",
+        "labels": "Edema|Cardiomegaly|Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Diffuse bilateral airspace opacities that are more confluent centrally with peripheral sparing. No cardiomegaly. No pleural effusion. Normal heart size. No pneumothorax. IMPRESSION: Noncardiogenic pulmonary edema (ARDS pattern). Clinical correlation needed.",
+        "labels": "Edema|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Mild interstitial pulmonary edema with peribronchial cuffing, indistinct vascular margins, and Kerley A and B lines. Cardiomegaly is mild. No pleural effusion. No pneumothorax. IMPRESSION: Mild interstitial pulmonary edema, early congestive heart failure.",
+        "labels": "Edema|Cardiomegaly",
+    },
+    {
+        "diagnosis": "FINDINGS: Asymmetric left-sided pulmonary edema with a perihilar distribution. Patchy airspace opacities in the left upper and lower lobes. No pleural effusion. Normal heart size. No pneumothorax. IMPRESSION: Asymmetric pulmonary edema. Consider acute mitral regurgitation or localized lung pathology.",
+        "labels": "Edema|Infiltration",
+    },
+
+    # ── Tuberculosis ──
+    {
+        "diagnosis": "FINDINGS: Right apical fibronodular opacities with associated volume loss and upward hilar retraction. Small calcified granuloma in the right apex. No cavitation. No pleural effusion. No pneumothorax. IMPRESSION: Right apical fibronodular changes consistent with prior granulomatous disease, likely old TB. No evidence of active disease.",
+        "labels": "Fibrosis|Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Cavitary lesion in the right upper lobe with surrounding airspace disease. Tree-in-bud opacities in the right upper lobe. Right paratracheal adenopathy. No pleural effusion. IMPRESSION: Right upper lobe cavitary lesion with tree-in-bud opacities suspicious for active pulmonary tuberculosis. Sputum AFB recommended.",
+        "labels": "Infiltration|Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Military pattern with innumerable tiny 1-2 mm nodules diffusely throughout both lungs. No consolidation. No pleural effusion. Heart size normal. No pneumothorax. IMPRESSION: Military nodules, highly suspicious for miliary tuberculosis. Urgent clinical correlation and AFB cultures recommended.",
+        "labels": "Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Left apical pleural thickening with an associated fibrotic band extending to the hilum. Calcified left hilar lymph node. Surrounding parenchyma shows traction bronchiectasis. No effusion. IMPRESSION: Old tuberculosis with apical pleural thickening and fibrotic scarring. Stable appearance.",
+        "labels": "Fibrosis|Pleural_Thickening",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral upper lobe fibrocavitary disease with thick-walled cavities and surrounding retraction. Elevated hila. Compensatory hyperinflation of lower lobes. No pneumothorax. Minimal pleural thickening. IMPRESSION: Chronic fibrocavitary TB with bilateral upper lobe involvement. Active disease cannot be excluded.",
+        "labels": "Fibrosis|Infiltration",
+    },
+
+    # ── Bronchiectasis ──
+    {
+        "diagnosis": "FINDINGS: Dilated, thickened bronchi seen end-on as tram-track opacities in the bilateral lower lobes. Bronchial wall thickening. Some mucus plugging. No consolidation. No pleural effusion. No pneumothorax. IMPRESSION: Bilateral lower lobe bronchiectasis with peribronchial thickening and mucus plugging.",
+        "labels": "Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Severe cystic bronchiectasis in both lower lobes with thin-walled cysts and air-fluid levels. Signet-ring sign present. Scattered tree-in-bud opacities. No pneumothorax. No pleural effusion. IMPRESSION: Cystic bronchiectasis with superinfection. Clinical correlation and sputum culture recommended.",
+        "labels": "Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Tram-track and ring-shadow opacities in the right middle lobe and lingula. Volume loss in the right middle lobe. No pleural effusion. No pneumothorax. IMPRESSION: Right middle lobe and lingular bronchiectasis with volume loss.",
+        "labels": "Infiltration|Atelectasis",
+    },
+    {
+        "diagnosis": "FINDINGS: Central bronchiectasis with dilated, mucus-filled bronchi appearing as glove-finger opacities radiating from the hila. Tree-in-bud opacities in the surrounding parenchyma. No effusion. No pneumothorax. IMPRESSION: Central bronchiectasis with mucoid impaction. Consider allergic bronchopulmonary aspergillosis (ABPA) if asthmatic.",
+        "labels": "Infiltration",
+    },
+
+    # ── Rib Fractures / Trauma ──
+    {
+        "diagnosis": "FINDINGS: Minimally displaced fracture of the right lateral 7th rib. Small associated right pleural effusion. No pneumothorax. Lungs are otherwise clear. IMPRESSION: Right 7th rib fracture with small adjacent pleural effusion.",
+        "labels": "Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Multiple left-sided rib fractures involving ribs 4-8 with a flail segment. Large left hemothorax with nearly complete opacification of the left hemithorax. Mediastinal shift to the right. No pneumothorax. IMPRESSION: Left flail chest with massive hemothorax. Cardiac contusion may be present.",
+        "labels": "Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Fracture of the left clavicle with inferior displacement. No pneumothorax. No hemothorax. Lungs are clear. Heart and mediastinum normal. IMPRESSION: Isolated left clavicular fracture. No thoracic injury.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Multiple bilateral old rib fractures with callus formation and cortical irregularity. No acute fracture line identified. No pneumothorax. No pleural effusion. Lungs clear. IMPRESSION: Multiple old healed rib fractures. No acute traumatic findings.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Superior sulcus opacity on the right with thickened pleura and erosion of the posterior right first rib. Correlate clinically for Pancoast syndrome. No effusion. No pneumothorax. IMPRESSION: Right apical mass with rib destruction concerning for Pancoast tumor.",
+        "labels": "Mass",
+    },
+    {
+        "diagnosis": "FINDINGS: Sternal fracture with mild displacement noted on lateral view. No pneumothorax. No hemothorax. No mediastinal widening. Heart size normal. IMPRESSION: Sternal fracture without underlying cardiac or aortic injury.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Wide mediastinum measuring >8 cm on AP view. Loss of aortic knob contour. Left apical cap. Left pleural effusion. No pneumothorax. IMPRESSION: Traumatic aortic injury suggested by mediastinal widening, apical cap, and left pleural effusion. Urgent CT aortogram and surgical consultation required.",
+        "labels": "Effusion",
+    },
+
+    # ── Hiatal Hernia ──
+    {
+        "diagnosis": "FINDINGS: Large retrocardiac air-fluid level behind the heart. Visible herniated stomach contents above the diaphragm. Mild compressive atelectasis of the left lower lobe. Heart size normal. No pneumothorax. IMPRESSION: Large hiatal hernia with stomach herniated into the chest. No acute obstructive findings.",
+        "labels": "Hernia",
+    },
+    {
+        "diagnosis": "FINDINGS: Moderate hiatal hernia seen on the lateral view as a retrocardiac opacity with air-fluid level. The hernia contains both stomach and colon. No evidence of obstruction. No pleural effusion. IMPRESSION: Large hiatal hernia containing stomach and colon, likely chronic.",
+        "labels": "Hernia",
+    },
+    {
+        "diagnosis": "FINDINGS: Small hiatal hernia noted incidentally. No air-fluid level. No associated atelectasis. Normal cardiomediastinal silhouette. IMPRESSION: Small hiatal hernia, otherwise normal chest.",
+        "labels": "Hernia",
+    },
+    {
+        "diagnosis": "FINDINGS: Large hiatal hernia with a mixed gas and soft tissue density in the retrocardiac region. The herniated stomach shows evidence of volvulus with two air-fluid levels at different heights. No pneumothorax. No pleural effusion. IMPRESSION: Large hiatal hernia with gastric volvulus. Urgent surgical evaluation recommended.",
+        "labels": "Hernia",
+    },
+
+    # ── Pneumoperitoneum / Free Air ──
+    {
+        "diagnosis": "FINDINGS: Free air under the right hemidiaphragm on upright PA view. Air outlining the liver. No pneumothorax. Lungs are clear. Heart size normal. IMPRESSION: Pneumoperitoneum suggesting hollow viscus perforation. Urgent surgical evaluation recommended.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Rigler sign visible with air outlining both sides of the bowel wall. Free air under both hemidiaphragms. Visible falciform ligament. No pneumothorax. IMPRESSION: Massive pneumoperitoneum. Emergency surgical consultation required.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Football sign in the supine AP view with large oval lucency over the upper abdomen. Air outlining the peritoneal cavity. No pneumothorax. IMPRESSION: Large pneumoperitoneum (football sign). Suspect perforated viscus.",
+        "labels": "",
+    },
+
+    # ── Aortic Aneurysm / Dissection ──
+    {
+        "diagnosis": "FINDINGS: Prominent thoracic aortic knob with calcification. Mediastinal width is at the upper limits of normal. No definite widening. Lungs are clear. IMPRESSION: Tortuous calcified aorta without definite aneurysm.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Wide mediastinum with loss of the aortic knob contour. Tracheal deviation to the right. Left pleural effusion. No pneumothorax. Heart size normal. IMPRESSION: Wide mediastinum concerning for aortic dissection. CT aortogram recommended urgently.",
+        "labels": "Effusion",
+    },
+    {
+        "diagnosis": "FINDINGS: Prominent descending thoracic aorta with calcified walls. Measured 4.5 cm in diameter. No dissection flap visible on X-ray. No pleural effusion. Lungs clear. IMPRESSION: Descending thoracic aortic aneurysm measuring 4.5 cm. CT with contrast recommended for further characterization.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Calcified aortic arch aneurysm projecting to the right of the trachea. No definite dissection. No pleural effusion. No pneumothorax. IMPRESSION: Aortic arch aneurysm. Echocardiogram or CT for further evaluation recommended.",
+        "labels": "",
+    },
+
+    # ── Pericardial Effusion ──
+    {
+        "diagnosis": "FINDINGS: Globular enlargement of the cardiac silhouette with a water-bottle configuration. Clear sharp borders. Lungs are clear. No pleural effusion. Pulmonary vascularity is normal. No pneumothorax. IMPRESSION: Large pericardial effusion giving a water-bottle heart appearance. Echocardiogram recommended.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Mild enlargement of the cardiac silhouette. Epicardial fat pad sign visible with a lucent line separating the heart from the pericardium. Lungs clear. No pleural effusion. IMPRESSION: Small to moderate pericardial effusion. Echocardiogram for confirmation.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Massive globular cardiomegaly with clear lungs. No pulmonary vascular congestion. No pleural effusion. No pneumothorax. IMPRESSION: Large pericardial effusion. Rule out pericardial tamponade. Urgent echocardiogram indicated.",
+        "labels": "",
+    },
+
+    # ── Pulmonary Embolism ──
+    {
+        "diagnosis": "FINDINGS: Hampton's hump: wedge-shaped pleural-based opacity in the right costophrenic angle. Small right pleural effusion. No pneumothorax. Heart size normal. IMPRESSION: Wedge-shaped opacity consistent with pulmonary infarct, likely secondary to pulmonary embolism. CT pulmonary angiogram recommended.",
+        "labels": "Effusion|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Westermark sign: oligemia of the left lung with decreased vascular markings. Prominent central pulmonary artery on the left (Fleischner sign). No pleural effusion. No pneumothorax. IMPRESSION: Westermark sign suggesting pulmonary embolism in the left pulmonary artery. Urgent CT pulmonary angiography recommended.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Enlarged right descending pulmonary artery (Palla sign). Right lower lobe opacity consistent with infarct. Small right pleural effusion. No pneumothorax. IMPRESSION: Right pulmonary artery enlargement with adjacent infarct, concerning for pulmonary embolism. CTPA recommended.",
+        "labels": "Effusion|Infiltration",
+    },
+
+    # ── Pediatric / Congenital ──
+    {
+        "diagnosis": "FINDINGS: Thymic sail sign present with a triangular soft tissue density projecting from the right superior mediastinum. Lungs clear. Heart size normal. No pneumothorax. IMPRESSION: Normal thymus with sail sign. No abnormality in this pediatric chest.",
+        "labels": "No Finding",
+    },
+    {
+        "diagnosis": "FINDINGS: Scimitar sign: curved tubular opacity in the right lower lobe coursing toward the right cardiophrenic angle. Right lung is hypoplastic. Mediastinal shift to the right. Heart is dextroposed. IMPRESSION: Scimitar syndrome (hypogenetic right lung syndrome with partial anomalous pulmonary venous return).",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Boot-shaped heart with upturned cardiac apex, prominent right ventricle, and concave pulmonary artery segment (coeur en sabot). Decreased pulmonary vascularity. No pleural effusion. IMPRESSION: Tetralogy of Fallot with classic boot-shaped heart. Surgical evaluation recommended.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Egg-on-side cardiac silhouette with narrow vascular pedicle. Increased pulmonary vascularity. No pleural effusion. No pneumothorax. IMPRESSION: Transposition of the great arteries with egg-on-side heart. Neonatal cardiology evaluation required.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Figure-of-3 sign in the aortic knob with rib notching of the posterior inferior aspects of ribs 3-8 bilaterally. Heart size normal. Lungs clear. No pleural effusion. IMPRESSION: Coarctation of the aorta with classic figure-of-3 sign and rib notching due to collateral circulation.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Dilated azygos vein seen as a comma-shaped density at the right tracheobronchial angle. No mediastinal mass. Lungs are clear. Heart size normal. IMPRESSION: Prominent azygos vein, likely due to azygos continuation of the IVC or congenital anomaly.",
+        "labels": "",
+    },
+
+    # ── Pneumoconiosis ──
+    {
+        "diagnosis": "FINDINGS: Small rounded nodular opacities (p, q, r type) in the upper and mid lung zones bilaterally. Eggshell calcifications of hilar lymph nodes. No pleural plaques. No pneumothorax. No pleural effusion. IMPRESSION: Simple silicosis with eggshell calcification of hilar nodes.",
+        "labels": "Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral upper lobe predominantly large opacities (progressive massive fibrosis) with surrounding emphysema. Hilar retraction and architectural distortion. Eggshell calcification of hilar nodes. No pneumothorax. IMPRESSION: Complicated silicosis with progressive massive fibrosis (PMF).",
+        "labels": "Fibrosis|Mass",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral diaphragmatic and lateral pleural plaques with calcification. No pleural effusion. Lungs are clear. No pneumothorax. Heart size normal. IMPRESSION: Bilateral pleural plaques due to asbestos exposure. No evidence of asbestosis or mesothelioma at this time.",
+        "labels": "Pleural_Thickening",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral lower lobe interstitial fibrosis with subpleural lines and honeycombing. Bilateral calcified pleural plaques. No pleural effusion. No pneumothorax. IMPRESSION: Asbestosis with parenchymal fibrosis and bilateral pleural plaques.",
+        "labels": "Fibrosis|Pleural_Thickening",
+    },
+    {
+        "diagnosis": "FINDINGS: Diffuse bilateral fine nodular opacities in a mid-to-upper lung zone predominance. Minimal hilar adenopathy. No pleural involvement. No pneumothorax. IMPRESSION: Simple coal workers pneumoconiosis (CWP) with upper lobe predominant nodular opacities.",
+        "labels": "Nodule",
+    },
+
+    # ── Rare/Obscure Conditions ──
+    {
+        "diagnosis": "FINDINGS: Extensive bilateral ground-glass opacities with interlobular septal thickening (crazy-paving pattern) in a geographic distribution. No pleural effusion. No pneumothorax. IMPRESSION: Crazy-paving pattern, suspicious for pulmonary alveolar proteinosis. BAL and HRCT recommended.",
+        "labels": "Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral diffuse micronodular opacities with a mid-to-upper lung predominance. Multiple thin-walled cysts of varying sizes. No pneumothorax. No pleural effusion. IMPRESSION: Langerhans cell histiocytosis (LCH) with characteristic cysts and nodules in a smoker. HRCT for confirmation.",
+        "labels": "Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral large thin-walled cysts with normal intervening lung parenchyma. Predominantly lower lobe distribution. No pneumothorax. No pleural effusion. No nodules. IMPRESSION: Lymphangioleiomyomatosis (LAM) with bilateral thin-walled cysts. HRCT recommended. Rule out tuberous sclerosis.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral perihilar and lower zone opacities with consolidation and air bronchograms. Spontaneous pneumothorax noted on the right. No pleural effusion. IMPRESSION: Spontaneous pneumothorax with underlying consolidation. Consider pulmonary contusion or alveolar hemorrhage.",
+        "labels": "Pneumothorax|Consolidation",
+    },
+    {
+        "diagnosis": "FINDINGS: Diffuse bilateral airspace opacities with central predominance and air bronchograms. Normal heart size. No pleural effusion. No pneumothorax. IMPRESSION: Acute respiratory distress syndrome (ARDS) with bilateral diffuse airspace disease.",
+        "labels": "Infiltration|Consolidation",
+    },
+    {
+        "diagnosis": "FINDINGS: Swyer-James syndrome: hyperlucent left lung with diminished vascular markings. Expiratory film shows air trapping on the left. Left lung is smaller than the right. Heart shifted to the left. No pneumothorax. IMPRESSION: Unilateral hyperlucent lung (Swyer-James/MacLeod syndrome), likely post-infectious in childhood.",
+        "labels": "Emphysema",
+    },
+    {
+        "diagnosis": "FINDINGS: Diffuse bilateral ground-glass opacities and consolidation with a peripheral predominance, sparing the costophrenic angles. Reverse halo sign (atoll sign) visible. No pleural effusion. No pneumothorax. IMPRESSION: Organizing pneumonia pattern. Clinical and histopathological correlation recommended.",
+        "labels": "Consolidation|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Multiple cavitary nodules of varying sizes in both lungs with thick walls. Some nodules show feeding vessel signs. No pleural effusion. No pneumothorax. IMPRESSION: Septic pulmonary emboli with cavitation. Sought for source of infection.",
+        "labels": "Nodule|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral basal and peripheral ground-glass opacities with consolidation. Areas of sparing. Normal heart size. Small left pleural effusion. No pneumothorax. IMPRESSION: Acute eosinophilic pneumonia or cryptogenic organizing pneumonia. Clinical history of eosinophilia needed.",
+        "labels": "Consolidation|Effusion|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Extensive bilateral reticulonodular opacities with a basal predominance. Bronchiectasis and architectural distortion. No pneumothorax. No pleural effusion. IMPRESSION: Rheumatoid arthritis-associated interstitial lung disease with a usual interstitial pneumonia (UIP) pattern.",
+        "labels": "Fibrosis|Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Scleroderma lung: bilateral lower lobe ground-glass and reticular opacities with early honeycombing. Dilated esophagus with air-fluid level. No pleural effusion. No pneumothorax. IMPRESSION: Systemic sclerosis-associated interstitial lung disease with NSIP pattern and esophageal dilatation.",
+        "labels": "Fibrosis|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral symmetric lower lobe consolidation with air bronchograms. Kerley B lines and small bilateral pleural effusions present. Heart size normal. No pneumothorax. IMPRESSION: Acute interstitial pneumonia (AIP/Hamman-Rich syndrome). Rapid progression suggests this entity.",
+        "labels": "Consolidation|Effusion|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Multiple subcentimeter nodules with a perilymphatic distribution and bilateral hilar adenopathy. No pleural effusion. No pneumothorax. Lungs otherwise clear. IMPRESSION: Stage I pulmonary sarcoidosis with bilateral hilar lymphadenopathy (1-2-3 sign). No parenchymal involvement.",
+        "labels": "Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral hilar and right paratracheal lymphadenopathy with associated upper lobe reticulonodular opacities. No honeycombing. No pleural effusion. No pneumothorax. IMPRESSION: Stage II pulmonary sarcoidosis with bilateral hilar adenopathy and parenchymal involvement.",
+        "labels": "Nodule|Fibrosis",
+    },
+    {
+        "diagnosis": "FINDINGS: Diffuse bilateral fine nodules with upper lobe predominance. Mild hilar adenopathy. No fibrosis. No pleural effusion. IMPRESSION: Hypersensitivity pneumonitis presenting with subacute stage. Exposure history correlation recommended.",
+        "labels": "Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Extensive bilateral consolidations and ground-glass opacities with air bronchograms and air-fluid levels. No definite pleural effusion. No pneumothorax. IMPRESSION: Diffuse alveolar hemorrhage. Consider vasculitis or coagulopathy. Clinical correlation with hemoptysis status urgent.",
+        "labels": "Consolidation",
+    },
+    {
+        "diagnosis": "FINDINGS: Unilateral right perihilar mass with obstructive pneumonitis and Golden S sign. Right upper lobe volume loss. No pleural effusion. No pneumothorax. IMPRESSION: Central right lung mass with obstructive atelectasis (Golden S sign suspicious for lung carcinoma). Bronchoscopy and CT recommended.",
+        "labels": "Mass|Atelectasis",
+    },
+    {
+        "diagnosis": "FINDINGS: Multiple randomly distributed nodules throughout both lungs with some showing halo sign (ground-glass surrounding nodule). No cavitation. No pleural effusion. No pneumothorax. IMPRESSION: Random nodules with halo sign. Consider fungal infection (aspergillosis, mucormycosis) in immunocompromised patient.",
+        "labels": "Nodule|Infiltration",
+    },
+    {
+        "diagnosis": "FINDINGS: Air crescent sign in a preexisting cavity with a round opacity inside. Right upper lobe cavity with intracavitary mass. No pleural effusion. No pneumothorax. IMPRESSION: Aspergilloma (fungus ball) in a preexisting cavity with air crescent sign.",
+        "labels": "Nodule|Mass",
+    },
+    {
+        "diagnosis": "FINDINGS: Diffuse bilateral fine miliary nodules with a random distribution. Bilateral hilar adenopathy. No pleural effusion. No pneumothorax. IMPRESSION: Miliary tuberculosis versus metastatic disease. Clinical history and sputum studies recommended.",
+        "labels": "Nodule",
+    },
+    {
+        "diagnosis": "FINDINGS: Situs inversus with cardiac apex, aortic knob, and stomach bubble on the right. Dextrocardia with normal cardiac situs. Lungs clear. No pneumothorax. IMPRESSION: Dextrocardia with situs inversus totalis. Consider Kartagener syndrome if bronchiectasis present.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Pectus excavatum with depressed sternum, right-sided cardiac displacement, and increased retrosternal airspace. Heart size normal. Lungs clear. No pneumothorax. No pleural effusion. IMPRESSION: Pectus excavatum deformity with cardiac displacement. Otherwise normal chest.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Absence of the right pectoralis muscle with hyperlucency of the right hemithorax. No mediastinal shift. No lung herniation. Lungs clear. No pneumothorax. IMPRESSION: Poland syndrome with absent right pectoralis muscle and right-sided hyperlucency.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Cervical rib arising from the C7 vertebra bilaterally. No associated thoracic outlet syndrome findings. Lungs clear. Heart size normal. No pneumothorax. IMPRESSION: Bilateral cervical ribs, an incidental finding.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Azygos lobe with a visible azygos fissure forming a tear-drop shaped opacity at the right apex. The azygos vein is seen at the base of the fissure. Lungs clear. No other abnormality. IMPRESSION: Azygos lobe, a normal variant. No pathological significance.",
+        "labels": "No Finding",
+    },
+    {
+        "diagnosis": "FINDINGS: Diffuse interstitial pulmonary calcification with bilateral dense nodular and conglomerate opacities. No pleural effusion. No pneumothorax. Heart size normal. IMPRESSION: Metastatic pulmonary calcification, likely secondary to chronic renal failure or hyperparathyroidism.",
+        "labels": "Nodule|Fibrosis",
+    },
+    {
+        "diagnosis": "FINDINGS: Unilateral right lung hyperlucency with diminished vascular markings. Expiratory views show air trapping. No mediastinal shift. No pneumothorax. IMPRESSION: Swyer-James syndrome (post-infectious obliterative bronchiolitis) causing right lung hyperlucency.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral symmetrical consolidation with air bronchograms in the lower lobes. Rapid progression over 24 hours. No pleural effusion. No pneumothorax. Normal heart size. IMPRESSION: Rapid progression of bilateral airspace disease suspicious for ARDS or diffuse alveolar hemorrhage. Clinical history critical.",
+        "labels": "Consolidation",
+    },
+    {
+        "diagnosis": "FINDINGS: Large bulla in the right upper lobe occupying more than 30% of the hemithorax. Compressed adjacent lung. No pneumothorax. No pleural effusion. Remaining lung is hyperexpanded and shows emphysematous change. IMPRESSION: Giant bulla in the right upper lobe. Consider bullectomy if symptomatic.",
+        "labels": "Emphysema",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral apical pleural thickening with associated fibrotic band. No definite mass. No cavitation. No calcification. No pleural effusion. No pneumothorax. IMPRESSION: Biapical pleural thickening, likely post-inflammatory. If new or progressive, consider Pancoast tumor and CT correlation.",
+        "labels": "Pleural_Thickening|Fibrosis",
+    },
+    {
+        "diagnosis": "FINDINGS: Extensive bilateral nodular and reticular opacities in a perihilar distribution with traction bronchiectasis. No honeycombing. No pleural effusion. No pneumothorax. IMPRESSION: Lymphangitic carcinomatosis. Known history of malignancy correlates with this appearance.",
+        "labels": "Nodule|Fibrosis",
+    },
+    {
+        "diagnosis": "FINDINGS: Air-fluid level in a preexisting cystic lesion in the right lower lobe. Surrounding consolidation. The cyst appears to have an intracavitary mass. No pneumothorax. Small pleural effusion on the right. IMPRESSION: Infected lung cyst with possible intracavitary mass, likely an aspergilloma or lung abscess.",
+        "labels": "Consolidation|Effusion|Mass",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilateral subpleural reticular opacities with basal honeycombing and traction bronchiectasis. Progressive from prior exam. No pneumothorax. No pleural effusion. Normal heart size. IMPRESSION: Idiopathic pulmonary fibrosis with decline and honeycombing progression. Pulmonary function correlation recommended.",
+        "labels": "Fibrosis",
+    },
+    {
+        "diagnosis": "FINDINGS: Fat-fluid level and air-fluid level within a large loculated pleural collection on the left. Pneumothorax is absent. The underlying lung is partially compressed. IMPRESSION: Empyema necessitans or complicated parapneumonic effusion with air-fluid level. Diagnostic thoracentesis recommended.",
+        "labels": "Effusion|Pneumothorax",
+    },
+    {
+        "diagnosis": "FINDINGS: Bilaterally enlarged pulmonary arteries with pruning of peripheral vessels. Right ventricular enlargement. No pleural effusion. No pneumothorax. IMPRESSION: Pulmonary arterial hypertension with enlarged central pulmonary arteries and right ventricular enlargement.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Spontaneous pneumomediastinum with air outlining the mediastinal structures, including the thymus (spinnaker sail sign), heart border, and great vessels. Subcutaneous emphysema in the neck. No pneumothorax. No pleural effusion. IMPRESSION: Pneumomediastinum with subcutaneous emphysema. Likely due to airway rupture or Valsalva.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Extensive subcutaneous emphysema in the chest wall extending to the neck. No pneumothorax. No pneumomediastinum. Multiple chest tubes in place. Lungs are partially aerated. IMPRESSION: Subcutaneous emphysema, likely from chest tube or traumatic air leak.",
+        "labels": "",
+    },
+    {
+        "diagnosis": "FINDINGS: Bochdalek hernia defect on the left with herniated abdominal contents (stomach and bowel) visible in the left hemithorax. Mediastinal shift to the right. No pneumothorax. No pleural effusion. IMPRESSION: Bochdalek hernia containing abdominal viscera. Congenital or acquired diaphragmatic hernia.",
+        "labels": "Hernia",
+    },
+    {
+        "diagnosis": "FINDINGS: Morgagni hernia in the right anterior cardiophrenic angle with omental fat herniating through the foramen of Morgagni. No bowel obstruction. No pleural effusion. No pneumothorax. IMPRESSION: Morgagni hernia containing omental fat. Usually asymptomatic and benign.",
+        "labels": "Hernia",
+    },
+    {
+        "diagnosis": "FINDINGS: Eventration of the right hemidiaphragm with marked elevation. The underlying lung shows compressive atelectasis. No pneumothorax. No pleural effusion. Heart shifted to the left. IMPRESSION: Right hemidiaphragm eventration with adjacent atelectasis. May mimic diaphragmatic rupture.",
+        "labels": "Atelectasis",
+    },
+    {
+        "diagnosis": "FINDINGS: Diaphragmatic rupture on the left with stomach herniated into the left hemithorax. The nasogastric tube is coiled above the diaphragm. Mediastinal shift to the right. No pneumothorax. Left pleural effusion. IMPRESSION: Left diaphragmatic rupture with gastric herniation. Emergency surgical repair needed.",
+        "labels": "Effusion|Hernia",
+    },
+]
+
+# Add more rows with varied symptoms cycling through existing images
+def generate_extended_rows():
+    rows = []
+    img_idx = 0
+    num_images = len(IMAGES)
+
+    symptom_idx = 0
+    for diag in DIAGNOSIS_ENTRIES:
+        img_path = IMAGES[img_idx % num_images]
+        img_idx += 1
+
+        symptom = SYMPTOM_TEMPLATES[symptom_idx % len(SYMPTOM_TEMPLATES)]
+        symptom_idx += 1
+
+        labels = diag["labels"]
+        diagnosis = diag["diagnosis"]
+
+        rows.append([img_path, "augmented", symptom, diagnosis, labels])
+
+    return rows
+
+
+if __name__ == "__main__":
+    initial_count = count_rows()
+    print(f"Current dataset rows: {initial_count}")
+
+    new_rows = generate_extended_rows()
+    print(f"Adding {len(new_rows)} new rows...")
+
+    with open(CSV_PATH, "a", newline="", encoding="utf-8") as f:
+        writer = csv.writer(f)
+        writer.writerows(new_rows)
+
+    final_count = count_rows()
+    print(f"New dataset rows: {final_count}")
+    print(f"Added {final_count - initial_count} rows")
+    print("Done!")

launch.ps1

@@ -0,0 +1,112 @@
+<#
+.SYNOPSIS
+  One-click launcher for MedicalAI - Light Weight
+.DESCRIPTION
+  Checks for Python, sets up venv, installs deps, and launches the app.
+  Double-click this file or run: powershell -File launch.ps1
+#>
+
+$ErrorActionPreference = "Stop"
+$ProjectRoot = Split-Path -Parent $MyInvocation.MyCommand.Path
+Set-Location $ProjectRoot
+
+$Host.UI.RawUI.WindowTitle = "MedicalAI - Light Weight"
+
+# ── Check Python ──
+$python = $null
+foreach ($cmd in @("python", "python3", "py")) {
+    try {
+        $v = & $cmd --version 2>&1
+        if ($v -match "Python 3\.(1[0-9]|[0-9]+)") {
+            $python = $cmd
+            break
+        }
+    } catch {}
+}
+if (-not $python) {
+    Write-Host "Python 3.10+ is required but not found." -ForegroundColor Red
+    Write-Host "Download from: https://www.python.org/downloads/" -ForegroundColor Yellow
+    Write-Host "Make sure to check 'Add Python to PATH' during installation." -ForegroundColor Yellow
+    Read-Host "Press Enter to exit"
+    exit 1
+}
+Write-Host "Using: $(& $python --version)" -ForegroundColor Green
+
+# ── Virtual Environment ──
+$venvPath = Join-Path $ProjectRoot ".venv"
+if (-not (Test-Path $venvPath)) {
+    Write-Host "Creating virtual environment..." -ForegroundColor Cyan
+    & $python -m venv $venvPath
+    if (-not $?) { throw "Failed to create venv" }
+}
+
+# ── Activate ──
+$activate = Join-Path $venvPath "Scripts\Activate.ps1"
+. $activate
+
+# ── Install Dependencies ──
+$reqPath = Join-Path $ProjectRoot "requirements.txt"
+if (Test-Path $reqPath) {
+    Write-Host "Installing dependencies..." -ForegroundColor Cyan
+    pip install -q -r $reqPath 2>&1 | Out-Null
+    if (-not $?) {
+        Write-Host "Retrying with full output..." -ForegroundColor Yellow
+        pip install -r $reqPath
+    }
+}
+
+# ── Check / Generate Default Models ──
+$defaultClassifier = Join-Path $ProjectRoot "models\default\fusion_classifier.onnx"
+$checkpoint = Join-Path $ProjectRoot "checkpoints\fusion_model.pth"
+$onnxFull = Join-Path $ProjectRoot "checkpoints\onnx_full\fusion_full.onnx"
+
+if ((-not (Test-Path $checkpoint)) -and (-not (Test-Path $onnxFull)) -and (-not (Test-Path $defaultClassifier))) {
+    Write-Host "No models found. Generating default models..." -ForegroundColor Yellow
+    python setup_default.py
+    Write-Host "Default models generated. The app will work immediately." -ForegroundColor Green
+}
+
+if (Test-Path $checkpoint) {
+    Write-Host "Trained model found." -ForegroundColor Green
+} elseif (Test-Path $onnxFull) {
+    Write-Host "ONNX pipeline found." -ForegroundColor Green
+} elseif (Test-Path $defaultClassifier) {
+    Write-Host "Default model found. Train a proper model for accurate results." -ForegroundColor Yellow
+}
+
+# ── Ask how to launch ──
+Write-Host ""
+Write-Host "MedicalAI - Light Weight" -ForegroundColor Cyan
+Write-Host "========================" -ForegroundColor Cyan
+Write-Host "1) Web UI (recommended - opens in browser)"
+Write-Host "2) Command-line interface (CLI)"
+Write-Host "3) API Server (for website/app integration)"
+Write-Host ""
+
+$choice = Read-Host "Select (1, 2, or 3)"
+
+if ($choice -eq "2") {
+    Write-Host "Launching CLI..." -ForegroundColor Green
+    python run.py
+} elseif ($choice -eq "3") {
+    # Ensure API deps are installed
+    try {
+        Import-Module python -ErrorAction Stop
+        python -c "import fastapi" 2>$null
+    } catch {
+        Write-Host "Installing API dependencies..." -ForegroundColor Cyan
+        python -m pip install "fastapi[standard]" uvicorn 2>&1 | Out-Null
+    }
+    Write-Host "Launching API server on http://127.0.0.1:8000 ..." -ForegroundColor Green
+    Write-Host "Your website can connect to: http://127.0.0.1:8000/api" -ForegroundColor Yellow
+    python quantization.py --mode serve-api
+} else {
+    Write-Host "Launching web UI..." -ForegroundColor Green
+    python web_ui.py
+}
+
+# ── Keep window open on crash ──
+if (-not $?) {
+    Write-Host "App exited with error. Press Enter to close." -ForegroundColor Red
+    Read-Host
+}

launch.sh

@@ -0,0 +1,85 @@
+#!/usr/bin/env bash
+# One-click launcher for MedicalAI on Linux/Mac
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+cd "$SCRIPT_DIR"
+
+# ── Check Python ──
+PYTHON=""
+for cmd in python3 python; do
+    if command -v "$cmd" &>/dev/null; then
+        ver=$("$cmd" --version 2>&1)
+        if echo "$ver" | grep -qE "Python 3\.(1[0-9]|[0-9]+)"; then
+            PYTHON="$cmd"
+            break
+        fi
+    fi
+done
+
+if [ -z "$PYTHON" ]; then
+    echo "Error: Python 3.10+ is required but not found."
+    echo "Install from: https://www.python.org/downloads/"
+    read -rp "Press Enter to exit"
+    exit 1
+fi
+echo "Using: $($PYTHON --version)"
+
+# ── Virtual Environment ──
+if [ ! -d ".venv" ]; then
+    echo "Creating virtual environment..."
+    $PYTHON -m venv .venv
+fi
+
+source .venv/bin/activate
+
+# ── Install Dependencies ──
+if [ -f "requirements.txt" ]; then
+    echo "Installing dependencies..."
+    pip install -q -r requirements.txt 2>/dev/null || pip install -r requirements.txt
+fi
+
+# ── Check / Generate Default Models ──
+if [ ! -f "checkpoints/fusion_model.pth" ] && \
+   [ ! -f "checkpoints/onnx_full/fusion_full.onnx" ] && \
+   [ ! -f "models/default/fusion_classifier.onnx" ]; then
+    echo "No models found. Generating default models..."
+    python setup_default.py
+    echo "Default models generated. The app will work immediately."
+fi
+
+if [ -f "checkpoints/fusion_model.pth" ]; then
+    echo "Trained model found."
+elif [ -f "checkpoints/onnx_full/fusion_full.onnx" ]; then
+    echo "ONNX pipeline found."
+elif [ -f "models/default/fusion_classifier.onnx" ]; then
+    echo "Default model found. Train for accurate results."
+fi
+
+# ── Launch ──
+echo ""
+echo "MedicalAI - Light Weight"
+echo "========================"
+echo "1) Web UI (recommended - opens in browser)"
+echo "2) Command-line interface (CLI)"
+echo "3) API Server (for website/app integration)"
+echo ""
+read -rp "Select (1, 2, or 3): " choice
+
+if [ "$choice" = "2" ]; then
+    echo "Launching CLI..."
+    python run.py
+elif [ "$choice" = "3" ]; then
+    echo "Launching API server on http://127.0.0.1:8000 ..."
+    echo "Your website can connect to: http://127.0.0.1:8000/api"
+    pip install -q "fastapi[standard]" uvicorn 2>/dev/null || pip install "fastapi[standard]" uvicorn
+    python quantization.py --mode serve-api
+else
+    if python -c "import gradio" 2>/dev/null; then
+        python web_ui.py
+    else
+        echo "Installing gradio..."
+        pip install gradio
+        python web_ui.py
+    fi
+fi

models/default/fusion_classifier.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91e573923c145522cc0c2a166e055439e6022f042dc917f287f41210275245f3
+size 4852

models/default/fusion_classifier.onnx.data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:546b236c5ee8445995e0acf16b566f7f021f538d128678a33f747eda52f7354c
+size 1376256

models/default/labels.json

@@ -0,0 +1 @@
+["No Finding", "Atelectasis", "Cardiomegaly", "Effusion", "Infiltration", "Mass", "Nodule", "Pneumonia", "Pneumothorax", "Consolidation", "Edema", "Emphysema", "Fibrosis", "Pleural_Thickening", "Hernia"]

models/default/symptoms.txt

@@ -0,0 +1,8 @@
+What abnormality is present in this chest X-ray?
+Patient presents with shortness of breath and cough.
+Fever and productive cough for 5 days.
+Chest pain and dyspnea on exertion.
+Routine pre-operative chest X-ray.
+Trauma patient. Evaluate for pneumothorax or fractures.
+Patient with history of smoking. Evaluate for lung pathology.
+Immunocompromised patient with fever.

optimize.py

@@ -0,0 +1,399 @@
+import gc
+import os
+import threading
+
+import psutil
+import torch
+from PIL import Image
+
+MODEL_DIR = "./blip-xray-finetuned"
+CHECKPOINT_DIR = "./checkpoints"
+CHECKPOINT_PATH = os.path.join(CHECKPOINT_DIR, "fusion_model.pth")
+ONNX_DIR = os.path.join(MODEL_DIR, "onnx")
+DEFAULT_MODEL_DIR = os.path.join("models", "default")
+DEFAULT_CLASSIFIER_ONNX = os.path.join(DEFAULT_MODEL_DIR, "fusion_classifier.onnx")
+DEFAULT_LABELS_PATH = os.path.join(DEFAULT_MODEL_DIR, "labels.json")
+ONNX_FULL_DIR = os.path.join(CHECKPOINT_DIR, "onnx_full")
+ONNX_FULL_PATH = os.path.join(ONNX_FULL_DIR, "fusion_full.onnx")
+ONNX_FULL_LABELS = os.path.join(ONNX_FULL_DIR, "labels.json")
+
+NIH_LABELS = [
+    "No Finding", "Atelectasis", "Cardiomegaly", "Effusion", "Infiltration",
+    "Mass", "Nodule", "Pneumonia", "Pneumothorax", "Consolidation",
+    "Edema", "Emphysema", "Fibrosis", "Pleural_Thickening", "Hernia",
+]
+
+_loaded_blip = None
+_loaded_fusion = None
+
+# ── CPU Thread Control ────────────────────────────────────────
+
+def set_cpu_threads(n=None):
+    if n is None:
+        n = max(1, psutil.cpu_count(logical=True) // 2)
+    os.environ["OMP_NUM_THREADS"] = str(n)
+    os.environ["MKL_NUM_THREADS"] = str(n)
+    os.environ["NUMEXPR_NUM_THREADS"] = str(n)
+    torch.set_num_threads(n)
+    return n
+
+# ── Memory ────────────────────────────────────────────────────
+
+def get_memory_usage():
+    proc = psutil.Process()
+    mem = proc.memory_info()
+    return {
+        "rss_mb": mem.rss / 1024 / 1024,
+        "vms_mb": mem.vms / 1024 / 1024,
+    }
+
+def clear_memory():
+    global _loaded_blip, _loaded_fusion
+    _loaded_blip = None
+    _loaded_fusion = None
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    if torch.backends.mps.is_available():
+        torch.mps.empty_cache()
+
+# ── Device ─────────────────────────────────────────────────────
+
+def get_device():
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+def use_fp16():
+    return get_device() in ("cuda", "mps")
+
+# ── BLIP / Vision inference ───────────────────────────────────
+
+def infer_blip(image_path, use_onnx=True):
+    global _loaded_blip
+    image = Image.open(image_path).convert("RGB")
+
+    if use_onnx and os.path.exists(os.path.join(ONNX_DIR, "model.onnx")):
+        return _infer_blip_onnx(image)
+    return _infer_blip_pytorch(image)
+
+def _infer_blip_pytorch(image):
+    global _loaded_blip
+    if _loaded_blip is None:
+        from transformers import BlipProcessor, BlipForConditionalGeneration
+        model_name = MODEL_DIR if os.path.exists(MODEL_DIR) else "Salesforce/blip-image-captioning-base"
+        _loaded_blip = {
+            "processor": BlipProcessor.from_pretrained(model_name),
+            "model": BlipForConditionalGeneration.from_pretrained(model_name).eval(),
+        }
+        if use_fp16() and hasattr(torch, "float16"):
+            try:
+                _loaded_blip["model"] = _loaded_blip["model"].half()
+            except Exception:
+                pass
+
+    p, m = _loaded_blip["processor"], _loaded_blip["model"]
+    inputs = p(images=image, return_tensors="pt")
+    if use_fp16():
+        inputs = {k: v.half() if v.dtype == torch.float32 else v for k, v in inputs.items()}
+    with torch.no_grad():
+        out = m.generate(**inputs, max_length=64)
+    return p.decode(out[0], skip_special_tokens=True)
+
+def _infer_blip_onnx(image):
+    from optimum.onnxruntime import ORTModelForVision2Seq
+    from transformers import BlipProcessor
+    processor = BlipProcessor.from_pretrained(ONNX_DIR)
+    model = ORTModelForVision2Seq.from_pretrained(ONNX_DIR, provider="CPUExecutionProvider")
+    inputs = processor(images=image, return_tensors="np")
+    out = model.generate(**inputs, max_length=64)
+    return processor.decode(out[0], skip_special_tokens=True)
+
+# ── Fusion / Symptom Check inference ──────────────────────────
+
+def get_available_models():
+    """Return a dict describing which models are available."""
+    return {
+        "trained_pytorch": os.path.exists(CHECKPOINT_PATH),
+        "default_classifier": os.path.exists(DEFAULT_CLASSIFIER_ONNX),
+        "onnx_full_pipeline": os.path.exists(ONNX_FULL_PATH),
+    }
+
+
+def _infer_fusion_default(image_path, symptoms):
+    """Fallback: use default ONNX classifier with stock PyTorch encoders."""
+    global _loaded_fusion
+    if _loaded_fusion is None:
+        import json
+        from transformers import CLIPModel, CLIPProcessor, AutoTokenizer, AutoModel
+        from training import DiagnosisFusionModel
+
+        label_list = NIH_LABELS
+        model = DiagnosisFusionModel(num_conditions=len(label_list))
+        # Reset classifier to random weights if no checkpoint
+        if not os.path.exists(CHECKPOINT_PATH):
+            for layer in model.classifier:
+                if hasattr(layer, "reset_parameters"):
+                    layer.reset_parameters()
+        _loaded_fusion = {
+            "model": model.eval(),
+            "label_list": label_list,
+        }
+
+    image = Image.open(image_path).convert("RGB")
+    m = _loaded_fusion["model"]
+    label_list = _loaded_fusion["label_list"]
+
+    with torch.no_grad():
+        logits = m([image], [symptoms])
+        probs = torch.softmax(logits, dim=-1)
+        confidence, predicted = torch.max(probs, dim=-1)
+
+    return label_list[predicted.item()], confidence.item()
+
+
+def infer_fusion(image_path, symptoms):
+    global _loaded_fusion
+
+    # Priority 1: Trained PyTorch model
+    if _loaded_fusion is None and os.path.exists(CHECKPOINT_PATH):
+        from training import DiagnosisFusionModel
+        checkpoint = torch.load(CHECKPOINT_PATH, weights_only=False)
+        label_list = checkpoint.get("label_list", [])
+        if label_list:
+            model = DiagnosisFusionModel(num_conditions=len(label_list))
+            model.classifier.load_state_dict(checkpoint["model_state"])
+            _loaded_fusion = {
+                "model": model.eval(),
+                "label_list": label_list,
+            }
+
+    # Priority 2: Default model (stock encoders + fresh classifier)
+    if _loaded_fusion is None:
+        return _infer_fusion_default(image_path, symptoms)
+
+    image = Image.open(image_path).convert("RGB")
+    m = _loaded_fusion["model"]
+    label_list = _loaded_fusion["label_list"]
+
+    with torch.no_grad():
+        logits = m([image], [symptoms])
+        probs = torch.softmax(logits, dim=-1)
+        confidence, predicted = torch.max(probs, dim=-1)
+
+    return label_list[predicted.item()], confidence.item()
+
+# ── Full ONNX Pipeline Export ─────────────────────────────────
+
+def _ensure_onnx_deps():
+    try:
+        import onnx  # noqa: F401
+        return True
+    except ImportError:
+        from rich.console import Console
+        console = Console()
+        console.print("[yellow]Installing ONNX dependencies...[/yellow]")
+        import subprocess, sys
+        subprocess.check_call([
+            sys.executable, "-m", "pip", "install",
+            "onnx", "onnxruntime", "onnxscript",
+        ])
+        return True
+
+def _load_fusion_model():
+    from training import DiagnosisFusionModel
+    checkpoint = torch.load(CHECKPOINT_PATH, weights_only=False)
+    label_list = checkpoint.get("label_list", [])
+    if not label_list:
+        return None, None
+    model = DiagnosisFusionModel(num_conditions=len(label_list))
+    model.classifier.load_state_dict(checkpoint["model_state"])
+    model.eval()
+    return model, label_list
+
+class _FullFusionONNXWrapper(torch.nn.Module):
+    """Wraps the full fusion pipeline so torch.onnx.export can trace it end-to-end."""
+    def __init__(self, model):
+        super().__init__()
+        self.image_encoder = model.image_encoder.vision_model
+        self.symptom_encoder = model.symptom_encoder
+        self.classifier = model.classifier
+        self.image_proj = model.image_encoder.visual_projection
+
+    def forward(self, pixel_values, input_ids, attention_mask):
+        vision_outputs = self.image_encoder(pixel_values)
+        image_features = self.image_proj(vision_outputs.pooler_output)
+        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+
+        text_outputs = self.symptom_encoder(input_ids, attention_mask=attention_mask)
+        text_features = text_outputs.last_hidden_state.mean(dim=1)
+
+        combined = torch.cat([image_features, text_features], dim=-1)
+        return self.classifier(combined)
+
+def export_full_fusion_onnx(output_dir=None):
+    """Export the entire fusion pipeline (image + text -> logits) to a single ONNX file."""
+    if output_dir is None:
+        output_dir = os.path.join(CHECKPOINT_DIR, "onnx_full")
+    os.makedirs(output_dir, exist_ok=True)
+
+    from rich.console import Console
+    console = Console()
+
+    model, label_list = _load_fusion_model()
+    if model is None:
+        console.print("[red]No model or labels found. Train first.[/red]")
+        return
+
+    _ensure_onnx_deps()
+
+    wrapper = _FullFusionONNXWrapper(model).eval()
+
+    dummy_pixel = torch.randn(1, 3, 224, 224)
+    dummy_ids = torch.randint(0, 100, (1, 64), dtype=torch.long)
+    dummy_mask = torch.ones(1, 64, dtype=torch.long)
+
+    console.print("[cyan]Exporting full fusion pipeline to ONNX...[/cyan]")
+
+    torch.onnx.export(
+        wrapper,
+        (dummy_pixel, dummy_ids, dummy_mask),
+        os.path.join(output_dir, "fusion_full.onnx"),
+        input_names=["pixel_values", "input_ids", "attention_mask"],
+        output_names=["logits"],
+        opset_version=14,
+        dynamic_axes={
+            "input_ids": {0: "batch_size", 1: "seq_len"},
+            "attention_mask": {0: "batch_size", 1: "seq_len"},
+            "pixel_values": {0: "batch_size"},
+            "logits": {0: "batch_size"},
+        },
+        dynamo=False,
+    )
+
+    import json
+    with open(os.path.join(output_dir, "labels.json"), "w") as f:
+        json.dump(label_list, f)
+
+    console.print(f"[green]Full ONNX model saved to {output_dir}/fusion_full.onnx[/green]")
+    console.print(f"[green]Labels saved to {output_dir}/labels.json[/green]")
+    console.print(f"[green]Model has {len(label_list)} output classes.[/green]")
+
+def infer_fusion_onnx(image_path, symptoms, model_dir=None):
+    """Run inference using the full ONNX pipeline. No PyTorch needed beyond preprocessing.
+
+    Searches for models in this order:
+      1. checkpoints/onnx_full/  (trained full pipeline)
+      2. models/default/         (default shipped classifier)
+    """
+    import json
+    import numpy as np
+    import onnxruntime as ort
+    from transformers import CLIPProcessor, AutoTokenizer
+
+    # Find the best available ONNX model + labels
+    candidates = [
+        (model_dir, "fusion_full.onnx", "labels.json"),
+        (ONNX_FULL_DIR, "fusion_full.onnx", "labels.json"),
+        (DEFAULT_MODEL_DIR, "fusion_classifier.onnx", "labels.json"),
+    ]
+
+    onnx_path = None
+    labels_path = None
+    for d, m, l in candidates:
+        if d is None:
+            continue
+        mp = os.path.join(d, m)
+        lp = os.path.join(d, l)
+        if os.path.exists(mp) and os.path.exists(lp):
+            onnx_path = mp
+            labels_path = lp
+            break
+
+    if onnx_path is None:
+        return None, "No ONNX model found. Run 'python setup_default.py' or 'python quantization.py --mode export-full'."
+
+    with open(labels_path) as f:
+        label_list = json.load(f)
+
+    clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+    tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+
+    image = Image.open(image_path).convert("RGB")
+    img_inputs = clip_processor(images=image, return_tensors="np")
+    pixel_values = img_inputs["pixel_values"].astype(np.float32)
+
+    tok_inputs = tokenizer(symptoms, return_tensors="np", padding="max_length", truncation=True, max_length=64)
+    input_ids = tok_inputs["input_ids"].astype(np.int64)
+    attention_mask = tok_inputs["attention_mask"].astype(np.int64)
+
+    session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+    logits = session.run(None, {
+        "pixel_values": pixel_values,
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+    })[0]
+
+    probs = np.exp(logits - logits.max(axis=-1, keepdims=True))
+    probs = probs / probs.sum(axis=-1, keepdims=True)
+    predicted = np.argmax(probs, axis=-1)
+    confidence = float(probs[0, predicted[0]])
+
+    return label_list[predicted[0]], confidence
+
+# ── Legacy Quantization (kept for backward compat) ──────────
+
+def quantize_blip(output_dir=None):
+    from rich.console import Console
+    console = Console()
+    console.print("[yellow]BLIP ONNX quantization requires a newer version of optimum.[/yellow]")
+    console.print("[yellow]Run: pip install --upgrade optimum[/yellow]")
+    console.print("[yellow]The system uses PyTorch automatically until then.[/yellow]")
+
+def quantize_fusion(output_dir=None):
+    if output_dir is None:
+        output_dir = os.path.join(CHECKPOINT_DIR, "onnx")
+    os.makedirs(output_dir, exist_ok=True)
+
+    if not os.path.exists(CHECKPOINT_PATH):
+        print("No fusion checkpoint found. Train first.")
+        return
+
+    from rich.console import Console
+    console = Console()
+
+    try:
+        import onnxscript
+    except ImportError:
+        console.print("[yellow]'onnxscript' is required for ONNX export.[/yellow]")
+        import questionary
+        if questionary.confirm("Install onnxscript now?", default=True).ask():
+            import subprocess, sys
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "onnxscript"])
+        else:
+            console.print("[yellow]Skipped. The system works fine without ONNX export.[/yellow]")
+            return
+
+    console.print("[cyan]The fusion model uses frozen CLIP + BERT encoders.[/cyan]")
+    console.print("[cyan]Exporting the classifier head only to ONNX (encoders stay in PyTorch).[/cyan]")
+
+    from training import DiagnosisFusionModel, load_label_list
+
+    label_list = load_label_list()
+    checkpoint = torch.load(CHECKPOINT_PATH, weights_only=False)
+    model = DiagnosisFusionModel(num_conditions=len(label_list))
+    model.classifier.load_state_dict(checkpoint["model_state"])
+    model.eval()
+
+    dummy = torch.randn(1, 512 + 768)
+    torch.onnx.export(
+        model.classifier,
+        dummy,
+        os.path.join(output_dir, "fusion_classifier.onnx"),
+        input_names=["features"],
+        output_names=["logits"],
+        opset_version=14,
+    )
+    console.print(f"[green]  ONNX classifier saved to {output_dir}[/green]")

push_to_hub.py

@@ -0,0 +1,132 @@
+"""Push this project to the Hugging Face Hub.
+
+A bare `hf upload <repo> .` would push .venv (1.4 GB), 9,688 dataset images and
+__pycache__ to the Hub -- `hf upload` does not read .gitignore. This script runs
+the upload with the right excludes and publishes MODEL_CARD.md as the Hub's
+README.md, leaving the GitHub README.md alone.
+
+    python push_to_hub.py                 # dry run: list what would be uploaded
+    python push_to_hub.py --push          # do it
+    python push_to_hub.py --push --data   # include ./data (~181 MB, 9.7k images)
+    python push_to_hub.py --push --create-pr
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+
+REPO_ID = "GAD-Research-Lab/MedicalAI-Light-Weight"
+ROOT = os.path.dirname(os.path.abspath(__file__))
+
+# `hf upload` matches these with fnmatch against repo-relative POSIX paths.
+EXCLUDE = [
+    ".venv/*",
+    ".git/*",
+    "**/__pycache__/*",
+    "*.pyc",
+    "build/*",
+    "dist/*",
+    "*.spec",
+    "update_config.json",
+    "results.csv",
+    "analyzing_images_for_ai.md",          # local scratch notes (gitignored)
+    "blip-xray-finetuned/xray_blip.pth",   # {'epoch': N} stub, no weights -- misleading on the Hub
+    "MODEL_CARD.md",   # uploaded separately, as README.md
+    "README.md",       # GitHub README; the model card takes its place on the Hub
+]
+DATA_EXCLUDE = ["data/*"]
+
+
+def hf_executable():
+    """Prefer the `hf` next to the running interpreter, so a venv run stays in its venv."""
+    bindir = os.path.dirname(sys.executable)
+    for candidate in (os.path.join(bindir, "hf.exe"), os.path.join(bindir, "hf")):
+        if os.path.exists(candidate):
+            return candidate
+    return "hf"
+
+
+def files_to_upload(exclude):
+    """Mirror hf upload's own filtering so the dry run is accurate."""
+    from huggingface_hub.utils import filter_repo_objects
+
+    paths = []
+    for dirpath, dirnames, filenames in os.walk(ROOT):
+        dirnames[:] = [d for d in dirnames if d not in (".git", ".venv", "__pycache__")]
+        for name in filenames:
+            full = os.path.join(dirpath, name)
+            paths.append(os.path.relpath(full, ROOT).replace(os.sep, "/"))
+    return sorted(filter_repo_objects(paths, allow_patterns=None, ignore_patterns=exclude))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--push", action="store_true", help="Actually upload (default is a dry run)")
+    ap.add_argument("--data", action="store_true", help="Include ./data (~181 MB of images)")
+    ap.add_argument("--create-pr", action="store_true", help="Open a PR instead of committing to main")
+    ap.add_argument("--private", action="store_true", help="Create the repo private if it does not exist")
+    ap.add_argument("--repo-id", default=REPO_ID)
+    ap.add_argument("--message", default="Upload models, ONNX exports and application code")
+    args = ap.parse_args()
+
+    exclude = list(EXCLUDE) + ([] if args.data else DATA_EXCLUDE)
+
+    files = files_to_upload(exclude)
+    total = sum(os.path.getsize(os.path.join(ROOT, f)) for f in files)
+    print(f"{len(files)} files, {total / 1e9:.2f} GB -> {args.repo_id}\n")
+    for f in files:
+        size = os.path.getsize(os.path.join(ROOT, f))
+        print(f"  {size / 1e6:>9.2f} MB  {f}" if size > 1e6 else f"  {'':>9}     {f}")
+
+    if not args.push:
+        print("\nDry run. Re-run with --push to upload.")
+        return 0
+
+    hf = hf_executable()
+    common = ["--repo-type", "model"]
+    if args.create_pr:
+        common.append("--create-pr")
+    if args.private:
+        common.append("--private")
+
+    # 1. Model card first, so the repo is never briefly published without one.
+    print("\n-> Uploading MODEL_CARD.md as README.md")
+    card = subprocess.run(
+        [hf, "upload", args.repo_id, os.path.join(ROOT, "MODEL_CARD.md"), "README.md",
+         "--commit-message", "Add model card", *common],
+        cwd=ROOT,
+    )
+    if card.returncode != 0:
+        print("Model card upload failed; stopping before the bulk upload.", file=sys.stderr)
+        return card.returncode
+
+    # 2. Everything else.
+    print("\n-> Uploading project files")
+    bulk = subprocess.run(
+        [hf, "upload", args.repo_id, ".", "--commit-message", args.message,
+         "--exclude", *exclude, *common],
+        cwd=ROOT,
+    )
+    if bulk.returncode != 0:
+        print("Bulk upload FAILED.", file=sys.stderr)
+        return bulk.returncode
+
+    # Never trust the exit code alone -- confirm against the Hub. A 403 on the LFS
+    # endpoint can still leave small files committed, which looks like success.
+    from huggingface_hub import HfApi
+
+    remote = set(HfApi().list_repo_files(args.repo_id))
+    missing = [f for f in files if f not in remote and f != "MODEL_CARD.md"]
+    if missing:
+        print(f"\n{len(missing)} file(s) did NOT land on the Hub:", file=sys.stderr)
+        for f in missing:
+            print(f"  {f}", file=sys.stderr)
+        return 1
+
+    print(f"\nVerified {len(files)} files: https://huggingface.co/{args.repo_id}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

quantization.py

@@ -0,0 +1,307 @@
+import argparse
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+from optimize import (
+    clear_memory,
+    get_available_models,
+    get_device,
+    get_memory_usage,
+    infer_blip,
+    infer_fusion,
+    infer_fusion_onnx,
+    set_cpu_threads,
+)
+
+DEFAULT_MODEL_DIR = "./models/default"
+CHECKPOINT_PATH = "./checkpoints/fusion_model.pth"
+ONNX_FULL_DIR = "./checkpoints/onnx_full"
+
+def _ensure_default_models():
+    if os.path.exists(os.path.join(DEFAULT_MODEL_DIR, "fusion_classifier.onnx")):
+        return
+    if os.path.exists(CHECKPOINT_PATH):
+        return
+    print("No models found. Generating default models...")
+    subprocess.check_call(
+        [sys.executable, "setup_default.py"],
+        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+    )
+
+def _get_model_source():
+    models = get_available_models()
+    if models["onnx_full_pipeline"]:
+        return "ONNX (full pipeline)"
+    if models["trained_pytorch"]:
+        return "Trained PyTorch"
+    if models["default_classifier"]:
+        return "Default (random weights)"
+    return "NOT AVAILABLE"
+
+# ── FastAPI App (lazy-loaded) ────────────────────────────────
+
+app = None
+_executor = None
+
+def _get_app():
+    global app, _executor
+    if app is not None:
+        return app
+
+    import asyncio
+    import tempfile
+    import time
+    from concurrent.futures import ThreadPoolExecutor
+    from contextlib import asynccontextmanager
+    from typing import Optional
+
+    from fastapi import FastAPI, File, Form, UploadFile, HTTPException
+    from fastapi.middleware.cors import CORSMiddleware
+    from pydantic import BaseModel
+
+    _executor = ThreadPoolExecutor(max_workers=2)
+
+    @asynccontextmanager
+    async def _app_lifespan(fapp: FastAPI):
+        set_cpu_threads()
+        _ensure_default_models()
+        yield
+        clear_memory()
+
+    a = FastAPI(
+        title="MedicalAI - Light Weight API",
+        description="REST API for chest X-ray analysis. Upload images for radiology captions or symptom-based diagnosis.",
+        version="1.0.0",
+        lifespan=_app_lifespan,
+    )
+
+    a.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    class HealthResponse(BaseModel):
+        model_config = {"protected_namespaces": ()}
+        status: str
+        device: str
+        model_source: str
+        memory: dict
+        models: dict
+
+    class VisionResponse(BaseModel):
+        caption: str
+        inference_time_ms: float
+
+    class SymptomResponse(BaseModel):
+        diagnosis: str
+        confidence: float
+        inference_time_ms: float
+
+    @a.get("/health", response_model=HealthResponse)
+    @a.get("/api/health", response_model=HealthResponse)
+    async def health():
+        models = get_available_models()
+        return HealthResponse(
+            status="ok",
+            device=get_device().upper(),
+            model_source=_get_model_source(),
+            memory=get_memory_usage(),
+            models=models,
+        )
+
+    @a.post("/api/vision", response_model=VisionResponse)
+    async def analyze_vision(file: UploadFile = File(...)):
+        ext = Path(file.filename).suffix if file.filename else ".jpg"
+        with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as f:
+            content = await file.read()
+            f.write(content)
+            path = f.name
+        loop = asyncio.get_event_loop()
+        try:
+            t0 = time.perf_counter()
+            caption = await loop.run_in_executor(_executor, infer_blip, path, False)
+            elapsed = (time.perf_counter() - t0) * 1000
+            return VisionResponse(caption=caption, inference_time_ms=round(elapsed, 1))
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+        finally:
+            os.unlink(path)
+            await loop.run_in_executor(None, clear_memory)
+
+    @a.post("/api/symptom-check", response_model=SymptomResponse)
+    async def analyze_symptom(
+        file: UploadFile = File(...),
+        symptoms: str = Form("No symptoms provided"),
+        use_onnx: Optional[bool] = None,
+    ):
+        ext = Path(file.filename).suffix if file.filename else ".jpg"
+        with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as f:
+            content = await file.read()
+            f.write(content)
+            path = f.name
+        loop = asyncio.get_event_loop()
+        try:
+            if use_onnx is None:
+                models = get_available_models()
+                use_onnx = models.get("onnx_full_pipeline", False)
+
+            t0 = time.perf_counter()
+            if use_onnx:
+                diagnosis, confidence = await loop.run_in_executor(
+                    _executor, infer_fusion_onnx, path, symptoms
+                )
+            else:
+                diagnosis, confidence = await loop.run_in_executor(
+                    _executor, infer_fusion, path, symptoms
+                )
+            elapsed = (time.perf_counter() - t0) * 1000
+
+            if diagnosis is None:
+                raise HTTPException(status_code=500, detail=confidence)
+
+            return SymptomResponse(
+                diagnosis=diagnosis,
+                confidence=round(confidence, 4),
+                inference_time_ms=round(elapsed, 1),
+            )
+        except HTTPException:
+            raise
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+        finally:
+            os.unlink(path)
+            await loop.run_in_executor(None, clear_memory)
+
+    @a.get("/api/models")
+    async def list_models():
+        return get_available_models()
+
+    app = a
+    return app
+
+# ── Original CLI Modes ──────────────────────────────────────
+
+def quantize_fusion():
+    from optimize import quantize_fusion as _do
+    _do()
+
+def export_full():
+    from optimize import export_full_fusion_onnx
+    export_full_fusion_onnx()
+
+def optimize_all():
+    from rich.console import Console
+    console = Console()
+    console.print("[bold cyan]Full Optimization Pipeline[/bold cyan]")
+    console.print()
+    console.print("[cyan]Quantizing fusion model...[/cyan]")
+    from optimize import quantize_fusion as qf
+    qf()
+    console.print()
+    console.print("[cyan]Exporting full ONNX pipeline...[/cyan]")
+    from optimize import export_full_fusion_onnx as ef
+    ef()
+    console.print()
+    console.print("[green]Optimization complete![/green]")
+    console.print("  Fusion classifier: ./checkpoints/onnx/fusion_classifier.onnx")
+    console.print("  Full pipeline:     ./checkpoints/onnx_full/fusion_full.onnx")
+
+def set_threads():
+    from optimize import set_cpu_threads as sct, get_memory_usage as gmu
+    n = sct()
+    mem = gmu()
+    print(f"CPU threads set to {n}")
+    print(f"Current RAM: {mem['rss_mb']:.0f} MB")
+
+def show_status():
+    import torch
+    mem = get_memory_usage()
+    print(f"Device:        {get_device().upper()}")
+    print(f"FP16 mode:     {'ON' if get_device() in ('cuda', 'mps') else 'OFF'}")
+    print(f"Process RAM:   {mem['rss_mb']:.0f} MB")
+    print(f"Torch threads: {torch.get_num_threads()}")
+    fusion_onnx = os.path.exists("./checkpoints/onnx/fusion_classifier.onnx")
+    fusion_pt = os.path.exists(CHECKPOINT_PATH)
+    fusion_full = os.path.exists(ONNX_FULL_DIR + "/fusion_full.onnx")
+    print(f"Fusion ONNX (classifier): {'yes' if fusion_onnx else 'no'}")
+    print(f"Fusion ONNX (full):       {'yes' if fusion_full else 'no'}")
+    print(f"Fusion .pth:              {'yes' if fusion_pt else 'no'}")
+    print(f"BLIP model:    {'fine-tuned' if os.path.exists('./blip-xray-finetuned') else 'stock (no fine-tune)'}")
+
+def explain():
+    from rich.console import Console
+    from rich.table import Table
+    from rich.panel import Panel
+    console = Console()
+    console.print(Panel.fit("[bold cyan]Model Optimization - How It Works[/bold cyan]"))
+    console.print()
+    console.print("[bold]Deployment options:[/bold]")
+    console.print()
+    t = Table(title="Deployment Options")
+    t.add_column("Option", style="cyan", width=22)
+    t.add_column("What it is", style="white")
+    t.add_column("Best for", style="green")
+    t.add_row("PyTorch (default)", "Full model in PyTorch.", "Development, GPU users")
+    t.add_row("ONNX Classifier", "Exports classifier head only.", "Minor CPU speedup")
+    t.add_row("ONNX Full Pipeline", "Exports entire pipeline to ONNX.", "Production, no-PyTorch")
+    console.print(t)
+    console.print()
+    t2 = Table(title="Fusion Model Components")
+    t2.add_column("Part", style="cyan", width=18)
+    t2.add_column("Role", style="white")
+    t2.add_column("Size", style="green")
+    t2.add_row("CLIP encoder", "Image -> 512 features", "~600 MB")
+    t2.add_row("Bio_ClinicalBERT", "Symptoms -> 768 features", "~400 MB")
+    t2.add_row("Classifier head", "1280 -> 256 -> N classes", "~0.5 MB")
+    console.print(t2)
+
+# ── API Server Mode ─────────────────────────────────────────
+
+def serve_api(host="127.0.0.1", port=8000, reload=False):
+    global app
+    _get_app()
+    import uvicorn
+    print(f"MedicalAI API Server starting on http://{host}:{port}")
+    print(f"Device: {get_device().upper()}")
+    print(f"Docs:   http://{host}:{port}/docs")
+    print(f"Health: http://{host}:{port}/health")
+    print()
+    print("Your website can connect to this API using the endpoints above.")
+    uvicorn.run("quantization:app", host=host, port=port, reload=reload)
+
+# ── Main ─────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MedicalAI - Export & Web API server")
+    parser.add_argument("--mode",
+                        choices=[
+                            "quantize-fusion", "export-full", "optimize-all",
+                            "set-threads", "status", "explain",
+                            "serve-api",
+                        ],
+                        default="status")
+    parser.add_argument("--host", default="127.0.0.1", help="Host for serve-api (default: 127.0.0.1)")
+    parser.add_argument("--port", type=int, default=8000, help="Port for serve-api (default: 8000)")
+    parser.add_argument("--reload", action="store_true", help="Auto-reload for development")
+    args = parser.parse_args()
+
+    if args.mode == "quantize-fusion":
+        quantize_fusion()
+    elif args.mode == "export-full":
+        export_full()
+    elif args.mode == "optimize-all":
+        optimize_all()
+    elif args.mode == "set-threads":
+        set_threads()
+    elif args.mode == "explain":
+        explain()
+    elif args.mode == "serve-api":
+        serve_api(host=args.host, port=args.port, reload=args.reload)
+    else:
+        show_status()

requirements.txt

@@ -0,0 +1,18 @@
+torch>=2.0.0
+transformers>=4.36.0
+datasets>=2.14.0
+pillow>=10.0.0
+rich>=13.0.0
+questionary>=2.0.0
+psutil>=5.9.0
+tqdm>=4.65.0
+optimum>=1.12.0
+onnxscript>=0.7.0
+
+# Optional (auto-installed on first use):
+# opencv-python  -- camera capture
+# pydicom        -- DICOM file loading
+# nltk           -- BLEU evaluation
+# gradio         -- Web UI (pip install gradio)
+# fastapi        -- REST API server (pip install "fastapi[standard]")
+# uvicorn        -- ASGI server for the API

run.py

@@ -0,0 +1,153 @@
+import os
+
+import questionary
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+
+import capture
+from optimize import (
+    clear_memory,
+    get_device,
+    get_memory_usage,
+    infer_blip,
+    infer_fusion,
+    set_cpu_threads,
+)
+
+console = Console()
+
+def run_vision(image_path):
+    console.print("[cyan]Running Vision analysis...[/cyan]")
+    try:
+        caption = infer_blip(image_path)
+    except Exception as e:
+        console.print(f"[red]Error: {e}[/red]")
+        return
+
+    console.print()
+    console.print(Panel(f"[bold green]{caption}[/bold green]", title="Generated Radiology Caption"))
+    return caption
+
+def run_symptom_check(image_path):
+    symptoms = questionary.text("Enter patient symptoms / clinical indication:").ask()
+    if not symptoms:
+        symptoms = "No symptoms provided"
+
+    try:
+        diagnosis, confidence = infer_fusion(image_path, symptoms)
+    except Exception as e:
+        console.print(f"[red]Error: {e}[/red]")
+        return
+
+    if diagnosis is None:
+        console.print(f"[red]{confidence}[/red]")
+        return
+
+    table = Table(title="Diagnosis Result")
+    table.add_column("Prediction", style="cyan")
+    table.add_column("Confidence", style="green")
+    table.add_row(diagnosis, f"{confidence:.1%}")
+    console.print(table)
+
+    if confidence < 0.75:
+        console.print("[yellow]Low confidence — consider follow-up.[/yellow]")
+
+    return diagnosis, confidence
+
+def show_status():
+    from training import info as training_info
+
+    console.print("[bold cyan]--- System Status ---[/bold cyan]")
+    mem = get_memory_usage()
+    console.print(f"Device:        [green]{get_device().upper()}[/green]")
+    console.print(f"Process RAM:   {mem['rss_mb']:.0f} MB")
+    console.print()
+
+    console.print("[bold cyan]--- Fusion Model Data ---[/bold cyan]")
+    training_info()
+    console.print()
+
+    cap_dir = capture.IMAGES_DIR
+    count = len(list(cap_dir.glob("*"))) if cap_dir.exists() else 0
+    console.print(f"Captured images: {count} in {cap_dir}")
+
+def install_deps():
+    deps = []
+    try:
+        import cv2
+    except ImportError:
+        deps.append("opencv-python")
+    try:
+        import pydicom
+    except ImportError:
+        deps.append("pydicom")
+    try:
+        import nltk
+    except ImportError:
+        deps.append("nltk")
+
+    if not deps:
+        console.print("[green]All optional dependencies are already installed.[/green]")
+        return
+
+    import subprocess
+    import sys
+    console.print(f"[yellow]Installing: {' '.join(deps)}[/yellow]")
+    subprocess.check_call([sys.executable, "-m", "pip", "install", *deps])
+    console.print("[green]Done.[/green]")
+
+def main():
+    set_cpu_threads()
+
+    console.print(Panel.fit("[bold cyan]MedicalAI - Light Weight[/bold cyan]"))
+    console.print()
+
+    while True:
+        choice = questionary.select(
+            "What would you like to do?",
+            choices=[
+                "Vision — Generate report from X-ray",
+                "Symptom Check — Diagnose from X-ray + symptoms",
+                "System Status & Data Info",
+                "Install optional deps (camera, DICOM, BLEU)",
+                "Exit",
+            ],
+            pointer=">",
+        ).ask()
+
+        if choice == "Exit":
+            console.print("[bold red]Exiting...[/bold red]")
+            clear_memory()
+            break
+
+        if choice == "Install optional deps (camera, DICOM, BLEU)":
+            install_deps()
+            continue
+
+        if choice == "System Status & Data Info":
+            show_status()
+            console.print()
+            continue
+
+        image_path, msg = capture.pick_image()
+        if image_path is None:
+            console.print(f"[red]{msg}[/red]")
+            continue
+        console.print(f"[dim]{msg}[/dim]")
+
+        if choice.startswith("Vision"):
+            run_vision(image_path)
+        elif choice.startswith("Symptom Check"):
+            run_symptom_check(image_path)
+
+        clear_memory()
+        console.print()
+        again = questionary.confirm("Do another?").ask()
+        if not again:
+            break
+
+    clear_memory()
+
+if __name__ == "__main__":
+    main()

setup_default.py

@@ -0,0 +1,182 @@
+"""
+Generate default ONNX models so the app works immediately after install.
+
+Run: python setup_default.py
+
+Creates:
+  - models/default/labels.json          (15 standard NIH classes)
+  - models/default/fusion_classifier.onnx (small MLP with random weights)
+  - models/default/fusion_full.onnx      (if full pipeline export possible)
+"""
+import json
+import os
+import shutil
+
+DEFAULT_DIR = os.path.join("models", "default")
+NIH_LABELS = [
+    "No Finding", "Atelectasis", "Cardiomegaly", "Effusion", "Infiltration",
+    "Mass", "Nodule", "Pneumonia", "Pneumothorax", "Consolidation",
+    "Edema", "Emphysema", "Fibrosis", "Pleural_Thickening", "Hernia",
+]
+
+
+def _ensure_dir(path):
+    os.makedirs(path, exist_ok=True)
+
+
+def _export_classifier_onnx():
+    """Export a minimal ONNX classifier head (random weights, correct shape)."""
+    import torch
+    import torch.nn as nn
+
+    classifier = nn.Sequential(
+        nn.Linear(512 + 768, 256),
+        nn.ReLU(),
+        nn.Dropout(0.2),
+        nn.Linear(256, len(NIH_LABELS)),
+    )
+    classifier.eval()
+
+    dummy = torch.randn(1, 512 + 768)
+    onnx_path = os.path.join(DEFAULT_DIR, "fusion_classifier.onnx")
+    torch.onnx.export(
+        classifier,
+        dummy,
+        onnx_path,
+        input_names=["features"],
+        output_names=["logits"],
+        opset_version=14,
+    )
+    return onnx_path
+
+
+def _export_full_onnx():
+    """Try to export a full-pipeline ONNX using stock CLIP + BERT encoders.
+
+    This requires transformers to be installed. The full model is large
+    (~1 GB) but runs standalone with ONNX Runtime (no torch).
+    """
+    try:
+        import torch
+        import torch.nn as nn
+        from transformers import CLIPModel, AutoModel
+
+        class _MinimalPipeline(nn.Module):
+            def __init__(self, num_classes):
+                super().__init__()
+                clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+                bert = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+                self.vision_encoder = clip.vision_model
+                self.visual_projection = clip.visual_projection
+                self.text_encoder = bert
+                self.text_pooler = bert.pooler
+                self.classifier = nn.Sequential(
+                    nn.Linear(512 + 768, 256),
+                    nn.ReLU(),
+                    nn.Dropout(0.2),
+                    nn.Linear(256, num_classes),
+                )
+                self.classifier.eval()
+
+            def forward(self, pixel_values, input_ids, attention_mask):
+                v_out = self.vision_encoder(pixel_values)
+                v_feat = self.visual_projection(v_out.pooler_output)
+                v_feat = v_feat / v_feat.norm(dim=-1, keepdim=True)
+
+                t_out = self.text_encoder(input_ids, attention_mask=attention_mask)
+                t_feat = self.text_pooler(t_out.last_hidden_state[:, 0, :])
+
+                combined = torch.cat([v_feat, t_feat], dim=-1)
+                return self.classifier(combined)
+
+        model = _MinimalPipeline(len(NIH_LABELS)).eval()
+        dummy_pixel = torch.randn(1, 3, 224, 224)
+        dummy_ids = torch.randint(0, 100, (1, 64), dtype=torch.long)
+        dummy_mask = torch.ones(1, 64, dtype=torch.long)
+
+        onnx_path = os.path.join(DEFAULT_DIR, "fusion_full.onnx")
+        torch.onnx.export(
+            model,
+            (dummy_pixel, dummy_ids, dummy_mask),
+            onnx_path,
+            input_names=["pixel_values", "input_ids", "attention_mask"],
+            output_names=["logits"],
+            opset_version=14,
+            dynamic_axes={
+                "input_ids": {0: "batch_size", 1: "seq_len"},
+                "attention_mask": {0: "batch_size", 1: "seq_len"},
+                "pixel_values": {0: "batch_size"},
+                "logits": {0: "batch_size"},
+            },
+        )
+        return onnx_path
+    except Exception as e:
+        return None
+
+
+def _save_labels():
+    path = os.path.join(DEFAULT_DIR, "labels.json")
+    with open(path, "w") as f:
+        json.dump(NIH_LABELS, f)
+    return path
+
+
+def _save_symptoms():
+    """Save default symptom prompts for the Vision model."""
+    path = os.path.join(DEFAULT_DIR, "symptoms.txt")
+    templates = [
+        "What abnormality is present in this chest X-ray?",
+        "Patient presents with shortness of breath and cough.",
+        "Fever and productive cough for 5 days.",
+        "Chest pain and dyspnea on exertion.",
+        "Routine pre-operative chest X-ray.",
+        "Trauma patient. Evaluate for pneumothorax or fractures.",
+        "Patient with history of smoking. Evaluate for lung pathology.",
+        "Immunocompromised patient with fever.",
+    ]
+    with open(path, "w") as f:
+        f.write("\n".join(templates))
+    return path
+
+
+def main():
+    from rich.console import Console
+    console = Console()
+
+    _ensure_dir(DEFAULT_DIR)
+
+    # Remove old default files
+    for f in os.listdir(DEFAULT_DIR):
+        fp = os.path.join(DEFAULT_DIR, f)
+        try:
+            if os.path.isfile(fp):
+                os.remove(fp)
+        except Exception:
+            pass
+
+    console.print("[bold cyan]Generating default models...[/bold cyan]")
+
+    labels_path = _save_labels()
+    console.print(f"  [green]{labels_path}[/green]")
+
+    symp_path = _save_symptoms()
+    console.print(f"  [green]{symp_path}[/green]")
+
+    cls_path = _export_classifier_onnx()
+    size = os.path.getsize(cls_path) / 1024
+    console.print(f"  [green]{cls_path}[/green] ({size:.0f} KB)")
+
+    full_path = _export_full_onnx()
+    if full_path:
+        size = os.path.getsize(full_path) / 1024 / 1024
+        console.print(f"  [green]{full_path}[/green] ({size:.0f} MB)")
+    else:
+        console.print("  [yellow]Full pipeline ONNX export skipped (no transformers or CUDA)[/yellow]")
+        console.print("  [yellow]  Run 'python quantization.py --mode export-full' after training for this.[/yellow]")
+
+    console.print()
+    console.print("[green]Default models ready. The app will use these until you train a proper model.[/green]")
+
+
+if __name__ == "__main__":
+    main()

training.py

@@ -0,0 +1,349 @@
+import argparse
+import csv
+import os
+import random
+
+import torch
+import torch.nn as nn
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader, random_split
+
+DATA_DIR = "./data"
+CSV_PATH = os.path.join(DATA_DIR, "dataset.csv")
+IMAGES_DIR = os.path.join(DATA_DIR, "images")
+CHECKPOINT_DIR = "./checkpoints"
+CHECKPOINT_PATH = os.path.join(CHECKPOINT_DIR, "fusion_model.pth")
+CONFIDENCE_THRESHOLD = 0.75
+
+CSV_COLUMNS = ["image_path", "source", "symptoms", "diagnosis", "labels"]
+
+def load_label_list():
+    if not os.path.exists(CSV_PATH):
+        return []
+    labels = set()
+    with open(CSV_PATH, newline="", encoding="utf-8") as f:
+        for row in csv.DictReader(f):
+            d = row.get("diagnosis", "").strip().lower()
+            if d:
+                labels.add(d)
+    return sorted(labels)
+
+def prepare_data():
+    from datasets import load_dataset
+    os.makedirs(IMAGES_DIR, exist_ok=True)
+
+    file_exists = os.path.exists(CSV_PATH)
+    if not file_exists:
+        with open(CSV_PATH, "w", newline="", encoding="utf-8") as f:
+            writer = csv.writer(f)
+            writer.writerow(CSV_COLUMNS)
+
+    rows_written = 0
+    sources_done = []
+
+    # ── IU-Xray (image + question + report) ──
+    print("Downloading IU-Xray from Hugging Face...")
+    iuxray = load_dataset("ayyuce/Indiana_University_Chest_X-ray_Collection", split="train")
+    written = 0
+    for i, example in enumerate(iuxray):
+        symptoms = (example.get("question") or "").strip()
+        diagnosis = (example.get("report") or "").strip()
+        image = example.get("image")
+        if not symptoms or not diagnosis or image is None:
+            continue
+        image_path = os.path.join(IMAGES_DIR, f"iu_xray_{i}.jpg")
+        image.convert("RGB").save(image_path)
+        with open(CSV_PATH, "a", newline="", encoding="utf-8") as f:
+            writer = csv.writer(f)
+            writer.writerow([image_path, "iu_xray", symptoms, diagnosis, ""])
+        written += 1
+    print(f"  IU-Xray: {written} rows")
+    rows_written += written
+    sources_done.append(f"iu_xray ({written})")
+
+    # ── NIH Chest X-ray (image + disease labels) ──
+    print("Downloading NIH Chest X-ray from Hugging Face...")
+    nih = load_dataset("g-ronimo/NIH-Chest-X-ray-dataset_resized300px", split="train", streaming=True)
+    label_names = [
+        "No Finding", "Atelectasis", "Cardiomegaly", "Effusion", "Infiltration",
+        "Mass", "Nodule", "Pneumonia", "Pneumothorax", "Consolidation",
+        "Edema", "Emphysema", "Fibrosis", "Pleural_Thickening", "Hernia"
+    ]
+    written = 0
+    for i, example in enumerate(nih):
+        if written >= 3000:
+            break
+        image = example.get("image")
+        label_indices = example.get("labels", [])
+        if image is None or not label_indices:
+            continue
+        label_str = "|".join(label_names[idx] for idx in label_indices)
+        primary_diagnosis = label_names[label_indices[0]]
+        image_path = os.path.join(IMAGES_DIR, f"nih_{i}.jpg")
+        image.convert("RGB").save(image_path)
+        with open(CSV_PATH, "a", newline="", encoding="utf-8") as f:
+            writer = csv.writer(f)
+            writer.writerow([image_path, "nih", "", primary_diagnosis, label_str])
+        written += 1
+        if written % 500 == 0:
+            print(f"  NIH progress: {written}...")
+    print(f"  NIH: {written} rows")
+    rows_written += written
+    sources_done.append(f"nih ({written})")
+
+    print(f"Done. Total: {rows_written} rows written to {CSV_PATH}")
+    print(f"Sources: {', '.join(sources_done)}")
+
+def add_data(image_path, symptoms, diagnosis, labels=""):
+    os.makedirs(DATA_DIR, exist_ok=True)
+    file_exists = os.path.exists(CSV_PATH)
+    with open(CSV_PATH, "a", newline="", encoding="utf-8") as f:
+        writer = csv.writer(f)
+        if not file_exists:
+            writer.writerow(CSV_COLUMNS)
+        writer.writerow([image_path, "user", symptoms, diagnosis, labels])
+    print(f"Added 1 row to {CSV_PATH}: diagnosis='{diagnosis}'")
+
+class FusionDataset(Dataset):
+    def __init__(self, csv_path, label_list):
+        self.rows = []
+        with open(csv_path, newline="", encoding="utf-8") as f:
+            for row in csv.DictReader(f):
+                d = row.get("diagnosis", "").strip().lower()
+                if d:
+                    self.rows.append(row)
+        self.label_list = label_list
+
+    def __len__(self):
+        return len(self.rows)
+
+    def __getitem__(self, idx):
+        row = self.rows[idx]
+        image = Image.open(row["image_path"]).convert("RGB")
+        symptoms = row.get("symptoms", "").strip()
+        label_idx = self.label_list.index(row["diagnosis"].strip().lower())
+        return image, symptoms, label_idx
+
+def collate_fn(batch):
+    images = [item[0] for item in batch]
+    symptoms = [item[1] for item in batch]
+    labels = torch.tensor([item[2] for item in batch], dtype=torch.long)
+    return images, symptoms, labels
+
+class DiagnosisFusionModel(nn.Module):
+    def __init__(self, num_conditions):
+        super().__init__()
+        from transformers import CLIPModel, CLIPProcessor, AutoTokenizer, AutoModel
+        self.image_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+        self.image_encoder = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+        self.symptom_tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+        self.symptom_encoder = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+        for param in self.image_encoder.parameters():
+            param.requires_grad = False
+        for param in self.symptom_encoder.parameters():
+            param.requires_grad = False
+        self.classifier = nn.Sequential(
+            nn.Linear(512 + 768, 256),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(256, num_conditions),
+        )
+
+    def encode_images(self, images):
+        inputs = self.image_processor(images=images, return_tensors="pt")
+        with torch.no_grad():
+            return self.image_encoder.get_image_features(**inputs)
+
+    def encode_symptoms(self, symptom_texts):
+        inputs = self.symptom_tokenizer(
+            symptom_texts, return_tensors="pt", padding=True, truncation=True, max_length=64
+        )
+        with torch.no_grad():
+            outputs = self.symptom_encoder(**inputs)
+            return outputs.last_hidden_state.mean(dim=1)
+
+    def forward(self, images, symptom_texts):
+        image_vecs = self.encode_images(images)
+        symptom_vecs = self.encode_symptoms(symptom_texts)
+        combined = torch.cat([image_vecs, symptom_vecs], dim=-1)
+        return self.classifier(combined)
+
+def train(epochs, batch_size, lr, val_split, use_amp, grad_accum):
+    from rich.console import Console
+    from rich.table import Table
+    from rich.progress import Progress, BarColumn, TextColumn, TimeElapsedColumn
+    _console = Console()
+
+    has_gpu = torch.cuda.is_available()
+    use_amp = use_amp and has_gpu
+    scaler = torch.cuda.amp.GradScaler() if use_amp else None
+
+    label_list = load_label_list()
+    if not label_list:
+        _console.print("[red]No data found. Run --mode prepare-data or --mode add-data first.[/red]")
+        return
+
+    dataset = FusionDataset(CSV_PATH, label_list)
+    val_size = max(int(val_split * len(dataset)), 1)
+    train_size = len(dataset) - val_size
+    train_subset, val_subset = random_split(dataset, [train_size, val_size])
+
+    train_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
+    val_loader = DataLoader(val_subset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
+
+    _console.print(f"[bold cyan]Training Setup[/bold cyan]")
+    _console.print(f"  Classes: {len(label_list)}")
+    _console.print(f"  Train/Val: {len(train_subset)}/{len(val_subset)}")
+    _console.print(f"  Batch size: {batch_size}  Grad accum: {grad_accum}")
+    _console.print(f"  Device: {'GPU' if has_gpu else 'CPU'}  AMP: {'ON' if use_amp else 'OFF'}")
+
+    model = DiagnosisFusionModel(num_conditions=len(label_list))
+    if has_gpu:
+        model = model.cuda()
+    optimizer = torch.optim.AdamW(model.classifier.parameters(), lr=lr)
+    loss_fn = nn.CrossEntropyLoss()
+
+    for epoch in range(epochs):
+        _console.print(f"\n[bold yellow]Epoch {epoch + 1}/{epochs}[/bold yellow]")
+        _console.print("-" * 40)
+
+        # ── Train ──
+        model.train()
+        train_loss = 0.0
+        optimizer.zero_grad()
+        train_progress = Progress(
+            TextColumn("[cyan]  Train[/cyan]"),
+            BarColumn(),
+            TextColumn("{task.completed}/{task.total}"),
+            TextColumn("[green]{task.fields[loss]:.4f}[/green]"),
+            TimeElapsedColumn(),
+            transient=True,
+        )
+        with train_progress:
+            task = train_progress.add_task("", total=len(train_loader), loss=0.0)
+            for i, (images, symptoms, labels) in enumerate(train_loader):
+                if has_gpu:
+                    labels = labels.cuda()
+                with torch.amp.autocast("cuda", enabled=use_amp):
+                    logits = model(images, symptoms)
+                    loss = loss_fn(logits, labels)
+                loss = loss / grad_accum
+                if use_amp:
+                    scaler.scale(loss).backward()
+                else:
+                    loss.backward()
+
+                if (i + 1) % grad_accum == 0 or (i + 1) == len(train_loader):
+                    if use_amp:
+                        scaler.step(optimizer)
+                        scaler.update()
+                    else:
+                        optimizer.step()
+                    optimizer.zero_grad()
+
+                train_loss += loss.item() * grad_accum
+                train_progress.update(task, advance=1, loss=loss.item() * grad_accum)
+
+        avg_train_loss = train_loss / len(train_loader)
+
+        # ── Validation ��─
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for images, symptoms, labels in val_loader:
+                if has_gpu:
+                    labels = labels.cuda()
+                logits = model(images, symptoms)
+                loss = loss_fn(logits, labels)
+                val_loss += loss.item()
+
+        avg_val_loss = val_loss / len(val_loader)
+
+        table = Table(show_header=False, box=None)
+        table.add_column("Metric", style="cyan")
+        table.add_column("Value", style="green")
+        table.add_row("Train loss", f"{avg_train_loss:.4f}")
+        table.add_row("Val loss",   f"{avg_val_loss:.4f}")
+        _console.print(table)
+
+    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+    torch.save({"model_state": model.classifier.state_dict(), "label_list": label_list}, CHECKPOINT_PATH)
+    _console.print(f"[green]Saved checkpoint to {CHECKPOINT_PATH}[/green]")
+
+def test(batch_size):
+    if not os.path.exists(CHECKPOINT_PATH):
+        print("No checkpoint found. Run --mode train first.")
+        return
+    checkpoint = torch.load(CHECKPOINT_PATH, weights_only=False)
+    label_list = checkpoint["label_list"]
+    model = DiagnosisFusionModel(num_conditions=len(label_list))
+    model.classifier.load_state_dict(checkpoint["model_state"])
+    model.eval()
+    dataset = FusionDataset(CSV_PATH, label_list)
+    test_size = max(int(0.2 * len(dataset)), 1)
+    _, test_subset = random_split(dataset, [len(dataset) - test_size, test_size])
+    loader = DataLoader(test_subset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
+    correct = 0
+    inconclusive = 0
+    total = 0
+    with torch.no_grad():
+        for images, symptoms, labels in loader:
+            logits = model(images, symptoms)
+            probs = torch.softmax(logits, dim=-1)
+            confidence, predicted = torch.max(probs, dim=-1)
+            for i in range(len(labels)):
+                total += 1
+                if confidence[i].item() < CONFIDENCE_THRESHOLD:
+                    inconclusive += 1
+                elif predicted[i].item() == labels[i].item():
+                    correct += 1
+    print(f"Tested on {total} held-out examples")
+    print(f"Correct (above confidence threshold): {correct} ({100 * correct / total:.1f}%)")
+    print(f"Flagged as inconclusive / needs follow-up: {inconclusive} ({100 * inconclusive / total:.1f}%)")
+
+def info():
+    if not os.path.exists(CSV_PATH):
+        print("No dataset.csv found. Run --mode prepare-data first.")
+        return
+    sources = {}
+    total = 0
+    with open(CSV_PATH, newline="", encoding="utf-8") as f:
+        for row in csv.DictReader(f):
+            src = row.get("source", "unknown")
+            sources[src] = sources.get(src, 0) + 1
+            total += 1
+    print(f"Dataset: {CSV_PATH}")
+    print(f"Total rows: {total}")
+    for src, count in sorted(sources.items()):
+        print(f"  {src}: {count}")
+    print(f"Images dir: {IMAGES_DIR}")
+    img_count = len([x for x in os.listdir(IMAGES_DIR) if os.path.isfile(os.path.join(IMAGES_DIR, x))]) if os.path.exists(IMAGES_DIR) else 0
+    print(f"Images: {img_count}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train/test the medical image+symptom fusion model")
+    parser.add_argument("--mode", required=True, choices=["prepare-data", "add-data", "train", "test", "info"])
+    parser.add_argument("--image", help="Path to an image file (for --mode add-data)")
+    parser.add_argument("--symptoms", help="Symptom description text (for --mode add-data)")
+    parser.add_argument("--diagnosis", help="Diagnosis label (for --mode add-data)")
+    parser.add_argument("--epochs", type=int, default=5)
+    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--val_split", type=float, default=0.15, help="Fraction of data for validation")
+    parser.add_argument("--use_amp", action="store_true", help="Enable mixed precision (GPU only)")
+    parser.add_argument("--grad_accum", type=int, default=1, help="Gradient accumulation steps")
+    args = parser.parse_args()
+
+    if args.mode == "prepare-data":
+        prepare_data()
+    elif args.mode == "add-data":
+        if not (args.image and args.symptoms and args.diagnosis):
+            print("--mode add-data requires --image, --symptoms, and --diagnosis")
+        else:
+            add_data(args.image, args.symptoms, args.diagnosis)
+    elif args.mode == "train":
+        train(args.epochs, args.batch_size, args.lr, args.val_split, args.use_amp, args.grad_accum)
+    elif args.mode == "test":
+        test(args.batch_size)
+    elif args.mode == "info":
+        info()

update.py

@@ -0,0 +1,261 @@
+"""
+Update script — pulls the newest models, UI files, and dependencies.
+
+Usage:
+  python update.py              # interactive menu
+  python update.py --all        # update everything
+  python update.py --models     # download latest ONNX models
+  python update.py --code       # git pull latest source
+  python update.py --deps       # upgrade pip packages
+
+The default model source is a Hugging Face repo.
+Configure with: python update.py --set-url <HF_REPO_ID>
+"""
+import argparse
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+CONFIG_FILE = "update_config.json"
+DEFAULT_HF_REPO = "your-org/medicalai-models"
+DEFAULT_DIR = Path("models") / "default"
+ONNX_FULL_DIR = Path("checkpoints") / "onnx_full"
+
+
+def load_config():
+    if os.path.exists(CONFIG_FILE):
+        with open(CONFIG_FILE) as f:
+            return json.load(f)
+    return {"hf_repo": DEFAULT_HF_REPO, "auto_update": True}
+
+
+def save_config(cfg):
+    with open(CONFIG_FILE, "w") as f:
+        json.dump(cfg, f, indent=2)
+
+
+def _ensure_dir(path):
+    os.makedirs(path, exist_ok=True)
+
+
+def _file_size(path):
+    return os.path.getsize(path) / 1024 / 1024
+
+
+def update_models_from_hf():
+    """Download latest ONNX models from Hugging Face."""
+    cfg = load_config()
+    repo = cfg["hf_repo"]
+
+    from rich.console import Console
+    console = Console()
+
+    if repo == DEFAULT_HF_REPO and "your-org" in repo:
+        console.print("[yellow]HF_REPO not set. Skipping model download.[/yellow]")
+        console.print("[yellow]Set your model repo: python update.py --set-url your-org/your-repo[/yellow]")
+        console.print("[yellow]Or train locally: python training.py --mode prepare-data && python training.py --mode train[/yellow]")
+        return
+
+    try:
+        import requests
+    except ImportError:
+        console.print("[red]'requests' required. Install: pip install requests[/red]")
+        return
+
+    _ensure_dir(ONNX_FULL_DIR)
+    _ensure_dir(DEFAULT_DIR)
+
+    files_to_download = [
+        ("fusion_full.onnx", ONNX_FULL_DIR / "fusion_full.onnx"),
+        ("labels.json", ONNX_FULL_DIR / "labels.json"),
+        ("fusion_classifier.onnx", DEFAULT_DIR / "fusion_classifier.onnx"),
+    ]
+
+    base_url = f"https://huggingface.co/{repo}/resolve/main"
+
+    for fname, dest in files_to_download:
+        url = f"{base_url}/{fname}"
+        console.print(f"[cyan]Downloading {fname}...[/cyan]")
+        try:
+            resp = requests.get(url, stream=True, timeout=30)
+            resp.raise_for_status()
+            with open(dest, "wb") as f:
+                for chunk in resp.iter_content(8192):
+                    f.write(chunk)
+            console.print(f"  [green]Saved {dest} ({_file_size(dest):.1f} MB)[/green]")
+        except Exception as e:
+            console.print(f"  [red]Failed: {e}[/red]")
+
+    console.print("[green]Model update complete.[/green]")
+
+
+def update_code():
+    """Pull latest source code from git."""
+    from rich.console import Console
+    console = Console()
+
+    if not os.path.exists(".git"):
+        console.print("[yellow]Not a git repository. Skipping code update.[/yellow]")
+        return
+
+    try:
+        result = subprocess.run(
+            ["git", "pull", "--ff-only"],
+            capture_output=True, text=True, timeout=60,
+        )
+        if result.returncode == 0:
+            console.print(f"[green]{result.stdout}[/green]")
+        else:
+            console.print(f"[yellow]{result.stderr}[/yellow]")
+    except Exception as e:
+        console.print(f"[red]Git pull failed: {e}[/red]")
+
+
+def update_deps():
+    """Upgrade all pip packages to latest compatible versions."""
+    from rich.console import Console
+    console = Console()
+
+    req = "requirements.txt"
+    if not os.path.exists(req):
+        console.print("[yellow]No requirements.txt found.[/yellow]")
+        return
+
+    console.print("[cyan]Upgrading dependencies...[/cyan]")
+    try:
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "install", "--upgrade", "-r", req],
+            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+        )
+        console.print("[green]Dependencies upgraded.[/green]")
+    except Exception as e:
+        console.print(f"[red]Upgrade failed: {e}[/red]")
+
+
+def update_all():
+    from rich.console import Console
+    console = Console()
+
+    console.print("[bold cyan]Full Update[/bold cyan]")
+    console.print()
+
+    console.print("[cyan]Step 1: Updating code...[/cyan]")
+    update_code()
+
+    console.print("[cyan]Step 2: Upgrading dependencies...[/cyan]")
+    update_deps()
+
+    console.print("[cyan]Step 3: Downloading latest models...[/cyan]")
+    update_models_from_hf()
+
+    console.print()
+    console.print("[green]Update complete![/green]")
+    console.print("  Run [bold]python quantization.py --mode status[/bold] to verify.")
+
+
+def set_repo_url(url):
+    cfg = load_config()
+    cfg["hf_repo"] = url
+    save_config(cfg)
+    print(f"Hugging Face repo set to: {url}")
+
+
+def show_status():
+    cfg = load_config()
+    print(f"Update config: {CONFIG_FILE}")
+    print(f"  HF repo:     {cfg['hf_repo']}")
+    print(f"  Auto update: {cfg['auto_update']}")
+    print()
+    print("Default models:")
+    for f in ["fusion_classifier.onnx", "labels.json"]:
+        p = DEFAULT_DIR / f
+        exists = os.path.exists(p)
+        size = f"({_file_size(p):.1f} MB)" if exists else ""
+        print(f"  {f}: {'yes' if exists else 'no'} {size}")
+    print()
+    print("Full ONNX pipeline:")
+    for f in ["fusion_full.onnx", "labels.json"]:
+        p = ONNX_FULL_DIR / f
+        exists = os.path.exists(p)
+        size = f"({_file_size(p):.1f} MB)" if exists else ""
+        print(f"  {f}: {'yes' if exists else 'no'} {size}")
+    print()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Update MedicalAI models, code, and deps")
+    parser.add_argument("--models", action="store_true", help="Download latest ONNX models")
+    parser.add_argument("--code", action="store_true", help="Git pull latest source")
+    parser.add_argument("--deps", action="store_true", help="Upgrade pip packages")
+    parser.add_argument("--all", action="store_true", help="Update everything")
+    parser.add_argument("--set-url", metavar="HF_REPO", help="Set Hugging Face model repo")
+    parser.add_argument("--status", action="store_true", help="Show update status")
+    args = parser.parse_args()
+
+    if args.set_url:
+        set_repo_url(args.set_url)
+        return
+
+    if args.status:
+        show_status()
+        return
+
+    if args.all:
+        update_all()
+        return
+
+    if args.models:
+        update_models_from_hf()
+        return
+
+    if args.code:
+        update_code()
+        return
+
+    if args.deps:
+        update_deps()
+        return
+
+    # Interactive mode
+    from rich.console import Console
+    import questionary
+
+    console = Console()
+    console.print("[bold cyan]MedicalAI - Update Manager[/bold cyan]")
+    console.print()
+
+    choice = questionary.select(
+        "What would you like to update?",
+        choices=[
+            "Everything (code + deps + models)",
+            "Models only (download latest ONNX)",
+            "Code only (git pull)",
+            "Dependencies only (pip upgrade)",
+            "Show update status",
+            "Set Hugging Face model repo",
+            "Cancel",
+        ],
+    ).ask()
+
+    if choice == "Everything (code + deps + models)":
+        update_all()
+    elif choice == "Models only (download latest ONNX)":
+        update_models_from_hf()
+    elif choice == "Code only (git pull)":
+        update_code()
+    elif choice == "Dependencies only (pip upgrade)":
+        update_deps()
+    elif choice == "Show update status":
+        show_status()
+    elif "Set Hugging Face" in choice:
+        repo = questionary.text("Enter Hugging Face repo (user/repo):").ask()
+        if repo:
+            set_repo_url(repo)
+    else:
+        console.print("[yellow]Cancelled.[/yellow]")
+
+
+if __name__ == "__main__":
+    main()

web_demo/index.html

@@ -0,0 +1,188 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+  <title>MedicalAI Web Demo</title>
+  <style>
+    * { box-sizing: border-box; margin: 0; padding: 0; }
+    body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; background: #f0f4f8; color: #1a202c; padding: 2rem; }
+    .container { max-width: 800px; margin: 0 auto; }
+    h1 { font-size: 1.75rem; margin-bottom: 0.25rem; }
+    .subtitle { color: #4a5568; margin-bottom: 2rem; }
+    .card { background: #fff; border-radius: 12px; padding: 1.5rem; margin-bottom: 1.5rem; box-shadow: 0 1px 3px rgba(0,0,0,0.1); }
+    .card h2 { font-size: 1.1rem; margin-bottom: 1rem; color: #2b6cb0; }
+    label { display: block; font-size: 0.875rem; font-weight: 600; margin-bottom: 0.375rem; color: #2d3748; }
+    input[type="file"], textarea, input[type="text"] { width: 100%; padding: 0.625rem; border: 1px solid #e2e8f0; border-radius: 8px; font-size: 0.9rem; }
+    textarea { resize: vertical; min-height: 60px; font-family: inherit; }
+    button { background: #2b6cb0; color: #fff; border: none; padding: 0.625rem 1.5rem; border-radius: 8px; font-size: 0.9rem; font-weight: 600; cursor: pointer; margin-top: 0.75rem; }
+    button:hover { background: #2c5282; }
+    button:disabled { opacity: 0.6; cursor: not-allowed; }
+    .result { margin-top: 1rem; padding: 1rem; background: #f7fafc; border-radius: 8px; border: 1px solid #e2e8f0; min-height: 40px; white-space: pre-wrap; }
+    .result .label { font-size: 0.75rem; text-transform: uppercase; letter-spacing: 0.05em; color: #718096; margin-bottom: 0.25rem; }
+    .result .value { font-size: 1rem; }
+    .badge { display: inline-block; font-size: 0.75rem; padding: 0.25rem 0.5rem; border-radius: 4px; font-weight: 600; }
+    .badge.online { background: #c6f6d5; color: #22543d; }
+    .badge.offline { background: #fed7d7; color: #822727; }
+    .row { display: flex; gap: 1rem; align-items: flex-end; }
+    .row > * { flex: 1; }
+    .status-bar { display: flex; gap: 1rem; align-items: center; flex-wrap: wrap; }
+    .preview img { max-width: 100%; max-height: 240px; border-radius: 8px; margin-top: 0.5rem; }
+    .code { background: #1a202c; color: #e2e8f0; padding: 1rem; border-radius: 8px; font-family: 'SF Mono', 'Fira Code', monospace; font-size: 0.8rem; overflow-x: auto; }
+    .loading { display: inline-block; width: 16px; height: 16px; border: 2px solid #e2e8f0; border-top-color: #2b6cb0; border-radius: 50%; animation: spin 0.6s linear infinite; margin-right: 0.5rem; vertical-align: middle; }
+    @keyframes spin { to { transform: rotate(360deg); } }
+  </style>
+</head>
+<body>
+  <div class="container">
+    <h1>MedicalAI Web Demo</h1>
+    <p class="subtitle">Connect your website to the MedicalAI API server</p>
+
+    <div class="status-bar" style="margin-bottom:1.5rem;">
+      <span id="healthBadge" class="badge offline">Checking API...</span>
+      <span id="modelInfo" style="font-size:0.85rem;color:#718096;"></span>
+    </div>
+
+    <div class="card">
+      <h2>Radiology Caption</h2>
+      <p style="font-size:0.85rem;color:#718096;margin-bottom:1rem;">Upload a chest X-ray and get an AI-generated radiology report.</p>
+      <label for="visionFile">X-ray Image</label>
+      <input type="file" id="visionFile" accept="image/*" />
+      <div id="visionPreview" class="preview"></div>
+      <button id="visionBtn" disabled>Analyze</button>
+      <div id="visionResult" class="result"><span style="color:#a0aec0;">Upload an image and click Analyze.</span></div>
+    </div>
+
+    <div class="card">
+      <h2>Symptom Check</h2>
+      <p style="font-size:0.85rem;color:#718096;margin-bottom:1rem;">Upload an X-ray with patient symptoms for diagnosis.</p>
+      <div class="row">
+        <div>
+          <label for="symptomFile">X-ray Image</label>
+          <input type="file" id="symptomFile" accept="image/*" />
+          <div id="symptomPreview" class="preview"></div>
+        </div>
+        <div>
+          <label for="symptoms">Symptoms / Clinical Indication</label>
+          <textarea id="symptoms" placeholder="e.g., shortness of breath, cough, fever..."></textarea>
+        </div>
+      </div>
+      <button id="symptomBtn" disabled>Diagnose</button>
+      <div id="symptomResult" class="result"><span style="color:#a0aec0;">Upload an image and enter symptoms.</span></div>
+    </div>
+
+    <div class="card">
+      <h2>How to connect your website</h2>
+      <p style="font-size:0.85rem;color:#718096;margin-bottom:1rem;">Use <code>fetch()</code> from JavaScript to call the API server running on your machine.</p>
+      <div class="code">// Example: Radiology Caption
+const form = new FormData();
+form.append('file', imageFile);
+
+const res = await fetch('http://YOUR_SERVER:8000/api/vision', {
+  method: 'POST',
+  body: form,
+});
+const data = await res.json();
+console.log(data.caption);
+
+// Example: Symptom Check
+const form2 = new FormData();
+form2.append('file', imageFile);
+form2.append('symptoms', 'cough, fever');
+
+const res2 = await fetch('http://YOUR_SERVER:8000/api/symptom-check', {
+  method: 'POST',
+  body: form2,
+});
+const data2 = await res2.json();
+console.log(data2.diagnosis, data2.confidence);</div>
+    </div>
+  </div>
+
+  <script>
+    const API_BASE = 'http://127.0.0.1:8000';
+
+    async function checkHealth() {
+      const badge = document.getElementById('healthBadge');
+      const info = document.getElementById('modelInfo');
+      try {
+        const res = await fetch(`${API_BASE}/health`);
+        if (!res.ok) throw new Error('Not OK');
+        const data = await res.json();
+        badge.className = 'badge online';
+        badge.textContent = 'API Online';
+        info.textContent = `Device: ${data.device} · Model: ${data.model_source}`;
+        document.getElementById('visionBtn').disabled = false;
+        document.getElementById('symptomBtn').disabled = false;
+      } catch {
+        badge.className = 'badge offline';
+        badge.textContent = 'API Offline';
+        info.textContent = 'Start the API server: python quantization.py --mode serve-api';
+      }
+    }
+    checkHealth();
+
+    function previewImage(input, previewId) {
+      const preview = document.getElementById(previewId);
+      preview.innerHTML = '';
+      if (input.files && input.files[0]) {
+        const img = document.createElement('img');
+        img.src = URL.createObjectURL(input.files[0]);
+        preview.appendChild(img);
+      }
+    }
+    document.getElementById('visionFile').addEventListener('change', function() {
+      previewImage(this, 'visionPreview');
+    });
+    document.getElementById('symptomFile').addEventListener('change', function() {
+      previewImage(this, 'symptomPreview');
+    });
+
+    async function callEndpoint(endpoint, fileInput, extraFields, resultId) {
+      const btn = resultId === 'visionResult' ? 'visionBtn' : 'symptomBtn';
+      const btnEl = document.getElementById(btn);
+      const resultEl = document.getElementById(resultId);
+
+      if (!fileInput.files || !fileInput.files[0]) {
+        resultEl.innerHTML = '<span style="color:#e53e3e;">Please select an image.</span>';
+        return;
+      }
+
+      btnEl.disabled = true;
+      btnEl.innerHTML = '<span class="loading"></span> Analyzing...';
+      resultEl.innerHTML = '<span class="loading"></span> Running inference...';
+
+      const form = new FormData();
+      form.append('file', fileInput.files[0]);
+      for (const [key, val] of Object.entries(extraFields)) {
+        form.append(key, val);
+      }
+
+      try {
+        const res = await fetch(`${API_BASE}${endpoint}`, { method: 'POST', body: form });
+        const data = await res.json();
+        if (!res.ok) throw new Error(data.detail || 'Request failed');
+
+        if (endpoint === '/api/vision') {
+          resultEl.innerHTML = `<div><div class="label">Radiology Caption</div><div class="value">${data.caption}</div><div style="font-size:0.75rem;color:#a0aec0;margin-top:0.5rem;">${data.inference_time_ms} ms</div></div>`;
+        } else {
+          resultEl.innerHTML = `<div><div class="label">Diagnosis</div><div class="value">${data.diagnosis}</div><div style="margin-top:0.5rem;"><span class="label">Confidence</span> <span class="value">${(data.confidence * 100).toFixed(1)}%</span></div><div style="font-size:0.75rem;color:#a0aec0;margin-top:0.5rem;">${data.inference_time_ms} ms</div></div>`;
+        }
+      } catch (err) {
+        resultEl.innerHTML = `<span style="color:#e53e3e;">Error: ${err.message}</span>`;
+      } finally {
+        btnEl.disabled = false;
+        btnEl.textContent = endpoint === '/api/vision' ? 'Analyze' : 'Diagnose';
+      }
+    }
+
+    document.getElementById('visionBtn').addEventListener('click', () => {
+      callEndpoint('/api/vision', document.getElementById('visionFile'), {}, 'visionResult');
+    });
+    document.getElementById('symptomBtn').addEventListener('click', () => {
+      const symptoms = document.getElementById('symptoms').value || 'No symptoms provided';
+      callEndpoint('/api/symptom-check', document.getElementById('symptomFile'), { symptoms }, 'symptomResult');
+    });
+  </script>
+</body>
+</html>

web_ui.py

@@ -0,0 +1,186 @@
+"""
+Browser-based UI for MedicalAI - Light Weight.
+Auto-installs gradio if missing.
+Auto-generates default models if none are trained.
+"""
+import importlib
+import os
+import subprocess
+import sys
+import tempfile
+
+from PIL import Image
+
+from optimize import (
+    clear_memory,
+    get_available_models,
+    get_device,
+    infer_blip,
+    infer_fusion,
+    infer_fusion_onnx,
+    set_cpu_threads,
+)
+
+CHECKPOINT_PATH = "./checkpoints/fusion_model.pth"
+ONNX_FULL_DIR = "./checkpoints/onnx_full"
+DEFAULT_MODEL_DIR = "./models/default"
+
+
+def _ensure_gradio():
+    try:
+        return importlib.import_module("gradio")
+    except ImportError:
+        from rich.console import Console
+        console = Console()
+        console.print("[yellow]'gradio' is required for the web UI.[/yellow]")
+        import questionary
+        install = questionary.confirm("Install gradio now?", default=True).ask()
+        if not install:
+            console.print("[red]gradio is required. Exiting.[/red]")
+            sys.exit(1)
+        console.print("[cyan]Installing gradio...[/cyan]")
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "gradio"])
+        return importlib.import_module("gradio")
+
+
+def _ensure_default_models():
+    """Generate default ONNX models if none exist at all."""
+    if os.path.exists(os.path.join(DEFAULT_MODEL_DIR, "fusion_classifier.onnx")):
+        return
+    if os.path.exists(CHECKPOINT_PATH):
+        return
+    print("No models found. Generating default models...")
+    subprocess.check_call(
+        [sys.executable, "setup_default.py"],
+        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+    )
+    print("Default models ready.")
+
+
+def analyze_vision(image):
+    if image is None:
+        return "Please upload an X-ray image."
+    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
+        path = f.name
+        Image.fromarray(image).save(path)
+    try:
+        caption = infer_blip(path, use_onnx=False)
+        return caption
+    except Exception as e:
+        return f"Error: {e}"
+    finally:
+        os.unlink(path)
+        clear_memory()
+
+
+def analyze_symptom(image, symptoms, use_onnx):
+    if image is None:
+        return "Please upload an X-ray image.", ""
+    if not symptoms:
+        symptoms = "No symptoms provided"
+    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
+        path = f.name
+        Image.fromarray(image).save(path)
+    try:
+        if use_onnx:
+            diagnosis, confidence = infer_fusion_onnx(path, symptoms)
+        else:
+            diagnosis, confidence = infer_fusion(path, symptoms)
+        if diagnosis is None:
+            return confidence, "No result"
+        return diagnosis, f"{confidence:.1%}"
+    except Exception as e:
+        return f"Error: {e}", ""
+    finally:
+        os.unlink(path)
+        clear_memory()
+
+
+def main():
+    set_cpu_threads()
+    _ensure_default_models()
+    gr = _ensure_gradio()
+
+    models = get_available_models()
+    has_trained = models["trained_pytorch"]
+    has_default = models["default_classifier"]
+    has_onnx_full = models["onnx_full_pipeline"]
+
+    model_source = "ONNX (full pipeline)" if has_onnx_full else \
+                   "Trained PyTorch" if has_trained else \
+                   "Default (random weights)" if has_default else \
+                   "NOT AVAILABLE"
+
+    print(f"Device: {get_device().upper()}")
+    print(f"Model:  {model_source}")
+    print()
+    print("Launching web UI... open the URL below in your browser.")
+
+    with gr.Blocks(title="MedicalAI - Light Weight", theme=gr.themes.Soft()) as demo:
+        gr.Markdown(
+            """
+            # MedicalAI - Light Weight
+            Upload a chest X-ray for AI-assisted analysis.
+            """
+        )
+
+        with gr.Tab("Vision - Generate Report"):
+            gr.Markdown("Upload an X-ray and get an AI-generated radiology caption.")
+            with gr.Row():
+                img_in = gr.Image(label="X-ray Image", type="numpy")
+            with gr.Row():
+                btn_vision = gr.Button("Generate Report", variant="primary")
+            with gr.Row():
+                caption_out = gr.Textbox(label="Radiology Caption", lines=6)
+
+            btn_vision.click(fn=analyze_vision, inputs=img_in, outputs=caption_out)
+
+        with gr.Tab("Symptom Check - Diagnosis"):
+            gr.Markdown("Upload an X-ray, enter symptoms, and get a diagnosis.")
+            with gr.Row():
+                img_in2 = gr.Image(label="X-ray Image", type="numpy")
+            with gr.Row():
+                symptoms_in = gr.Textbox(
+                    label="Patient Symptoms / Clinical Indication",
+                    placeholder="e.g., shortness of breath, cough, fever...",
+                    lines=3,
+                )
+            with gr.Row():
+                onnx_checkbox = gr.Checkbox(
+                    label="Use ONNX (no PyTorch backend)",
+                    value=has_onnx_full,
+                    interactive=has_onnx_full,
+                )
+            with gr.Row():
+                btn_diag = gr.Button("Diagnose", variant="primary")
+            with gr.Row():
+                with gr.Column():
+                    diag_out = gr.Textbox(label="Diagnosis", lines=4)
+                    conf_out = gr.Textbox(label="Confidence")
+
+            btn_diag.click(
+                fn=analyze_symptom,
+                inputs=[img_in2, symptoms_in, onnx_checkbox],
+                outputs=[diag_out, conf_out],
+            )
+
+        with gr.Tab("System Info"):
+            gr.Markdown(f"""
+            **Device:** `{get_device().upper()}`
+            **Model:** `{model_source}`
+            **Trained checkpoint:** {'Yes' if has_trained else 'No'}
+            **ONNX full pipeline:** {'Yes' if has_onnx_full else 'No'}
+            **Default model:** {'Yes' if has_default else 'No'}
+            **Dataset:** `{os.path.abspath('./data/dataset.csv')}`
+
+            ### 3 ways to improve accuracy
+            1. **Train** — `python training.py --mode prepare-data && python training.py --mode train`
+            2. **Update** — `python update.py --models` (downloads pre-trained model from Hugging Face)
+            3. **Export to ONNX** — `python quantization.py --mode export-full`
+            """)
+
+    demo.launch(server_name="127.0.0.1", server_port=7860, share=False)
+
+
+if __name__ == "__main__":
+    main()

xray_training.py

@@ -0,0 +1,273 @@
+import argparse
+import os
+
+import torch
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader, random_split
+from tqdm import tqdm
+
+MODEL_DIR = "./blip-xray-finetuned"
+CHECKPOINT_PATH = os.path.join(MODEL_DIR, "xray_blip.pth")
+ONNX_PATH = os.path.join(MODEL_DIR, "onnx")
+
+class RadiologyCaptionDataset(Dataset):
+    def __init__(self, hf_dataset, max_samples=None):
+        self.data = []
+        for i, example in enumerate(hf_dataset):
+            if max_samples and i >= max_samples:
+                break
+            image = example.get("image")
+            caption = example.get("caption", "").strip()
+            if image is None or not caption:
+                continue
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+            self.data.append((image, caption))
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+def collate_fn(batch, processor):
+    images = [item[0] for item in batch]
+    captions = [item[1] for item in batch]
+    encoding = processor(
+        images=images, text=captions, return_tensors="pt", padding=True, truncation=True, max_length=128
+    )
+    encoding["labels"] = encoding["input_ids"].clone()
+    return encoding
+
+def train(epochs, batch_size, lr, max_samples, resume, val_split, use_amp, grad_accum):
+    from transformers import BlipProcessor, BlipForConditionalGeneration
+    from datasets import load_dataset
+
+    has_gpu = torch.cuda.is_available()
+    use_amp = use_amp and has_gpu
+    scaler = torch.cuda.amp.GradScaler() if use_amp else None
+
+    hf_ds = load_dataset("eltorio/ROCOv2-radiology", split="train", streaming=True)
+    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+
+    if resume and os.path.exists(CHECKPOINT_PATH):
+        print(f"Resuming from checkpoint: {CHECKPOINT_PATH}")
+        model = BlipForConditionalGeneration.from_pretrained(MODEL_DIR)
+        start_epoch = torch.load(CHECKPOINT_PATH, weights_only=False).get("epoch", 0) + 1
+    else:
+        model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+        start_epoch = 0
+
+    if has_gpu:
+        model = model.cuda()
+
+    from functools import partial
+    _collate = partial(collate_fn, processor=processor)
+
+    full_dataset = RadiologyCaptionDataset(hf_ds, max_samples)
+    val_size = max(int(val_split * len(full_dataset)), 1)
+    train_size = len(full_dataset) - val_size
+    train_subset, val_subset = random_split(full_dataset, [train_size, val_size])
+
+    train_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True, collate_fn=_collate)
+    val_loader = DataLoader(val_subset, batch_size=batch_size, shuffle=False, collate_fn=_collate)
+    print(f"Dataset: {len(train_subset)} train + {len(val_subset)} val samples")
+    print(f"Device: {'GPU' if has_gpu else 'CPU'}  AMP: {'ON' if use_amp else 'OFF'}  Grad accum: {grad_accum}")
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+    model.train()
+
+    for epoch in range(start_epoch, epochs):
+        print(f"\n{'='*40}")
+        print(f"  Epoch {epoch + 1}/{epochs}")
+        print(f"{'='*40}")
+
+        # ── Train ──
+        total_loss = 0.0
+        optimizer.zero_grad()
+        pbar = tqdm(train_loader, desc=f"  Train")
+        for i, batch in enumerate(pbar):
+            pixel_values = batch.get("pixel_values")
+            input_ids = batch.get("input_ids")
+            attention_mask = batch.get("attention_mask")
+            labels = batch.get("labels")
+
+            if has_gpu:
+                pixel_values = pixel_values.cuda()
+                input_ids = input_ids.cuda()
+                attention_mask = attention_mask.cuda()
+                labels = labels.cuda()
+
+            with torch.amp.autocast("cuda", enabled=use_amp):
+                outputs = model(
+                    pixel_values=pixel_values,
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    labels=labels,
+                )
+                loss = outputs.loss / grad_accum
+
+            if use_amp:
+                scaler.scale(loss).backward()
+            else:
+                loss.backward()
+
+            if (i + 1) % grad_accum == 0 or (i + 1) == len(train_loader):
+                if use_amp:
+                    scaler.step(optimizer)
+                    scaler.update()
+                else:
+                    optimizer.step()
+                optimizer.zero_grad()
+
+            total_loss += loss.item() * grad_accum
+            pbar.set_postfix(loss=f"{loss.item() * grad_accum:.4f}")
+
+        avg_train_loss = total_loss / len(train_loader)
+
+        # ── Validation ──
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for batch in tqdm(val_loader, desc=f"  Val"):
+                pixel_values = batch.get("pixel_values")
+                input_ids = batch.get("input_ids")
+                attention_mask = batch.get("attention_mask")
+                labels = batch.get("labels")
+                outputs = model(
+                    pixel_values=pixel_values,
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    labels=labels,
+                )
+                val_loss += outputs.loss.item()
+        model.train()
+
+        avg_val_loss = val_loss / len(val_loader)
+        print(f"  Train loss: {avg_train_loss:.4f}  |  Val loss: {avg_val_loss:.4f}")
+
+        os.makedirs(MODEL_DIR, exist_ok=True)
+        model.save_pretrained(MODEL_DIR)
+        processor.save_pretrained(MODEL_DIR)
+        torch.save({"epoch": epoch}, CHECKPOINT_PATH)
+        print(f"  Checkpoint saved to {MODEL_DIR}")
+
+    print(f"\nTraining complete. Model saved to {MODEL_DIR}")
+
+def evaluate(batch_size, max_samples):
+    from transformers import BlipProcessor, BlipForConditionalGeneration
+    from datasets import load_dataset
+
+    if not os.path.exists(MODEL_DIR):
+        print(f"No model found at {MODEL_DIR}. Run training first.")
+        return
+
+    from functools import partial
+    print("Loading model and dataset...")
+    processor = BlipProcessor.from_pretrained(MODEL_DIR)
+    model = BlipForConditionalGeneration.from_pretrained(MODEL_DIR)
+    model.eval()
+
+    _collate = partial(collate_fn, processor=processor)
+    hf_ds = load_dataset("eltorio/ROCOv2-radiology", split="train", streaming=True)
+    dataset = RadiologyCaptionDataset(hf_ds, max_samples)
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=_collate)
+
+    try:
+        from nltk.translate.bleu_score import corpus_bleu, SmoothingFunction
+        smoothie = SmoothingFunction().method4
+    except ImportError:
+        print("nltk not installed. Skipping BLEU evaluation.")
+        print("Install with: pip install nltk")
+        return
+
+    print(f"Evaluating on {len(dataset)} samples...")
+    references = []
+    hypotheses = []
+    total_loss = 0.0
+
+    with torch.no_grad():
+        for batch in tqdm(loader, desc="Evaluating"):
+            pixel_values = batch.get("pixel_values")
+            input_ids = batch.get("input_ids")
+            attention_mask = batch.get("attention_mask")
+            labels = batch.get("labels")
+
+            outputs = model(
+                pixel_values=pixel_values,
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+            )
+            total_loss += outputs.loss.item()
+
+            generated_ids = model.generate(pixel_values=pixel_values, max_length=64)
+            for i in range(len(labels)):
+                ref = processor.decode(labels[i], skip_special_tokens=True)
+                hyp = processor.decode(generated_ids[i], skip_special_tokens=True)
+                if ref and hyp:
+                    references.append([ref.split()])
+                    hypotheses.append(hyp.split())
+
+    avg_loss = total_loss / len(loader)
+    bleu = corpus_bleu(references, hypotheses, smoothing_function=smoothie)
+    print(f"Average loss: {avg_loss:.4f}")
+    print(f"Corpus BLEU: {bleu:.4f}")
+
+    print("\nSample generations:")
+    for i in range(min(5, len(references))):
+        ref = " ".join(references[i][0])
+        hyp = " ".join(hypotheses[i])
+        print(f"  REF: {ref[:120]}")
+        print(f"  HYP: {hyp[:120]}")
+        print()
+
+def export_onnx():
+    print("[yellow]BLIP ONNX export requires a newer version of optimum.[/yellow]")
+    print("[yellow]Run: pip install --upgrade optimum[/yellow]")
+    print("[yellow]Until then, the system uses PyTorch directly (no speed difference for inference).[/yellow]")
+
+def generate(image_path):
+    from transformers import BlipProcessor, BlipForConditionalGeneration
+
+    if not os.path.exists(MODEL_DIR):
+        print(f"No model found at {MODEL_DIR}. Run training first.")
+        return
+
+    print(f"Loading model from {MODEL_DIR}...")
+    processor = BlipProcessor.from_pretrained(MODEL_DIR)
+    model = BlipForConditionalGeneration.from_pretrained(MODEL_DIR)
+    model.eval()
+
+    image = Image.open(image_path).convert("RGB")
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        out = model.generate(**inputs, max_length=64)
+    caption = processor.decode(out[0], skip_special_tokens=True)
+    print(f"Generated caption: {caption}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Fine-tune BLIP for radiology caption generation")
+    parser.add_argument("--mode", required=True, choices=["train", "evaluate", "export-onnx", "generate"])
+    parser.add_argument("--epochs", type=int, default=3)
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--lr", type=float, default=5e-5)
+    parser.add_argument("--max_samples", type=int, default=500, help="Max samples for training/eval (remove to use all)")
+    parser.add_argument("--resume", action="store_true", help="Resume from last checkpoint")
+    parser.add_argument("--val_split", type=float, default=0.1, help="Fraction of data for validation")
+    parser.add_argument("--use_amp", action="store_true", help="Enable mixed precision (GPU only)")
+    parser.add_argument("--grad_accum", type=int, default=1, help="Gradient accumulation steps")
+    parser.add_argument("--image", help="Path to image for --mode generate")
+    args = parser.parse_args()
+
+    if args.mode == "train":
+        train(args.epochs, args.batch_size, args.lr, args.max_samples, args.resume, args.val_split, args.use_amp, args.grad_accum)
+    elif args.mode == "evaluate":
+        evaluate(args.batch_size, args.max_samples)
+    elif args.mode == "export-onnx":
+        export_onnx()
+    elif args.mode == "generate":
+        if not args.image:
+            print("--mode generate requires --image <path>")
+        else:
+            generate(args.image)