Blame for server.py

"""Local model server for ScanLens.

Runs several image-classification models ("tests") and checks that each image
is the kind of scan a test was trained on before running it.

  GET  /tests        -> the available tests
  POST /tests/<id>   -> run one test on an image sent as the raw request body
  POST /tests/auto   -> work out the scan type and run every test made for it

Only safetensors weights are loaded; pickle (.bin) weights can run code when loaded.
"""
import hmac
import io
import json
import os
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "30")
# Xet transfers stalled on this machine; plain HTTPS downloads are reliable.
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")


def _token_from_env_file():
    """MedSigLIP is gated, so downloads need the HF_TOKEN kept in the app's .env.local."""
    path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".env.local")
    try:
        with open(path) as f:
            for line in f:
                if line.startswith("HF_TOKEN=") and line.strip()[9:].startswith("hf_"):
                    return line.strip()[9:]
    except OSError:
        pass
    return None


if not os.environ.get("HF_TOKEN") and _token_from_env_file():
    os.environ["HF_TOKEN"] = _token_from_env_file()

import numpy as np
import torch
from PIL import Image, UnidentifiedImageError
from transformers import AutoModel, AutoProcessor, pipeline

# Hosted (e.g. on Hugging Face Spaces) this runs with HOST=0.0.0.0 so it can be reached from outside.
HOST, PORT = os.environ.get("HOST", "127.0.0.1"), int(os.environ.get("PORT", "8000"))
# When set, every request must carry it in the X-Server-Secret header, so only the ScanLens app can
# use a public server. A custom header, because a private Space needs Authorization for its own token.
SECRET = os.environ.get("MODEL_SERVER_SECRET")
# "cuda" on the ZeroGPU Space (app.py sets it); this plain HTTP server runs on the CPU.
DEVICE = os.environ.get("DEVICE", "cpu")
MAX_BYTES = 25 * 1024 * 1024

# What the scan-type check can recognise. Only the types in references.json have tests; the rest
# exist so other images have somewhere to land instead of being forced into a test.
SCANS = {
    "chest-xray": "a chest X-ray radiograph, a front view of the whole chest showing the ribs, both lungs, the heart and the collarbones",
    "bone-xray": "an X-ray of a bone, hand, wrist, arm, leg or joint",
    "knee-xray": "an X-ray of the knee joint, showing the end of the thigh bone above the shin bone with the joint space between them",
    "brain-mri": "an MRI scan of the brain",
    "chest-ct": "a chest CT scan, a cross-sectional slice through the body seen from below, showing both lungs as dark ovals beside the spine",
    "retina-photo": "a retinal fundus photograph, a round orange-red image of the back of the eye showing the optic disc and blood vessels",
    "skin-photo": "a dermatoscope close-up photograph of a mole or skin lesion",
    "head-ct": "a CT scan of the head",
    "abdomen-ct": "a CT scan of the abdomen",
    "ultrasound": "an ultrasound image",
    "dental-xray": "a dental X-ray of teeth",
    "photo": "a photograph of a person, object or place",
    "document": "a screenshot, document or drawing",
}
SCAN_NAMES = {
    "chest-xray": "chest X-ray", "bone-xray": "bone X-ray", "brain-mri": "brain MRI", "chest-ct": "chest CT",
    "head-ct": "head CT", "abdomen-ct": "abdominal CT", "ultrasound": "ultrasound", "dental-xray": "dental X-ray",
    "photo": "photograph", "document": "screenshot or document",
    "knee-xray": "knee X-ray", "retina-photo": "retina photo", "skin-photo": "skin lesion photo",
}
# Colour photographs rather than scans: never split into slices.
PHOTO_SCANS = {"retina-photo", "skin-photo"}

