Spaces:
Running on Zero
Running on Zero
| """Builds references.json: the average MedSigLIP fingerprint of real example scans for each | |
| scan type that has tests. server.py uses it to tell scan types apart; text descriptions alone | |
| mistook photos of X-ray films and annotated X-rays for CT scans. | |
| .venv/bin/python build_references.py | |
| Downloads a few random images per type from public Hugging Face datasets into references/ | |
| (git-ignored), embeds them, and writes references.json. Needs HF_TOKEN (see server.py). | |
| """ | |
| import json | |
| import random | |
| import time | |
| import urllib.parse | |
| import urllib.request | |
| from pathlib import Path | |
| import torch | |
| from PIL import Image | |
| import server | |
| SOURCES = { | |
| "chest-xray": ["hf-vision/chest-xray-pneumonia", "sasi2004/chest-xrays-indiana-university"], | |
| "chest-ct": ["Mahadih534/Chest_CT-Scan_images-Dataset"], | |
| "bone-xray": ["Hemg/bone-fracture-detection"], | |
| "brain-mri": ["BJyotibrat/Masoud-Nickparvar-Brain-Tumor-MRI-Dataset", "Mahadih534/brain-tumor-MRI-dataset"], | |
| # Tight crops of the knee joint (the osteoarthritis grading set) and whole knee X-rays. | |
| "knee-xray": ["SilpaCS/kneeosteoarthritis", "Mary-m007/knee_classification01"], | |
| "retina-photo": ["bumbledeep/aptos", "bumbledeep/eyepacs"], | |
| # Dermatoscope images: the skin test was trained on these, not on ordinary photos of skin. | |
| "skin-photo": ["marmal88/skin_cancer", "MedOtter/ISIC2017"], | |
| } | |
| PER_DATASET = 6 | |
| HERE = Path(__file__).parent | |
| IMAGES = HERE / "references" | |
| def _get(url): | |
| for _ in range(3): | |
| try: | |
| with urllib.request.urlopen(url, timeout=40) as r: | |
| return json.load(r) | |
| except Exception: | |
| time.sleep(3) | |
| return None | |
| def download(): | |
| random.seed(7) | |
| for scan, datasets in SOURCES.items(): | |
| n = len(list(IMAGES.glob(f"{scan}-*.jpg"))) | |
| if n: | |
| continue # already downloaded | |
| for ds in datasets: | |
| q = urllib.parse.quote(ds) | |
| splits = (_get(f"https://datasets-server.huggingface.co/splits?dataset={q}") or {}).get("splits") | |
| if not splits: | |
| continue | |
| s = splits[0] | |
| size = _get(f"https://datasets-server.huggingface.co/size?dataset={q}") or {} | |
| total = next((x["num_rows"] for x in size.get("size", {}).get("splits", []) | |
| if x["split"] == s["split"] and x["config"] == s["config"]), 200) | |
| for offset in random.sample(range(total), min(PER_DATASET, total)): | |
| rows = (_get(f"https://datasets-server.huggingface.co/rows?dataset={q}&config={s['config']}" | |
| f"&split={s['split']}&offset={offset}&length=1") or {}).get("rows") | |
| if not rows: | |
| continue | |
| img = next((v for v in rows[0]["row"].values() if isinstance(v, dict) and "src" in v), None) | |
| if img: | |
| urllib.request.urlretrieve(img["src"], IMAGES / f"{scan}-{n}.jpg") | |
| n += 1 | |
| def main(): | |
| IMAGES.mkdir(exist_ok=True) | |
| download() | |
| check = server.ScanTypeCheck(server.GATE_MODEL, references=None) | |
| out = {"model": server.GATE_MODEL, "types": {}} | |
| for scan, datasets in SOURCES.items(): | |
| files = sorted(IMAGES.glob(f"{scan}-*.jpg")) | |
| vectors = torch.stack([check.embed(Image.open(f).convert("RGB")) for f in files]) | |
| mean = vectors.mean(0) | |
| out["types"][scan] = {"count": len(files), "datasets": datasets, "vector": [round(v, 6) for v in (mean / mean.norm()).tolist()]} | |
| print(f"{scan}: {len(files)} images") | |
| (HERE / "references.json").write_text(json.dumps(out)) | |
| print("wrote references.json") | |
| if __name__ == "__main__": | |
| main() | |