Spaces:
Running on Zero
Running on Zero
| """Builds public/samples/: a few labelled, anonymised scans from public datasets, shown as | |
| "Try a sample scan" on the upload page. Each keeps its dataset label as the known answer. | |
| .venv/bin/pip install pyarrow # only needed for this script | |
| .venv/bin/python build_samples.py | |
| Images already in public/samples/ are kept; delete one to download it again. | |
| """ | |
| import io | |
| import json | |
| import time | |
| import urllib.error | |
| import urllib.parse | |
| import urllib.request | |
| from pathlib import Path | |
| from PIL import Image | |
| HERE = Path(__file__).parent | |
| OUT = HERE.parent / "public" / "samples" | |
| HF = "https://huggingface.co/datasets" | |
| SOURCES = { | |
| "chest-xray": dict(repo="hf-vision/chest-xray-pneumonia", source="Chest X-ray pneumonia dataset (Kermany et al.)", licence="CC BY 4.0"), | |
| "brain-mri": dict(repo="BJyotibrat/Masoud-Nickparvar-Brain-Tumor-MRI-Dataset", source="Brain tumour MRI dataset (Masoud Nickparvar)", licence="CC BY 4.0"), | |
| "chest-ct": dict(repo="Mahadih534/Chest_CT-Scan_images-Dataset", source="Chest CT-scan images (Mohamed Hany)", licence="CC"), | |
| "chest-xray-tb": dict(repo="realsudarshan/chest-xray-tb-pneumonia", source="Chest X-ray TB, pneumonia and normal (compiled from public sets)", licence="CC BY 4.0"), | |
| "retina-photo": dict(repo="bumbledeep/eyepacs", source="EyePACS diabetic retinopathy photos (Kaggle 2015)", licence="MIT (as listed on the Hugging Face copy)"), | |
| "skin-photo": dict(repo="marmal88/skin_cancer", source="HAM10000 dermatoscopic images (Tschandl et al.)", licence="CC BY-NC 4.0"), | |
| "knee-xray": dict(repo="SilpaCS/kneeosteoarthritis", source="Knee osteoarthritis severity grading dataset (Chen, from the OAI)", licence="CC BY 4.0"), | |
| } | |
| # Row numbers of the parquet-only samples below (see from_rows). | |
| ROWS = {"chest-xray-tuberculosis": 1700, "retina-photo-moderate-dr": 12663, "retina-photo-no-dr": 12600, | |
| "skin-photo-melanoma": 1201, "skin-photo-nevus": 800, "knee-xray-kl3": 8000, "knee-xray-kl0": 500} | |
| # (sample id, scan type, dataset label, folder in the dataset repo). Chest X-rays come from a | |
| # parquet file instead of folders, so their folder is the class index. | |
| WANTED = [ | |
| ("chest-xray-normal", "chest-xray", "Normal", 0), | |
| ("chest-xray-pneumonia", "chest-xray", "Pneumonia", 1), | |
| ("brain-mri-glioma", "brain-mri", "Glioma", "Testing/glioma"), | |
| ("brain-mri-meningioma", "brain-mri", "Meningioma", "Testing/meningioma"), | |
| ("brain-mri-notumor", "brain-mri", "No tumour", "Testing/notumor"), | |
| ("chest-ct-adenocarcinoma", "chest-ct", "Adenocarcinoma", "test/adenocarcinoma"), | |
| ("chest-ct-normal", "chest-ct", "Normal", "test/normal"), | |
| # These datasets are only parquet files, so samples are picked by row through the datasets server: | |
| # (split, row number). A fifth item names the SOURCES entry when it isn't the scan type. | |
| ("chest-xray-tuberculosis", "chest-xray", "Tuberculosis", ("test", ROWS["chest-xray-tuberculosis"]), "chest-xray-tb"), | |
| ("retina-photo-moderate-dr", "retina-photo", "Moderate diabetic retinopathy", ("train", ROWS["retina-photo-moderate-dr"])), | |
| ("retina-photo-no-dr", "retina-photo", "No diabetic retinopathy", ("train", ROWS["retina-photo-no-dr"])), | |
| ("skin-photo-melanoma", "skin-photo", "Melanoma", ("test", ROWS["skin-photo-melanoma"])), | |
| ("skin-photo-nevus", "skin-photo", "Mole (melanocytic nevus)", ("test", ROWS["skin-photo-nevus"])), | |
| ("knee-xray-kl3", "knee-xray", "Kellgren-Lawrence grade 3", ("train", ROWS["knee-xray-kl3"])), | |
| ("knee-xray-kl0", "knee-xray", "Kellgren-Lawrence grade 0", ("train", ROWS["knee-xray-kl0"])), | |
| ] | |
| # Colour photographs stay in colour; scans are saved in grayscale. | |
| COLOUR = {"retina-photo", "skin-photo"} | |
| def get(url, tries=3): | |
| for attempt in range(tries): | |
| try: | |
| with urllib.request.urlopen(url, timeout=60) as r: | |
| return r.read() | |
| except urllib.error.URLError: # the datasets server sometimes answers 502 or 429 | |
| if attempt == tries - 1: | |
| raise | |
| time.sleep(10) | |
| def from_folder(repo, folder): | |
| """A file a third of the way through the folder, so reruns pick the same image.""" | |
| tree = json.loads(get(f"https://huggingface.co/api/datasets/{repo}/tree/main/{urllib.parse.quote(folder)}")) | |
| files = sorted(f["path"] for f in tree if f["type"] == "file") | |
| return get(f"{HF}/{repo}/resolve/main/{urllib.parse.quote(files[len(files) // 3])}") | |
| def from_parquet(repo, label): | |
| import pyarrow.parquet as pq | |
| data = get(f"{HF}/{repo}/resolve/main/data/validation-00000-of-00001.parquet") | |
| rows = pq.read_table(io.BytesIO(data)).to_pylist() | |
| return next(r for r in rows if r["label"] == label)["image"]["bytes"] | |
| def from_rows(repo, split, row): | |
| q = urllib.parse.quote(repo) | |
| config = json.loads(get(f"https://datasets-server.huggingface.co/splits?dataset={q}"))["splits"][0]["config"] | |
| rows = json.loads(get(f"https://datasets-server.huggingface.co/rows?dataset={q}&config={config}" | |
| f"&split={split}&offset={row}&length=1"))["rows"] | |
| image = next(v for v in rows[0]["row"].values() if isinstance(v, dict) and "src" in v) | |
| return get(image["src"]) | |
| def main(): | |
| OUT.mkdir(parents=True, exist_ok=True) | |
| manifest = [] | |
| for sid, scan, truth, where, *source in WANTED: | |
| src = SOURCES[source[0] if source else scan] | |
| path = OUT / f"{sid}.jpg" | |
| if not path.exists(): | |
| if isinstance(where, tuple): | |
| raw = from_rows(src["repo"], *where) | |
| elif isinstance(where, int): | |
| raw = from_parquet(src["repo"], where) | |
| else: | |
| raw = from_folder(src["repo"], where) | |
| img = Image.open(io.BytesIO(raw)).convert("RGB" if scan in COLOUR else "L") | |
| # Colour photos compress less well; smaller keeps each sample under about 150 KB. | |
| img.thumbnail((768, 768) if scan in COLOUR else (1024, 1024)) | |
| img.save(path, "JPEG", quality=85 if scan in COLOUR else 90) | |
| print("downloaded", path.name) | |
| manifest.append(dict(id=sid, file=path.name, scan=scan, truth=truth, source=src["source"], | |
| url=f"{HF}/{src['repo']}", licence=src["licence"])) | |
| (OUT / "samples.json").write_text(json.dumps(manifest, indent=2) + "\n") | |
| print(f"wrote {len(manifest)} samples") | |
| if __name__ == "__main__": | |
| main() | |