# scoring: "softmax" for models that pick one label, "sigmoid" for models that score each label on its own.
# Optional "scans": every scan type a test accepts, when more than its main "scan".
TESTS = {
    "chest-xray-pneumonia": {
        "name": "Pneumonia", "scan": "chest-xray", "scoring": "softmax",
        "model": "nickmuchi/vit-finetuned-chest-xray-pneumonia",
        # main only has pickle weights; this PR from Hugging Face's conversion bot adds safetensors.
        "revision": "refs/pr/1",
        "about": "Normal or pneumonia. Trained mostly on children's chest X-rays.",
    },
    "chest-xray-covid": {
        "name": "COVID-19 and pneumonia", "scan": "chest-xray", "scoring": "softmax",
        "model": "GGital/vit-Covid", "revision": "main",
        "about": "COVID-19, other pneumonia, or normal.",
    },
    "chest-xray-14": {
        "name": "14 chest conditions", "scan": "chest-xray", "scoring": "sigmoid",
        "model": "Sohaibsoussi/ViT-NIH-Chest-X-ray-dataset-small", "revision": "main",
        "about": "Scores each of 14 conditions separately, such as effusion, nodule and cardiomegaly.",
    },
    "chest-xray-tb": {
        "name": "Tuberculosis", "scan": "chest-xray", "scoring": "softmax",
        "model": "runaksh/chest_xray_tuberculosis_detection", "revision": "main",
        # Uploaded without its image settings; it was fine-tuned from this model, so it uses its settings.
        "processor": "google/vit-base-patch16-224-in21k",
        "about": "Normal or tuberculosis. Only knows these two, so other lung disease can show as tuberculosis: "
                 "in testing it called 3 of 15 pneumonia X-rays tuberculosis.",
    },
    "bone-xray-fracture": {
        # Not offered for knee X-rays: in testing it called 11 of 24 knees with no known fracture fractured.
        "name": "Bone fracture", "scan": "bone-xray", "scoring": "softmax",
        "model": "prithivMLmods/Bone-Fracture-Detection", "revision": "main",
        "about": "Fractured or not fractured.",
    },
    "knee-xray-oa": {
        "name": "Knee osteoarthritis", "scan": "knee-xray", "scoring": "softmax",
        "model": "thanhle829/kneexpert-xray-kl", "revision": "bf9e73749478e69e372c6b10e0d02b9b0c7cd786",
        # Published for the timm library; loaded by _timm_convnext instead.
        "loader": "timm-convnext", "file": "convnext_tiny_weighted_smoothed",
        "labels": {"KL0": "KL 0 (none)", "KL1": "KL 1 (doubtful)", "KL2": "KL 2 (mild)",
                   "KL3": "KL 3 (moderate)", "KL4": "KL 4 (severe)"},
        "about": "Kellgren-Lawrence grade 0 to 4 on a front-view knee X-ray. Trained on close crops of the knee "
                 "joint; often one grade off, especially between grades 0, 1 and 2.",
    },
    "retina-photo-dr": {
        "name": "Diabetic retinopathy", "scan": "retina-photo", "scoring": "softmax", "color": True,
        "model": "nadahh/APTOS2019DetectionViaLLMMConvNextV2", "revision": "main",
        "labels": {"0": "No DR", "1": "Mild", "2": "Moderate", "3": "Severe", "4": "Proliferative DR"},
        "about": "No diabetic retinopathy, or mild, moderate, severe or proliferative. Trained on one clinic's "
                 "photos (APTOS 2019); on other photos it told retinopathy from none well but often got the grade wrong.",
    },
    "skin-photo-lesion": {
        "name": "Skin lesion", "scan": "skin-photo", "scoring": "softmax", "color": True,
        "model": "spycoder/vit-base-patch16-224-in21k-enhanced-ham10000", "revision": "main",
        "labels": {"akiec": "Actinic keratosis", "bcc": "Basal cell carcinoma", "bkl": "Benign keratosis",
                   "df": "Dermatofibroma", "mel": "Melanoma", "nv": "Mole (melanocytic nevus)", "vasc": "Vascular lesion"},
        "about": "Seven kinds of skin lesion, including melanoma, basal cell carcinoma and moles. Made for "
                 "dermatoscope close-ups; much less accurate on images from other clinics, and not for ordinary photos.",
    },
    "brain-mri-tumour": {
        "name": "Brain tumour", "scan": "brain-mri", "scoring": "softmax",
        "model": "prithivMLmods/BrainTumor-Classification-Mini", "revision": "main",
        "about": "No tumour, glioma, meningioma or pituitary tumour.",
    },
    "chest-ct-lung-cancer": {
        "name": "Lung cancer", "scan": "chest-ct", "scoring": "softmax",
        "model": "oohtmeel/swin-tiny-patch4-finetuned-lung-cancer-ct-scans", "revision": "main",
        "about": "Positive or negative for lung cancer on a single CT slice.",
    },
}
# Google's medical SigLIP: trained on X-ray, CT, MRI and more, so it tells scan types apart far better
# than the general-purpose SigLIP. Gated: needs HF_TOKEN from an account that accepted its terms.
GATE_MODEL = os.environ.get("GATE_MODEL", "google/medsiglip-448")
# The scan type is whichever set of real example scans (references.json) the image most resembles.
# In testing, real scans matched their type at 0.60-0.93 similarity and a screenshot of a scan at 0.53.
# Below this the image isn't clearly any scan type, so no test runs.
MIN_SCAN_SCORE = float(os.environ.get("MIN_SCAN_SCORE", "0.58"))
# Fundus photos vary more in colour and framing than scans: real ones matched their type at 0.54-0.84,
# while everyday photos stayed below 0.36 and every other scan type below 0.45.
MIN_SCAN_SCORES = {"retina-photo": 0.5}
# Text descriptions still catch scan types that have no tests (ultrasound, head CT...) when confident.
MIN_OTHER_TEXT_SCORE = 0.5
REFERENCES = os.path.join(os.path.dirname(os.path.abspath(__file__)), "references.json")


