Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| """Train CellQuant on CVAT-annotated images. | |
| Examples | |
| -------- | |
| Final model on every annotated image: | |
| python train.py --data-dir /path/to/images_and_zips --out weights/cellquant.pt | |
| Leave-one-image-out fold (held-out image is only used for monitoring): | |
| python train.py --data-dir ... --holdout C110 --out runs/fold_C110.pt | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import math | |
| import os | |
| import sys | |
| import time | |
| import numpy as np | |
| import torch | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from cellquant.annotations import discover, load_image, rasterize # noqa: E402 | |
| from cellquant.data import AugConfig, Prefetcher, make_targets, negative_nucleus_points # noqa: E402 | |
| from cellquant.infer import Predictor, pick_device # noqa: E402 | |
| from cellquant.losses import multitask_loss # noqa: E402 | |
| from cellquant.metrics import evaluate_cells, evaluate_nostrils # noqa: E402 | |
| from cellquant.model import CellUNet, save_checkpoint # noqa: E402 | |
| from cellquant.postprocess import PostConfig, analyze # noqa: E402 | |
| def parse_args(argv=None): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--data-dir", required=True) | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--holdout", default="", help="image key to leave out (e.g. C110); empty = train on all") | |
| ap.add_argument("--iters", type=int, default=3000) | |
| ap.add_argument("--batch", type=int, default=4) | |
| ap.add_argument("--crop", type=int, default=512) | |
| ap.add_argument("--lr", type=float, default=3e-4) | |
| ap.add_argument("--weight-decay", type=float, default=1e-4) | |
| ap.add_argument("--warmup", type=int, default=100) | |
| ap.add_argument("--encoder", default="resnet34") | |
| ap.add_argument("--no-pretrained", action="store_true") | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--device", default="auto") | |
| ap.add_argument("--log-every", type=int, default=50) | |
| ap.add_argument("--val-every", type=int, default=500) | |
| ap.add_argument("--workers", type=int, default=3, help="augmentation threads") | |
| ap.add_argument("--p-focus", type=float, default=AugConfig.p_focus, | |
| help="share of crops forced to contain an antibody-negative nucleus (0 = off)") | |
| ap.add_argument("--ab-neg-weight", default="auto", | |
| help="loss weight of antibody-negative nucleus pixels: 'auto' = sqrt(pos px / neg px) clipped to [1, 5], or a number") | |
| return ap.parse_args(argv) | |
| def antibody_negative_weight(targets: list[dict], spec: str = "auto") -> float: | |
| """Class-balancing weight for the antibody head (see losses.masked_bce).""" | |
| if spec != "auto": | |
| return float(spec) | |
| pos = sum(float(((t["antibody"] > 0.5) & (t["cellmask"] > 0)).sum()) for t in targets) | |
| neg = sum(float(((t["antibody"] < 0.5) & (t["cellmask"] > 0)).sum()) for t in targets) | |
| if neg == 0: | |
| return 1.0 | |
| return float(np.clip(np.sqrt(pos / neg), 1.0, 5.0)) | |
| def lr_at(it: int, total: int, base: float, warmup: int) -> float: | |
| if it < warmup: | |
| return base * (it + 1) / warmup | |
| t = (it - warmup) / max(1, total - warmup) | |
| return base * (0.02 + 0.98 * 0.5 * (1 + math.cos(math.pi * t))) | |
| def quick_val(model: CellUNet, img: np.ndarray, R, device) -> dict: | |
| pred = Predictor(model, device=str(device), tta=1) | |
| res, _ = pred.predict(img, PostConfig()) | |
| ce = evaluate_cells(R.cell_inst, R.cell_labels, res.cell_inst, res.cell_labels) | |
| ne = evaluate_nostrils(R.nostril_inst, R.nostril_labels, res.nostril_inst, res.nostril_labels) | |
| return { | |
| "cell_det_f1": ce["detection"]["f1"], | |
| "cell_attr_both": ce["attribute_accuracy"]["both"], | |
| "nos_det_f1": ne["detection"]["f1"], | |
| "nos_type_acc": ne["type_accuracy"], | |
| "n_pred_cells": ce["n_pred"], "n_gt_cells": ce["n_gt"], | |
| } | |
| def main(argv=None): | |
| a = parse_args(argv) | |
| torch.manual_seed(a.seed) | |
| np.random.seed(a.seed) | |
| device = pick_device(a.device) | |
| samples, _ = discover(a.data_dir) | |
| if not samples: | |
| raise SystemExit("no annotated images found") | |
| train = [s for s in samples if s.key != a.holdout] | |
| held = [s for s in samples if s.key == a.holdout] | |
| if a.holdout and not held: | |
| raise SystemExit(f"holdout {a.holdout!r} not among {[s.key for s in samples]}") | |
| print(f"device={device} train={[s.key for s in train]} holdout={[s.key for s in held]}") | |
| images, targets, focus = [], [], [] | |
| for s in train: | |
| img = load_image(s.image_path) | |
| R = rasterize(s.annotation, img.shape[0], img.shape[1]) | |
| images.append(img) | |
| targets.append(make_targets(R)) | |
| focus.append(negative_nucleus_points(R)) | |
| ab_neg_weight = antibody_negative_weight(targets, a.ab_neg_weight) | |
| print(f"antibody-negative nuclei in training set: {sum(len(f) for f in focus)} " | |
| f"(loss weight {ab_neg_weight:.2f}, focus share {a.p_focus:.2f})") | |
| sampler = Prefetcher(images, targets, AugConfig(crop=a.crop, p_focus=a.p_focus), batch=a.batch, seed=a.seed, | |
| workers=a.workers, focus=focus if a.p_focus > 0 else None) | |
| held_img = held_R = None | |
| if held: | |
| held_img = load_image(held[0].image_path) | |
| held_R = rasterize(held[0].annotation, held_img.shape[0], held_img.shape[1]) | |
| model = CellUNet(a.encoder, pretrained=not a.no_pretrained).to(device) | |
| opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=a.weight_decay) | |
| history = {"train": [], "val": []} | |
| t0 = time.time() | |
| run = {} | |
| for it in range(a.iters): | |
| lr = lr_at(it, a.iters, a.lr, a.warmup) | |
| for g in opt.param_groups: | |
| g["lr"] = lr | |
| model.train() | |
| x, t = sampler.next() | |
| out = model(x.to(device)) | |
| loss, parts = multitask_loss(out, t, ab_neg_weight=ab_neg_weight) | |
| opt.zero_grad(set_to_none=True) | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) | |
| opt.step() | |
| for k, v in parts.items(): | |
| run[k] = run.get(k, 0.0) + v | |
| if (it + 1) % a.log_every == 0: | |
| avg = {k: v / a.log_every for k, v in run.items()} | |
| run = {} | |
| avg["iter"] = it + 1 | |
| avg["lr"] = lr | |
| avg["elapsed_s"] = time.time() - t0 | |
| history["train"].append(avg) | |
| print(f"it {it+1:5d} loss {avg['total']:.3f} cell {avg['ce_cell']:.3f}/{avg['dice_cell']:.3f} " | |
| f"nos {avg['ce_nos']:.3f}/{avg['dice_nos']:.3f} ab {avg['bce_ab']:.3f} cn {avg['bce_cn']:.3f} " | |
| f"dist {avg.get('mse_dist', 0.0):.4f} lr {lr:.2e} {avg['elapsed_s']:.0f}s", flush=True) | |
| if held and ((it + 1) % a.val_every == 0 or it + 1 == a.iters): | |
| model.eval() | |
| v = quick_val(model, held_img, held_R, device) | |
| v["iter"] = it + 1 | |
| history["val"].append(v) | |
| print(f" val[{a.holdout}] cell F1 {v['cell_det_f1']:.3f} attr {v['cell_attr_both']:.3f} " | |
| f"nos F1 {v['nos_det_f1']:.3f} type {v['nos_type_acc']:.3f} (pred {v['n_pred_cells']} / gt {v['n_gt_cells']} cells)", flush=True) | |
| sampler.close() | |
| model.eval().cpu() | |
| os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True) | |
| save_checkpoint(model, a.out, { | |
| "config": {**vars(a), "ab_neg_weight_used": ab_neg_weight}, | |
| "train_keys": [s.key for s in train], | |
| "holdout": a.holdout, | |
| "history": history, | |
| "labels": list(__import__("cellquant.labels", fromlist=["ALL_LABELS"]).ALL_LABELS), | |
| }) | |
| with open(os.path.splitext(a.out)[0] + "_history.json", "w") as f: | |
| json.dump(history, f, indent=1) | |
| print(f"saved {a.out} ({time.time()-t0:.0f}s)") | |
| if __name__ == "__main__": | |
| main() | |