class ScanTypeCheck:
    """Works out the scan type by comparing the image with example scans of each type with tests,
    and with text descriptions of other kinds of image."""

    def __init__(self, model_id, references=REFERENCES):
        self.processor = AutoProcessor.from_pretrained(model_id)
        self.model = AutoModel.from_pretrained(model_id, use_safetensors=True).eval()
        self.keys = list(SCANS)
        texts = [f"This is {SCANS[k]}." for k in self.keys]
        tokens = self.processor(text=texts, padding="max_length", return_tensors="pt")
        # The text side is worked out once here, on the CPU; only images go through the model later.
        with torch.no_grad():
            t = self.model.get_text_features(**tokens)
        self.text = t / t.norm(dim=-1, keepdim=True)
        self.scale, self.bias = float(self.model.logit_scale.exp()), float(self.model.logit_bias)
        self.model.to(DEVICE)
        self.refs = {}
        if references:
            with open(references) as f:
                data = json.load(f)
            if data["model"] != model_id:
                raise ValueError(f"{references} was built with {data['model']}; run build_references.py")
            self.refs = {k: torch.tensor(v["vector"]) for k, v in data["types"].items()}

    def embed(self, image):
        with torch.no_grad():
            v = self.model.get_image_features(**self.processor(images=image, return_tensors="pt").to(DEVICE))
        return (v / v.norm(dim=-1, keepdim=True))[0].cpu()

    def __call__(self, image):
        v = self.embed(image)
        logits = (v @ self.text.T) * self.scale + self.bias
        probs = logits.softmax(dim=-1)
        i = int(probs.argmax())
        text_scan, text_score = self.keys[i], float(probs[i])
        if text_scan not in self.refs and text_score >= MIN_OTHER_TEXT_SCORE:
            return {"scan": text_scan, "name": SCAN_NAMES[text_scan], "score": text_score, "clear": True}
        scan, score = max(((k, float(v @ ref)) for k, ref in self.refs.items()), key=lambda x: x[1])
        return {"scan": scan, "name": SCAN_NAMES[scan], "score": round(score, 3), "clear": score >= MIN_SCAN_SCORES.get(scan, MIN_SCAN_SCORE)}


_pipes = {}
_locks = {key: threading.Lock() for key in [*TESTS, "gate"]}


def _load(key):
    """Load a model once; later calls reuse it. Safe to call from several threads."""
    with _locks[key]:
        if key not in _pipes:
            if key == "gate":
                print(f"Loading scan-type check ({GATE_MODEL}) ...", flush=True)
                _pipes[key] = ScanTypeCheck(GATE_MODEL)
            else:
                t = TESTS[key]
                print(f"Loading {key} ({t['model']} @ {t['revision']}) ...", flush=True)
                if t.get("loader") == "timm-convnext":
                    _pipes[key] = _timm_convnext(t)
                else:
                    # "processor": a repo to take the image preprocessing from, for a model uploaded without its own.
                    extra = {"image_processor": t["processor"]} if t.get("processor") else {}
                    _pipes[key] = pipeline("image-classification", model=t["model"], revision=t["revision"], device=DEVICE,
                                           model_kwargs={"use_safetensors": True}, **extra)
            print(f"Ready: {key}", flush=True)
        return _pipes[key]


def _timm_convnext(t):
    """A pipeline for a ConvNeXt-tiny trained with timm and published as bare safetensors.

    Its weights are renamed to the transformers ConvNeXt layout, so timm itself isn't needed.
    """
    import re
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file
    from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor

    meta_path = hf_hub_download(t["model"], f"{t['file']}.json", revision=t["revision"])
    with open(meta_path) as f:
        meta = json.load(f)
    weights = load_file(hf_hub_download(t["model"], f"{t['file']}.safetensors", revision=t["revision"]))
    renames = [
        (r"^stem\.0\.", "convnext.embeddings.patch_embeddings."),
        (r"^stem\.1\.", "convnext.embeddings.layernorm."),
        (r"^stages\.(\d+)\.downsample\.", r"convnext.encoder.stages.\1.downsampling_layer."),
        (r"^stages\.(\d+)\.blocks\.(\d+)\.conv_dw\.", r"convnext.encoder.stages.\1.layers.\2.dwconv."),
        (r"^stages\.(\d+)\.blocks\.(\d+)\.norm\.", r"convnext.encoder.stages.\1.layers.\2.layernorm."),
        (r"^stages\.(\d+)\.blocks\.(\d+)\.mlp\.fc1\.", r"convnext.encoder.stages.\1.layers.\2.pwconv1."),
        (r"^stages\.(\d+)\.blocks\.(\d+)\.mlp\.fc2\.", r"convnext.encoder.stages.\1.layers.\2.pwconv2."),
        (r"^stages\.(\d+)\.blocks\.(\d+)\.gamma$", r"convnext.encoder.stages.\1.layers.\2.layer_scale_parameter"),
        (r"^head\.norm\.", "convnext.layernorm."),
        (r"^head\.fc\.", "classifier."),
    ]
    state = {}
    for k, v in weights.items():
        for pattern, new in renames:
            k = re.sub(pattern, new, k)
        state[k] = v
    # The publisher's calibration: logits / temperature, folded into the last layer.
    temperature = float(meta.get("temperature", 1.0))
    state["classifier.weight"] = state["classifier.weight"] / temperature
    state["classifier.bias"] = state["classifier.bias"] / temperature
    labels = meta["classes"]
    config = ConvNextConfig(depths=[3, 3, 9, 3], hidden_sizes=[96, 192, 384, 768], num_labels=len(labels),
                            id2label=dict(enumerate(labels)), label2id={l: i for i, l in enumerate(labels)})
    model = ConvNextForImageClassification(config).eval()
    model.load_state_dict(state, strict=True)
    # Resize the short side and centre-crop a square, keeping proportions: trained on square crops of
    # the joint, it graded wider knee images better this way than squeezed to a square.
    processor = ConvNextImageProcessor(size={"shortest_edge": meta["input_size"]}, crop_pct=1.0,
                                       image_mean=meta["normalize"]["mean"], image_std=meta["normalize"]["std"])
    return pipeline("image-classification", model=model, image_processor=processor, device=DEVICE)


def _preload():
    for key in ["gate", *TESTS]:
        try:
            _load(key)
        except Exception as e:  # keep serving; the failing test reports its error when used
            print(f"Failed to load {key}: {e!r}", flush=True)


def list_tests():
    return [
        {"id": k, "name": t["name"], "scan": t["scan"], "scan_name": SCAN_NAMES[t["scan"]],
         "model": t["model"], "about": t["about"], "ready": k in _pipes}
        for k, t in TESTS.items()
    ]


def detect_scan(image):
    return _load("gate")(image)


def _test_info(test_id):
    t = TESTS[test_id]
    return {"id": test_id, "name": t["name"], "scan": t["scan"], "scan_name": SCAN_NAMES[t["scan"]]}


def _scans(t):
    """Every scan type a test accepts; "scan" is the main one, shown in list_tests."""
    return t.get("scans", [t["scan"]])


def _slices(image, scan):
    """Photos are never grids of slices; a dark fundus border or skin folds could look like one."""
    return None if scan in PHOTO_SCANS else find_slices(image)


def _prepare(t, image):
    # Scans are grayscale; tinted stock images (blue X-rays, pink-highlighted tumours) fooled the
    # models in testing, and grayscale versions of the same images were read correctly.
    # Colour photos (retina, skin) keep their colour: it is much of what those models read.
    return image.convert("RGB") if t.get("color") else image.convert("L").convert("RGB")


def _label(t, label):
    """The label to show: some models name their classes with codes such as "nv", "2" or "KL3"."""
    return t.get("labels", {}).get(label, label)


# Top labels that mean "nothing found"; there is nothing to point at, so no heatmap is made for them.
# Compared after "labels" renames, so these are the names shown, lower-cased. The skin test has no
# "nothing found" answer: every label is a kind of lesion.
NORMAL_LABELS = {"normal", "no finding", "no tumor", "not fractured", "negative", "no dr", "kl 0 (none)"}

# Occlusion heatmap: grey out a PATCH x PATCH block of an 8 x 8 grid at every position and
# measure how much the model's evidence for its finding drops. Big drops mark where it looked.
# Evidence is measured as log-odds, not probability: at 99.9% the probability barely moves when
# part of the scan is covered, which left confident findings with no heatmap.
GRID, PATCH = 8, 2
# Below this largest drop in log-odds, covering any area barely changes the answer and the map is noise.
MIN_HEAT = 0.1
# Circle up to this many separate areas, each at least PEAK_SHARE as strong as the strongest.
MAX_PEAKS, PEAK_SHARE = 3, 0.5


def occlusion_heatmap(clf, image, label, scoring):
    model = clf.model
    index = {name: int(i) for i, name in model.config.id2label.items()}[label]
    x = clf.image_processor(images=image, return_tensors="pt")["pixel_values"].to(model.device)
    _, _, h, w = x.shape
    ch, cw = h // GRID, w // GRID

    def confidence(batch):
        with torch.no_grad():
            logits = model(pixel_values=batch).logits
        if scoring == "sigmoid":
            return logits[:, index]  # a sigmoid logit is already the log-odds
        others = torch.cat([logits[:, :index], logits[:, index + 1:]], dim=1)
        return logits[:, index] - torch.logsumexp(others, dim=1)

    base = float(confidence(x)[0])
    positions = [(gy, gx) for gy in range(GRID - PATCH + 1) for gx in range(GRID - PATCH + 1)]
    drops = []
    for start in range(0, len(positions), 16):
        batch = []
        for gy, gx in positions[start:start + 16]:
            covered = x.clone()
            # 0 after normalisation is the processor's mean colour, a neutral grey.
            covered[:, :, gy * ch:(gy + PATCH) * ch, gx * cw:(gx + PATCH) * cw] = 0
            batch.append(covered)
        drops += (base - confidence(torch.cat(batch))).clamp(min=0).tolist()

    total = [[0.0] * GRID for _ in range(GRID)]
    count = [[0] * GRID for _ in range(GRID)]
    for (gy, gx), drop in zip(positions, drops):
        for yy in range(gy, gy + PATCH):
            for xx in range(gx, gx + PATCH):
                total[yy][xx] += drop
                count[yy][xx] += 1
    grid = [[total[y][x] / count[y][x] for x in range(GRID)] for y in range(GRID)]
    peak = max(max(row) for row in grid)
    if peak < MIN_HEAT:
        return None
    # Strongest cells first; skip any within two cells of one already circled so circles don't overlap.
    cells = sorted(((grid[y][x], y, x) for y in range(GRID) for x in range(GRID)), reverse=True)
    peaks = []
    for value, y, x in cells:
        if value < peak * PEAK_SHARE or len(peaks) == MAX_PEAKS:
            break
        if all(max(abs(y - py), abs(x - px)) >= 2 for _, py, px in peaks):
            peaks.append((value, y, x))
    return {
        "label": label,
        # Rows top to bottom, each 0-1 relative to the strongest cell; covers the whole image.
        "grid": [[round(v / peak, 3) for v in row] for row in grid],
        "peaks": [{"x": (x + 0.5) / GRID, "y": (y + 0.5) / GRID, "strength": round(v / peak, 3)} for v, y, x in peaks],
        # Kept for analyses saved before multiple circles existed.
        "peak": {"x": (peaks[0][2] + 0.5) / GRID, "y": (peaks[0][1] + 0.5) / GRID},
        "drop": round(peak, 3),
    }


def _bands(profile, min_gap, min_len):
    """Runs of lines with content along one axis, separated by at least min_gap empty lines."""
    runs, start, end, gap = [], None, 0, 0
    for i, has_content in enumerate(profile > 0.02):
        if has_content:
            if start is None:
                start = i
            gap, end = 0, i
        elif start is not None:
            gap += 1
            if gap >= min_gap:
                runs.append((start, end + 1))
                start = None
    if start is not None:
        runs.append((start, end + 1))
    return [r for r in runs if r[1] - r[0] >= min_len]


def find_slices(image):
    """Boxes of the slices when the image is a grid of slices (a montage), else None.

    The models were trained on one slice at a time; shown a whole grid, the brain model called a
    clearly visible tumour "No Tumor" at 96%. Tested slice by slice, the same slices were flagged.
    """
    g = np.asarray(image.convert("L"), dtype=np.float32) / 255
    h, w = g.shape
    bright = g > 0.12
    rows = _bands(bright.mean(axis=1), max(2, h // 200), h // 12)
    cols = _bands(bright.mean(axis=0), max(2, w // 200), w // 12)
    # A grid has real gaps both across and down; one band either way is a single image
    # (white padding beside a single scan looks like a gap in one direction only).
    if len(rows) < 2 or len(cols) < 2:
        return None
    # Slices in a grid are about the same size; split bands where neighbouring slices touch.
    unit = float(np.median([b - a for a, b in rows + cols]))

    def split(runs):
        out = []
        for a, b in runs:
            k = max(1, round((b - a) / unit))
            out += [(a + (b - a) * i // k, a + (b - a) * (i + 1) // k) for i in range(k)]
        return out

    boxes = [(x0, y0, x1, y1) for y0, y1 in split(rows) for x0, x1 in split(cols)
             if bright[y0:y1, x0:x1].mean() > 0.05]
    return boxes if len(boxes) >= 4 else None


def _classify_slices(test_id, image, boxes):
    """Run a test on each slice of a grid and box the slices with a finding."""
    t = TESTS[test_id]
    clf = _load(test_id)
    crops = [_prepare(t, image.crop(b)) for b in boxes]
    per_slice = clf(crops, function_to_apply=t["scoring"], top_k=len(clf.model.config.id2label))
    per_slice = [[{"label": _label(t, p["label"]), "score": p["score"]} for p in r] for r in per_slice]
    flagged = [(b, r) for b, r in zip(boxes, per_slice) if r[0]["label"].lower() not in NORMAL_LABELS]
    if flagged:
        # The overall answer is the most confident flagged slice.
        results = max(flagged, key=lambda f: f[1][0]["score"])[1]
    else:
        # Nothing flagged: average every slice's scores.
        totals = {}
        for r in per_slice:
            for p in r:
                totals[p["label"]] = totals.get(p["label"], 0) + p["score"] / len(per_slice)
        results = sorted(({"label": k, "score": v} for k, v in totals.items()), key=lambda p: -p["score"])
    return {
        "test": _test_info(test_id), "model": t["model"],
        "predictions": [{"label": r["label"], "score": float(r["score"])} for r in results],
        # In image pixels, the same shape detection models return.
        "regions": [{"label": r[0]["label"], "score": float(r[0]["score"]),
                     "box": {"xmin": b[0], "ymin": b[1], "xmax": b[2], "ymax": b[3]}} for b, r in flagged],
        "slices": {"count": len(boxes), "flagged": len(flagged)},
        # Covering parts of a whole grid says nothing useful; the boxes show where instead.
        "heatmap": None, "heatmap_note": None,
    }


def _classify(test_id, image, slices=None):
    if slices:
        return _classify_slices(test_id, image, slices)
    t = TESTS[test_id]
    clf = _load(test_id)
    image = _prepare(t, image)
    results = clf(image, function_to_apply=t["scoring"], top_k=len(clf.model.config.id2label))
    top = results[0]["label"]
    finding = _label(t, top).lower() not in NORMAL_LABELS
    heatmap = occlusion_heatmap(clf, image, top, t["scoring"]) if finding else None
    if heatmap:
        heatmap["label"] = _label(t, top)
    return {
        "test": _test_info(test_id), "model": t["model"],
        "predictions": [{"label": _label(t, r["label"]), "score": float(r["score"])} for r in results],
        "heatmap": heatmap,
        # A finding that no area of the scan explains: the model is reacting to the whole image
        # (contrast, style, source), which in testing went with unreliable results.
        "heatmap_note": "diffuse" if finding and heatmap is None else None,
    }


def _unclear(scan):
    return 422, {
        "error": "This doesn't clearly look like a medical scan. Upload the scan image itself, "
                 "not a photo or screenshot of it, cropped to the scan.",
        "scan_check": scan,
    }


def run_test(test_id, image):
    t = TESTS[test_id]
    scan = detect_scan(image)
    if not scan["clear"]:
        return _unclear(scan)
    if scan["scan"] not in _scans(t):
        needs = " or ".join(SCAN_NAMES[s] for s in _scans(t))
        return 422, {
            "error": f"This looks like a {scan['name']}, but the {t['name']} test needs a {needs}. "
                     f"Choose a test for {scan['name']} scans, or upload a {needs}.",
            "test": _test_info(test_id), "scan_check": scan,
        }
    return 200, {"results": [_classify(test_id, image, _slices(image, scan["scan"]))], "scan_check": scan}


def run_matching(image):
    """For "not sure which test": work out the scan type, then run every test made for it."""
    scan = detect_scan(image)
    if not scan["clear"]:
        return _unclear(scan)
    matching = [k for k, t in TESTS.items() if scan["scan"] in _scans(t)]
    if not matching:
        have = sorted({SCAN_NAMES[s] for t in TESTS.values() for s in _scans(t)})
        return 422, {
            "error": f"This looks like a {scan['name']}. ScanLens doesn't have a test for that yet. "
                     f"It has tests for {', '.join(have[:-1])} and {have[-1]} scans.",
            "scan_check": scan,
        }
    slices = _slices(image, scan["scan"])
    return 200, {"results": [_classify(k, image, slices) for k in matching], "scan_check": scan}


class Handler(BaseHTTPRequestHandler):
    def _send(self, status, body):
        data = json.dumps(body).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def _allowed(self):
        if SECRET and not hmac.compare_digest(self.headers.get("X-Server-Secret", ""), SECRET):
            self._send(401, {"error": "Missing or wrong X-Server-Secret header."})
            return False
        return True

    def do_GET(self):
        if not self._allowed():
            return
        if self.path.rstrip("/") in ("", "/tests"):
            return self._send(200, list_tests())
        self._send(404, {"error": "Not found."})

    def do_POST(self):
        if not self._allowed():
            return
        parts = self.path.strip("/").split("/")
        if len(parts) != 2 or parts[0] != "tests" or (parts[1] not in TESTS and parts[1] != "auto"):
            return self._send(404, {"error": "Unknown test."})
        length = int(self.headers.get("Content-Length") or 0)
        if not 0 < length <= MAX_BYTES:
            return self._send(400, {"error": "Send one image of up to 25 MB as the request body."})
        try:
            image = Image.open(io.BytesIO(self.rfile.read(length))).convert("RGB")
        except UnidentifiedImageError:
            return self._send(400, {"error": "The request body isn't an image this server can read."})
        try:
            self._send(*(run_matching(image) if parts[1] == "auto" else run_test(parts[1], image)))
        except Exception as e:
            self._send(500, {"error": f"The model failed: {e}"})


if __name__ == "__main__":
    threading.Thread(target=_preload, daemon=True).start()
    print(f"Serving on http://{HOST}:{PORT} (models load in the background)", flush=True)
    ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()