"""Prepare, train, evaluate, and predict. See README.md for commands."""
import argparse
import json
import os
import random
import time
from pathlib import Path
# Keep downloaded weights with the project, not in a global cache.
os.environ.setdefault("HF_HOME", str(Path(".cache/huggingface").resolve()))
import numpy as np
import torch
from torch import nn
from torch.utils.data import ConcatDataset, DataLoader, RandomSampler, WeightedRandomSampler
from data import (
BACKGROUND,
IMAGE_NAMES,
FolderDataset,
GridDataset,
SupervisedGrids,
grid_boxes,
prepare,
prepare_jsonl,
prepare_folders,
read_rgb,
record_image_path,
target_name,
transform,
)
from models import BACKBONE, GridModel, select_target, transfer_backbone
def save_json(path, value):
Path(path).write_text(json.dumps(value, indent=2) + "\n")
def device_for(name):
if name == "auto":
name = (
"cuda"
if torch.cuda.is_available()
else "mps"
if torch.backends.mps.is_available()
else "cpu"
)
if name == "mps" and not torch.backends.mps.is_available():
raise RuntimeError(
"MPS is unavailable. Run in a terminal with Apple GPU access; use --device cpu only intentionally."
)
if name == "cuda" and not torch.cuda.is_available():
raise RuntimeError("CUDA is unavailable on this machine.")
return torch.device(name)
def empty_cache(device):
if device.type == "mps":
torch.mps.empty_cache()
elif device.type == "cuda":
torch.cuda.empty_cache()
def seed_everything(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.set_num_threads(4)
def records_for(manifest, tiles, split):
return [r for r in manifest["records"] if r["tiles"] == tiles and r["split"] == split]
def auxiliary_records(manifest, targets, split):
return [r for r in records_for(manifest, 16, split) if r["target"] in targets]
def folder_records(folders, kind, split):
return (
[r for r in folders["records"] if r["kind"] == kind and r["split"] == split]
if folders
else []
)
def cell_logits(model, images):
"""Logits shaped (batch, classes, cells); a separate tile is one cell."""
output = model(images)
return output.unsqueeze(-1) if model.tiles == 9 else output.flatten(2)
def partial_loss(logits, labels, weights, any_labels, any_weights):
"""Weighted BCE on known cells, plus presence anywhere (max over cells) for image-level labels."""
cells = nn.functional.binary_cross_entropy_with_logits(
logits, labels, weight=weights, reduction="sum"
)
anywhere = nn.functional.binary_cross_entropy_with_logits(
logits.amax(-1), any_labels, weight=any_weights, reduction="sum"
)
total = weights.sum() + any_weights.sum()
return (cells + anywhere) / total.clamp(min=1), total
@torch.inference_mode()
def evaluate_folders(model, loader, device):
model.eval()
loss, total, probabilities, labels, weights = 0.0, 0.0, [], [], []
for images, *supervision in loader:
supervision = [x.to(device) for x in supervision]
logits = cell_logits(model, images.to(device)).float()
batch_loss, batch_total = partial_loss(logits, *supervision)
loss += batch_loss.item() * batch_total.item()
total += batch_total.item()
probabilities.append(logits.amax(-1).sigmoid().cpu())
labels.append(supervision[2].cpu())
weights.append(supervision[3].cpu())
return (
loss / max(total, 1),
torch.cat(probabilities),
torch.cat(labels).bool(),
torch.cat(weights),
)
def folder_report(probabilities, labels, weights, classes, threshold):
"""Single-label folder images: recall per class, and how often each class fires on other classes' images."""
predicted = probabilities >= threshold
background = (weights == 1).all(1) & ~labels.any(1)
report = {
"threshold": threshold,
"images": len(labels),
"per_class": {},
"background_images": int(background.sum()),
"background_any_selected": float(predicted[background].any(1).float().mean())
if background.any()
else None,
}
for i, name in enumerate(classes):
positive, others = labels[:, i], ~labels[:, i] & labels.any(1)
report["per_class"][name] = {
"positives": int(positive.sum()),
"recall": float(predicted[positive, i].float().mean()) if positive.any() else None,
"fires_on_other_classes": float(predicted[others, i].float().mean())
if others.any()
else None,
}
return report
def thresholds_for(checkpoint, targets):
return torch.tensor(
[
checkpoint.get("thresholds", {}).get(
checkpoint["classes"][int(i)], checkpoint["threshold"]
)
for i in targets
]
).unsqueeze(1)
def grid_logits(model, images, targets):
if model.tiles == 9:
batch, cells = images.shape[:2]
logits = model(images.flatten(0, 1)).reshape(batch, cells, -1).transpose(1, 2)
else:
logits = model(images).flatten(2)
selected = select_target(logits, targets)
# Presence in any enlarged subregion counts as presence in the tile.
for target_id in getattr(model, "crop_target_ids", []):
mask = targets == target_id
if not mask.any():
continue
tile_images = images[mask].flatten(0, 1)
height, width = tile_images.shape[-2:]
crop_h, crop_w = height // 2, width // 2
positions = [
(0, 0),
(0, width - crop_w),
(height - crop_h, 0),
(height - crop_h, width - crop_w),
((height - crop_h) // 2, (width - crop_w) // 2),
]
scores = selected[mask]
for top, left in positions:
crop = tile_images[:, :, top : top + crop_h, left : left + crop_w]
crop = nn.functional.interpolate(
crop, size=(height, width), mode="bilinear", align_corners=False
)
enlarged = model(crop)[:, target_id].reshape_as(scores)
scores = torch.maximum(scores, enlarged)
selected[mask] = scores
return selected
@torch.inference_mode()
def evaluate_loader(model, loader, device):
model.eval()
predictions, labels, targets = [], [], []
loss, total = 0.0, 0
for images, target, truth in loader:
logits = grid_logits(model, images.to(device), target.to(device))
loss += nn.functional.binary_cross_entropy_with_logits(
logits, truth.to(device), reduction="sum"
).item()
total += truth.numel()
predictions.append(logits.sigmoid().cpu())
labels.append(truth)
targets.append(target)
if not total:
raise ValueError("Evaluation split is empty")
return loss / total, torch.cat(predictions), torch.cat(labels).bool(), torch.cat(targets)
def metrics(probabilities, labels, threshold):
predicted = probabilities >= threshold
tp = int((predicted & labels).sum())
fp = int((predicted & ~labels).sum())
fn = int((~predicted & labels).sum())
return {
"grids": len(labels),
"exact_grid_accuracy": float((predicted == labels).all(1).float().mean()),
"cell_accuracy": float((predicted == labels).float().mean()),
"cell_precision": tp / max(tp + fp, 1),
"cell_recall": tp / max(tp + fn, 1),
"cell_f1": 2 * tp / max(2 * tp + fp + fn, 1),
"tp": tp,
"fp": fp,
"fn": fn,
}
def choose_threshold(probabilities, labels):
candidates = [round(0.1 + 0.025 * i, 3) for i in range(33)]
def rank(threshold):
result = metrics(probabilities, labels, threshold)
return result["exact_grid_accuracy"], result["cell_f1"], -abs(threshold - 0.5)
return max(candidates, key=rank)
def load_checkpoint(path, device):
checkpoint = torch.load(path, map_location="cpu", weights_only=True)
model = GridModel(
checkpoint["tiles"],
len(checkpoint["classes"]),
pretrained=False,
backbone=checkpoint.get("backbone", BACKBONE),
)
model.load_state_dict(checkpoint["model"])
model.crop_target_ids = [
checkpoint["classes"].index(name) for name in checkpoint.get("crop_targets", [])
]
if model.crop_target_ids and checkpoint["tiles"] != 9:
raise ValueError("Crop views are supported only by tile classifiers")
return model.to(device).eval(), checkpoint
def train_one(args, tiles, manifest, folders=None):
seed_everything(args.seed)
device = device_for(args.device)
output = Path(args.output) / f"grid{tiles}"
if (output / "best.pt").exists():
raise ValueError(
f"{output}/best.pt already exists. Choose a new --output to preserve this run."
)
output.mkdir(parents=True, exist_ok=True)
cells = 1 if tiles == 9 else 16
kind = "tile" if tiles == 9 else "full"
classes = sorted(
{r["target"] for r in manifest["records"] if r["tiles"] == tiles}
| {t for r in folder_records(folders, kind, "train") for t in r["targets"]} - {BACKGROUND}
)
aux_targets = sorted({target_name(name) for name in args.aux_targets})
if aux_targets and (tiles != 9 or not set(aux_targets) <= set(classes)):
raise ValueError("Auxiliary targets must be existing classes in a 9-tile model")
# 4x4 photos are 450 px; 448 keeps small slivers at a size divisible by the stride.
size = args.size or (224 if tiles == 9 else 448)
backbone_lr = args.backbone_lr if args.backbone_lr is not None else 0.0001
finetune_head_lr = args.finetune_head_lr if args.finetune_head_lr is not None else 0.0003
config = {
"tiles": tiles,
"classes": classes,
"size": size,
"backbone": args.backbone,
"seed": args.seed,
"batch_size": args.batch_size,
"device": str(device),
"precision": "bfloat16 autocast" if args.amp else "float32",
"frozen_epochs": args.frozen_epochs,
"finetune_epochs": args.epochs,
"patience": args.patience,
"optimizer": "AdamW",
"weight_decay": 0.0001,
"head_lr": args.head_lr,
"finetune_head_lr": finetune_head_lr,
"finetune_backbone_lr": backbone_lr,
"warm_start": args.warm_start,
"init_backbone": args.init_backbone,
"aux_targets": aux_targets,
"folders": args.folders,
"folder_kind": kind if folders else None,
"folder_negative_weight": args.negative_weight if folders else None,
"old_fraction": args.old_fraction if folders else None,
"epoch_size": args.epoch_size,
"keep_uncertain": args.keep_uncertain,
"train_batchnorm": not args.freeze_bn,
"amp": args.amp,
"selection": (
"mean of answer-key and folder validation loss"
if folders
else "mean of main and auxiliary validation loss"
if aux_targets
else "main validation loss"
),
"torch": str(torch.__version__),
}
if args.amp and device.type != "cuda":
raise ValueError("--amp is supported only with --device cuda")
if device.type == "cuda":
torch.backends.cudnn.benchmark = True
torch.set_float32_matmul_precision("high")
save_json(output / "config.json", config)
save_json(output / "classes.json", classes)
save_json(output / "split.json", manifest)
training_records = records_for(manifest, tiles, "train") + auxiliary_records(
manifest, aux_targets, "train"
)
training = SupervisedGrids(
GridDataset(
manifest["root"],
training_records,
classes,
size,
training=True,
tiles=tiles == 9,
include_full_images=bool(aux_targets),
),
cells,
keep_uncertain=args.keep_uncertain,
flip=not args.no_flip,
)
validation = GridDataset(manifest["root"], records_for(manifest, tiles, "val"), classes, size)
workers = {
"num_workers": args.workers,
"persistent_workers": args.workers > 0,
"pin_memory": device.type == "cuda",
}
sampler, folder_loader = None, None
if folders:
folder_training = FolderDataset(
folders["root"],
folder_records(folders, kind, "train"),
classes,
size,
cells,
training=True,
negative_weight=args.negative_weight,
)
folder_validation = FolderDataset(
folders["root"],
folder_records(folders, kind, "val"),
classes,
size,
cells,
negative_weight=args.negative_weight,
)
# Square-root class balancing: rare classes are seen more often without memorizing a few images.
counts = {}
for record in folder_training.records:
counts[record["targets"][0]] = counts.get(record["targets"][0], 0) + 1
folder_weights = torch.tensor(
[counts[r["targets"][0]] ** -0.5 for r in folder_training.records]
)
folder_weights *= (1 - args.old_fraction) / folder_weights.sum()
old_weights = torch.full((len(training),), args.old_fraction / len(training))
sampler = WeightedRandomSampler(
torch.cat([old_weights, folder_weights]).double(),
args.epoch_size or len(training) + len(folder_training),
generator=torch.Generator().manual_seed(args.seed),
)
training = ConcatDataset([training, folder_training])
folder_loader = DataLoader(folder_validation, batch_size=64, **workers)
elif args.epoch_size:
sampler = RandomSampler(
training,
replacement=True,
num_samples=args.epoch_size,
generator=torch.Generator().manual_seed(args.seed),
)
train_loader = DataLoader(
training, batch_size=args.batch_size, shuffle=sampler is None, sampler=sampler, **workers
)
val_loader = DataLoader(validation, batch_size=16, **workers)
aux_loader = None
if aux_targets:
aux_validation = GridDataset(
manifest["root"],
auxiliary_records(manifest, aux_targets, "val"),
classes,
size,
tile_grids=True,
)
if not len(aux_validation):
raise ValueError("Auxiliary targets need validation examples")
aux_loader = DataLoader(aux_validation, batch_size=2)
if args.warm_start:
model, initial = load_checkpoint(args.warm_start, device)
config["backbone"] = initial.get("backbone", BACKBONE)
save_json(output / "config.json", config)
if initial["tiles"] != tiles or initial["classes"] != classes or initial["size"] != size:
raise ValueError("Warm-start model must match task, classes, and image size")
else:
model = GridModel(tiles, len(classes), backbone=args.backbone).to(device)
if args.init_backbone:
source = torch.load(args.init_backbone, map_location="cpu", weights_only=True)
print(
f"grid{tiles}: copied {transfer_backbone(model, source['model'])} backbone tensors "
f"from {args.init_backbone}",
flush=True,
)
model.freeze_backbone(True)
optimizer = torch.optim.AdamW(model.head.parameters(), lr=args.head_lr, weight_decay=0.0001)
history, best, stale = [], float("inf"), 0
start = time.perf_counter()
print(
f"grid{tiles}: {len(train_loader.dataset)} training items ({len(train_loader.sampler)} per epoch), "
f"{len(validation)} validation grids, device={device}",
flush=True,
)
for epoch in range(args.frozen_epochs + args.epochs):
frozen = epoch < args.frozen_epochs
if epoch == args.frozen_epochs:
model.freeze_backbone(False)
optimizer = torch.optim.AdamW(
[
{"params": model.backbone.parameters(), "lr": backbone_lr},
{"params": model.head.parameters(), "lr": finetune_head_lr},
],
weight_decay=0.0001,
)
stale = 0
model.train()
if args.freeze_bn:
# Keep pretrained batch-normalization statistics (for small datasets).
model.backbone.eval()
total_loss, supervised = 0.0, 0.0
epoch_start = time.perf_counter()
for images, *supervision in train_loader:
images = images.to(device, non_blocking=True)
supervision = [x.to(device, non_blocking=True) for x in supervision]
optimizer.zero_grad(set_to_none=True)
with torch.autocast("cuda", dtype=torch.bfloat16, enabled=args.amp):
logits = cell_logits(model, images)
loss, count = partial_loss(logits.float(), *supervision)
loss.backward()
optimizer.step()
total_loss += loss.item() * count.item()
supervised += count.item()
val_loss, probabilities, truth, _ = evaluate_loader(model, val_loader, device)
aux_loss = evaluate_loader(model, aux_loader, device)[0] if aux_loader is not None else None
folder_loss = (
evaluate_folders(model, folder_loader, device)[0] if folder_loader is not None else None
)
extra = [x for x in (aux_loss, folder_loss) if x is not None]
selection_loss = sum([val_loss] + extra) / (1 + len(extra))
row = {
"epoch": epoch + 1,
"phase": "head" if frozen else "finetune",
"train_loss": total_loss / max(supervised, 1),
"val_loss": val_loss,
"aux_val_loss": aux_loss,
"folder_val_loss": folder_loss,
"selection_loss": selection_loss,
"val_exact_at_0.5": metrics(probabilities, truth, 0.5)["exact_grid_accuracy"],
"seconds": time.perf_counter() - epoch_start,
}
history.append(row)
save_json(output / "history.json", history)
print(json.dumps({"grid": tiles, **row}), flush=True)
if selection_loss < best:
best, stale = selection_loss, 0
torch.save(
{
**config,
"epoch": epoch + 1,
"val_loss": val_loss,
"selection_loss": selection_loss,
"model": {k: v.detach().cpu().clone() for k, v in model.state_dict().items()},
},
output / "best.pt",
)
elif not frozen:
stale += 1
if not frozen and stale >= args.patience:
break
elapsed = time.perf_counter() - start
del optimizer, model
empty_cache(device)
model, checkpoint = load_checkpoint(output / "best.pt", device)
_, probabilities, truth, val_targets = evaluate_loader(model, val_loader, device)
threshold = choose_threshold(probabilities, truth)
checkpoint["threshold"] = threshold
checkpoint["thresholds"] = {}
if aux_loader is not None:
_, aux_probabilities, aux_truth, aux_ids = evaluate_loader(model, aux_loader, device)
for name in aux_targets:
mask = aux_ids == classes.index(name)
if mask.sum() >= 3:
checkpoint["thresholds"][name] = choose_threshold(
aux_probabilities[mask], aux_truth[mask]
)
save_json(
output / "aux_validation.json",
{
"source": "4x4 crops; proxy only for 3x3 performance",
"thresholds": checkpoint["thresholds"],
**metrics(aux_probabilities, aux_truth, thresholds_for(checkpoint, aux_ids)),
},
)
if folder_loader is not None:
_, folder_probabilities, folder_truth, folder_weights = evaluate_folders(
model, folder_loader, device
)
save_json(
output / "folder_validation.json",
{
"source": f"held-out {kind} images from {args.folders}; single-label, so 'fires_on_other_classes' "
"includes images where that target is genuinely visible",
**folder_report(
folder_probabilities, folder_truth, folder_weights, classes, threshold
),
},
)
checkpoint["training_seconds"] = elapsed
torch.save(checkpoint, output / "best.pt")
save_json(
output / "validation.json",
{
"threshold": threshold,
"thresholds": checkpoint["thresholds"],
**metrics(probabilities, truth, thresholds_for(checkpoint, val_targets)),
},
)
print(
f"grid{tiles}: saved best epoch {checkpoint['epoch']}; threshold={threshold}; training={elapsed:.1f}s",
flush=True,
)
def synchronize(device):
if device.type == "mps":
torch.mps.synchronize()
elif device.type == "cuda":
torch.cuda.synchronize()
def calibrate_crops(args):
"""Enable enlarged views only after a validation-only comparison."""
device = device_for(args.device)
seed_everything(42)
model, checkpoint = load_checkpoint(args.checkpoint, device)
name = target_name(args.target)
if checkpoint["tiles"] != 9 or name not in checkpoint["classes"]:
raise ValueError("Choose a supported target in a tile classifier")
if name in checkpoint.get("crop_targets", []):
raise ValueError("Crop views are already enabled for this target")
directory = Path(args.checkpoint).parent
manifest = json.loads((directory / "split.json").read_text())
records = auxiliary_records(manifest, [name], "val")
if len(records) < 3:
raise ValueError("Need at least three auxiliary validation grids")
dataset = GridDataset(
manifest["root"], records, checkpoint["classes"], checkpoint["size"], tile_grids=True
)
loader = DataLoader(dataset, batch_size=2)
_, probabilities, truth, _ = evaluate_loader(model, loader, device)
old_threshold = checkpoint.get("thresholds", {}).get(name, checkpoint["threshold"])
before = metrics(probabilities, truth, old_threshold)
model.crop_target_ids.append(checkpoint["classes"].index(name))
_, probabilities, truth, _ = evaluate_loader(model, loader, device)
cutoff = choose_threshold(probabilities, truth)
after = metrics(probabilities, truth, cutoff)
accepted = (
after["exact_grid_accuracy"] >= before["exact_grid_accuracy"]
and after["cell_f1"] > before["cell_f1"]
)
report = {
"target": name,
"source": "4x4 validation crops; not 3x3 test data",
"views": "full tile plus center and four corner half-size crops",
"before": {"threshold": old_threshold, **before},
"after": {"threshold": cutoff, **after},
"accepted": accepted,
}
if accepted:
backup = directory / "pre_crop_best.pt"
if backup.exists():
raise ValueError(f"Refusing to overwrite {backup}")
torch.save(checkpoint, backup)
checkpoint.setdefault("crop_targets", []).append(name)
checkpoint.setdefault("thresholds", {})[name] = cutoff
torch.save(checkpoint, args.checkpoint)
save_json(directory / "crop_calibration.json", report)
print(json.dumps(report, indent=2))
@torch.inference_mode()
def benchmark(model, dataset, device):
images, target, _ = dataset[0]
images, target = images.unsqueeze(0).to(device), torch.tensor([target], device=device)
for _ in range(5):
grid_logits(model, images, target)
synchronize(device)
timings = []
for _ in range(30):
start = time.perf_counter()
grid_logits(model, images, target).sigmoid().cpu()
synchronize(device)
timings.append((time.perf_counter() - start) * 1000)
return {
"median_ms": float(np.median(timings)),
"p95_ms": float(np.percentile(timings, 95)),
"scope": "one grid, warm model, includes output transfer; excludes image loading and preprocessing",
"device": str(device),
"iterations": len(timings),
}
def evaluate(args):
device = device_for(args.device)
seed_everything(42)
model, checkpoint = load_checkpoint(args.checkpoint, device)
directory = Path(args.checkpoint).parent
manifest = json.loads((directory / "split.json").read_text())
records = records_for(manifest, checkpoint["tiles"], args.split)
if not records:
raise ValueError(
f"No {args.split} grids in this manifest"
+ (" (run prepare with --review)" if args.split == "review" else "")
)
dataset = GridDataset(
args.data or manifest["root"], records, checkpoint["classes"], checkpoint["size"]
)
loss, probabilities, truth, targets = evaluate_loader(
model, DataLoader(dataset, batch_size=4), device
)
threshold = thresholds_for(checkpoint, targets)
report = {
"split": args.split,
"epoch": checkpoint["epoch"],
"threshold": checkpoint["threshold"],
"thresholds": checkpoint.get("thresholds", {}),
"loss": loss,
"training_seconds": checkpoint["training_seconds"],
"overall": metrics(probabilities, truth, threshold),
"per_class": {},
}
if args.split == "review":
# Same grids, answered by the labeler whose output the model was trained on.
right = [r["labeler_answers"] == r["answers"] for r in records if "labeler_answers" in r]
report["labeler_exact_grid_accuracy"] = sum(right) / len(right) if right else None
for i, name in enumerate(checkpoint["classes"]):
mask = targets == i
support = {
split: len(
[
r
for r in records_for(manifest, checkpoint["tiles"], split)
if r["target"] == name
]
)
for split in ("train", "val", "test", "review")
}
report["per_class"][name] = {
"support": support,
"metrics": metrics(probabilities[mask], truth[mask], threshold[mask])
if mask.any()
else None,
}
report["latency"] = benchmark(model, dataset, device)
save_json(directory / f"{args.split}_metrics.json", report)
save_json(
directory / f"{args.split}_predictions.json",
[
{
"id": r["id"],
"target": r["target"],
"answers": r["answers"],
"selected_cells": (p >= cutoff).nonzero().flatten().tolist(),
"probabilities": p.tolist(),
}
for r, p, cutoff in zip(records, probabilities, threshold)
],
)
print(json.dumps(report, indent=2))
@torch.inference_mode()
def predict(args):
device = device_for(args.device)
model, checkpoint = load_checkpoint(args.checkpoint, device)
name = target_name(args.target)
if name not in checkpoint["classes"]:
raise ValueError(f"Unknown target {name!r}; supported: {checkpoint['classes']}")
preprocess = transform(checkpoint["size"])
image = read_rgb(image_file(args.image))
if checkpoint["tiles"] == 9:
images = torch.stack([preprocess(image.crop(box)) for box in grid_boxes(*image.size, 9)])
else:
images = preprocess(image)
target = torch.tensor([checkpoint["classes"].index(name)], device=device)
probabilities = grid_logits(model, images.unsqueeze(0).to(device), target).sigmoid().cpu()[0]
cutoff = checkpoint.get("thresholds", {}).get(name, checkpoint["threshold"])
print(
json.dumps(
{
"target": name,
"threshold": cutoff,
"selected_cells": (probabilities >= cutoff).nonzero().flatten().tolist(),
"probabilities": probabilities.tolist(),
},
indent=2,
)
)
def export(args):
"""ONNX fp32 plus calibrated int8, checked against PyTorch on the old test grids."""
import onnx
from onnxruntime.quantization import (
CalibrationDataReader,
CalibrationMethod,
QuantFormat,
QuantType,
quantize_static,
)
from onnxruntime.quantization.shape_inference import quant_pre_process
from predict_onnx import GridPredictor
seed_everything(42)
model, checkpoint = load_checkpoint(args.checkpoint, torch.device("cpu"))
if model.crop_target_ids:
raise ValueError("Enlarged crop views are not supported by the ONNX export")
tiles, size = checkpoint["tiles"], checkpoint["size"]
directory = Path(args.checkpoint).parent
output = Path(args.output or directory / "onnx")
output.mkdir(parents=True, exist_ok=True)
manifest = json.loads((directory / "split.json").read_text())
root = Path(args.data or manifest["root"])
class Probabilities(nn.Module):
def __init__(self):
super().__init__()
self.model = model
def forward(self, images):
logits = self.model(images)
return (logits if tiles == 9 else logits.flatten(2)).sigmoid()
info = {
"tiles": tiles,
"size": size,
"classes": checkpoint["classes"],
"threshold": checkpoint["threshold"],
"thresholds": checkpoint.get("thresholds", {}),
"epoch": checkpoint["epoch"],
}
def with_metadata(path):
exported = onnx.load(path)
del exported.metadata_props[:]
exported.metadata_props.add(key="grid_vision", value=json.dumps(info))
onnx.save(exported, path)
fp32, prepared, int8 = (
output / f"grid{tiles}.onnx",
output / f"grid{tiles}.pre.onnx",
output / f"grid{tiles}.int8.onnx",
)
for path in (fp32, int8):
if path.exists():
raise ValueError(f"{path} already exists. Choose a new --output.")
torch.onnx.export(
Probabilities().eval(),
torch.zeros(9 if tiles == 9 else 1, 3, size, size),
fp32,
input_names=["images"],
output_names=["probabilities"],
opset_version=17,
dynamic_axes={"images": {0: "batch"}, "probabilities": {0: "batch"}},
dynamo=False,
)
with_metadata(fp32)
reference = GridPredictor(fp32)
calibration = random.Random(0).sample(
records_for(manifest, tiles, "train"),
min(args.calibration_grids, len(records_for(manifest, tiles, "train"))),
)
class Calibration(CalibrationDataReader):
def __init__(self):
self.batches = iter(reference.inputs(record_image_path(root, r)) for r in calibration)
def get_next(self):
batch = next(self.batches, None)
return None if batch is None else {"images": batch}
quant_pre_process(str(fp32), str(prepared))
# Percentile ranges and full-precision first/last layers changed the fewest baseline decisions.
layers = [
n.name for n in onnx.load(prepared).graph.node if n.op_type in ("Conv", "Gemm", "MatMul")
]
quantize_static(
str(prepared),
str(int8),
Calibration(),
quant_format=QuantFormat.QDQ,
per_channel=True,
activation_type=QuantType.QUInt8,
weight_type=QuantType.QInt8,
calibrate_method=CalibrationMethod.Percentile,
nodes_to_exclude=[layers[0], layers[-1]],
)
prepared.unlink()
with_metadata(int8)
records = records_for(manifest, tiles, "test")
variants = {"fp32": GridPredictor(fp32), "int8": GridPredictor(int8)}
probabilities = {name: [] for name in ["torch", *variants]}
labels, cutoffs = [], []
with torch.inference_mode():
for record in records:
inputs = reference.inputs(record_image_path(root, record))
target = checkpoint["classes"].index(record["target"])
scores = Probabilities()(torch.from_numpy(inputs)).numpy()
probabilities["torch"].append(scores.T[target] if tiles == 9 else scores[0, target])
for name, predictor in variants.items():
probabilities[name].append(predictor.probabilities(inputs)[target])
labels.append([i in record["answers"] for i in range(record["tiles"])])
cutoffs.append(info["thresholds"].get(record["target"], info["threshold"]))
labels, cutoffs = torch.tensor(labels), torch.tensor(cutoffs).unsqueeze(1)
tensors = {name: torch.tensor(np.array(values)) for name, values in probabilities.items()}
selected = {name: values >= cutoffs for name, values in tensors.items()}
report = {
"checkpoint": str(args.checkpoint),
"test_grids": len(records),
"preprocessing": "predict_onnx.preprocess (Pillow + numpy) for every variant",
}
for name, path in [("fp32", fp32), ("int8", int8)]:
timings = {}
sample = reference.inputs(record_image_path(root, records[0]))
for threads in (1, 4):
predictor = GridPredictor(path, threads)
for _ in range(5):
predictor.probabilities(sample)
runs = []
for _ in range(50):
start = time.perf_counter()
predictor.probabilities(sample)
runs.append((time.perf_counter() - start) * 1000)
timings[f"{threads}_thread_median_ms"] = float(np.median(runs))
report[name] = {
"file": str(path),
"megabytes": path.stat().st_size / 1e6,
"max_abs_difference_from_torch": float((tensors[name] - tensors["torch"]).abs().max()),
"grids_with_different_selection": int(
(selected[name] != selected["torch"]).any(1).sum()
),
"test": metrics(tensors[name], labels, cutoffs),
"cpu_latency_per_grid": {
**timings,
"scope": "one grid on this machine's CPU, excludes image loading",
},
}
report["torch_test"] = metrics(tensors["torch"], labels, cutoffs)
save_json(output / "export.json", report)
print(json.dumps(report, indent=2))
def seam_ratio(image, parts=3):
"""Separate tiles have sharp seams at the 1/parts lines; one continuous photo usually does not."""
pixels = np.asarray(image.convert("L"), dtype=np.float32)
height, width = pixels.shape
dx, dy = np.abs(np.diff(pixels, axis=1)).mean(0), np.abs(np.diff(pixels, axis=0)).mean(1)
def seams(differences, length):
return max(
differences[max(0, round(length * j / parts) - 2) : round(length * j / parts) + 1].max()
for j in range(1, parts)
)
return min(seams(dx, width), seams(dy, height)) / ((np.median(dx) + np.median(dy)) / 2 + 1e-6)
def image_file(source):
"""A challenge image, or the image inside a dataset item folder."""
source = Path(source)
if source.is_dir():
found = next((source / name for name in IMAGE_NAMES if (source / name).exists()), None)
if found is None:
raise ValueError(f"No image in {source}")
return found
return source
def layout(source):
"""9 or 16 tiles for a challenge image or item folder (300 px = 3x3, 450 px = 4x4)."""
metadata = Path(source) / "metadata.json"
if Path(source).is_dir() and metadata.is_file():
grid = json.loads(metadata.read_text()).get("grid")
if grid in ("3x3", "4x4"):
return 9 if grid == "3x3" else 16
image = read_rgb(image_file(source))
side = max(image.size)
if abs(side - 300) <= 30:
return 9
if abs(side - 450) <= 45:
return 16
return 9 if seam_ratio(image) >= 3 else 16
def solve(args):
"""Detect the layout and use the matching grid model."""
args.checkpoint = args.grid9 if layout(args.image) == 9 else args.grid16
predict(args)
def main():
parser = argparse.ArgumentParser(description=__doc__)
commands = parser.add_subparsers(dest="command", required=True)
prep = commands.add_parser("prepare")
prep.add_argument(
"--data",
required=True,
help="Folder of labeled <item>/{image.jpg, metadata.json} (or answer-key) folders",
)
prep.add_argument("--output", default="runs")
prep.add_argument("--seed", type=int, default=42)
prep.add_argument(
"--workers", type=int, default=os.cpu_count() or 1, help="Processes for scanning images"
)
prep.add_argument(
"--sam",
action="store_true",
help="Take 4x4 answers from <item>/sam.json (SAM 3) instead of metadata.json",
)
prep.add_argument(
"--review",
nargs="*",
default=[],
help="Reference answer files held out as a 'review' split",
)
prep_folders = commands.add_parser("prepare-folders", help="Ingest class-folder image datasets")
prep_folders.add_argument("--data", nargs="+", required=True, help="Class-folder datasets")
prep_folders.add_argument(
"--manifest", default="runs/split.json", help="Old split; its held-out images are excluded"
)
prep_folders.add_argument("--output", default="runs")
prep_folders.add_argument("--seed", type=int, default=42)
train = commands.add_parser("train")
train.add_argument("--manifest", default="runs/split.json")
train.add_argument("--output", default="runs/baseline")
train.add_argument("--tiles", choices=["9", "16", "both"], default="both")
train.add_argument("--device", choices=["auto", "cuda", "mps", "cpu"], default="auto")
train.add_argument("--batch-size", type=int, default=16)
train.add_argument("--frozen-epochs", type=int, default=3)
train.add_argument(
"--epochs", type=int, default=20, help="Maximum fine-tuning epochs after head training"
)
train.add_argument("--patience", type=int, default=5)
train.add_argument("--backbone-lr", type=float, help="Default: 1e-4")
train.add_argument("--finetune-head-lr", type=float, help="Default: 3e-4")
train.add_argument("--head-lr", type=float, default=0.001)
train.add_argument("--warm-start", help="Initialize from a compatible trained checkpoint")
train.add_argument(
"--aux-targets",
nargs="*",
default=[],
help="Use these targets' 4x4 crops to train a 9-tile classifier",
)
train.add_argument("--seed", type=int, default=42)
train.add_argument("--size", type=int, help="Input size. Default: 224 for 9 tiles, 448 for 16")
train.add_argument(
"--folders", help="Folder manifest from prepare-folders, e.g. runs/folders.json"
)
train.add_argument(
"--negative-weight",
type=float,
default=0.2,
help="Weight for assuming other targets are absent in a single-label folder image",
)
train.add_argument(
"--old-fraction", type=float, default=0.25, help="Share of samples from answer-key grids"
)
train.add_argument("--epoch-size", type=int, help="Training samples (tiles or grids) per epoch")
train.add_argument(
"--keep-uncertain",
action="store_true",
help="Also train on cells the labeler marked uncertain (ignored by default)",
)
train.add_argument(
"--no-flip", action="store_true", help="Turn off random left-right mirroring"
)
train.add_argument(
"--backbone",
default=BACKBONE,
help="Any timm model with a stride-16 stage, e.g. convnext_tiny.dinov3_lvd1689m",
)
train.add_argument(
"--init-backbone", help="Copy backbone weights from another checkpoint, e.g. the 3x3 model"
)
train.add_argument(
"--freeze-bn", action="store_true", help="Keep pretrained batch-normalization statistics"
)
train.add_argument(
"--workers",
type=int,
default=min(16, os.cpu_count() or 1),
help="DataLoader worker processes",
)
train.add_argument("--amp", action="store_true", help="bfloat16 autocast (CUDA only)")
test = commands.add_parser("evaluate")
test.add_argument("--checkpoint", required=True)
test.add_argument(
"--split",
choices=["test", "review"],
default="test",
help="review: the reference grids held out by prepare --review",
)
test.add_argument("--data", help="Override dataset root if moved")
test.add_argument("--device", choices=["auto", "cuda", "mps", "cpu"], default="auto")
prediction = commands.add_parser("predict")
prediction.add_argument("--checkpoint", required=True)
prediction.add_argument("--image", required=True, help="Image file or example folder")
prediction.add_argument("--target", required=True)
prediction.add_argument("--device", choices=["auto", "cuda", "mps", "cpu"], default="auto")
onnx_export = commands.add_parser(
"export", help="ONNX fp32 and int8 for CPU inference without PyTorch"
)
onnx_export.add_argument("--checkpoint", required=True)
onnx_export.add_argument("--output", help="Default: <checkpoint folder>/onnx")
onnx_export.add_argument("--data", help="Override dataset root if moved")
onnx_export.add_argument("--calibration-grids", type=int, default=100)
solver = commands.add_parser("solve", help="Detect the layout and use the matching model")
solver.add_argument("--image", required=True, help="Image file or example folder")
solver.add_argument("--target", required=True)
solver.add_argument("--grid9", default="runs/baseline/grid9/best.pt")
solver.add_argument("--grid16", default="runs/baseline/grid16/best.pt")
solver.add_argument("--device", choices=["auto", "cuda", "mps", "cpu"], default="cpu")
crops = commands.add_parser("calibrate-crops")
crops.add_argument("--checkpoint", required=True)
crops.add_argument("--target", required=True)
crops.add_argument("--device", choices=["auto", "cuda", "mps", "cpu"], default="auto")
args = parser.parse_args()
if args.command == "prepare":
source = Path(args.data)
if (source / "data" / "3x3").is_dir() or (source / "3x3").is_dir():
if args.sam or args.review:
parser.error("--sam and --review apply only to raw labeled folders")
_, counts = prepare_jsonl(source, args.output)
else:
_, counts = prepare(args.data, args.output, args.seed, args.workers, args.sam, args.review)
print(json.dumps(counts, indent=2))
elif args.command == "prepare-folders":
_, summary = prepare_folders(args.data, args.manifest, args.output, args.seed)
print(json.dumps(summary, indent=2))
elif args.command == "train":
if min(args.batch_size, args.epochs, args.patience) < 1 or args.frozen_epochs < 0:
parser.error(
"Batch size, epochs, and patience must be positive; frozen epochs cannot be negative"
)
if any(
lr is not None and (not np.isfinite(lr) or lr <= 0)
for lr in (args.backbone_lr, args.finetune_head_lr, args.head_lr)
):
parser.error("Learning rates must be finite and positive")
if not 0 < args.old_fraction < 1 or args.negative_weight < 0 or args.workers < 0:
parser.error("--old-fraction must be in (0, 1); weights and workers cannot be negative")
if not np.isfinite(args.negative_weight):
parser.error("--negative-weight must be finite")
if (args.epoch_size is not None and args.epoch_size < 1) or (
args.size is not None and args.size < 1
):
parser.error("--epoch-size and --size must be positive")
manifest = json.loads(Path(args.manifest).read_text())
folders = json.loads(Path(args.folders).read_text()) if args.folders else None
for tiles in [9, 16] if args.tiles == "both" else [int(args.tiles)]:
train_one(args, tiles, manifest, folders)
elif args.command == "evaluate":
evaluate(args)
elif args.command == "export":
if args.calibration_grids < 1:
parser.error("--calibration-grids must be positive")
export(args)
elif args.command == "solve":
solve(args)
elif args.command == "calibrate-crops":
calibrate_crops(args)
else:
predict(args)
if __name__ == "__main__":
main()
"""Prepare, train, evaluate, and predict. See README.md for commands."""
import argparse
import json
import os
import random
import time
from pathlib import Path
# Keep downloaded weights with the project, not in a global cache.
os.environ.setdefault("HF_HOME", str(Path(".cache/huggingface").resolve()))
import numpy as np
import torch
from torch import nn
from torch.utils.data import ConcatDataset, DataLoader, RandomSampler, WeightedRandomSampler
from data import (
BACKGROUND,
IMAGE_NAMES,
FolderDataset,
GridDataset,
SupervisedGrids,
grid_boxes,
prepare,
prepare_jsonl,
prepare_folders,
read_rgb,
record_image_path,
target_name,
transform,
)
from models import BACKBONE, GridModel, select_target, transfer_backbone
def save_json(path, value):
Path(path).write_text(json.dumps(value, indent=2) + "\n")
def device_for(name):
if name == "auto":
name = (
"cuda"
if torch.cuda.is_available()
else "mps"
if torch.backends.mps.is_available()
else "cpu"
)
if name == "mps" and not torch.backends.mps.is_available():
raise RuntimeError(
"MPS is unavailable. Run in a terminal with Apple GPU access; use --device cpu only intentionally."
)
if name == "cuda" and not torch.cuda.is_available():
raise RuntimeError("CUDA is unavailable on this machine.")
return torch.device(name)
def empty_cache(device):
if device.type == "mps":
torch.mps.empty_cache()
elif device.type == "cuda":
torch.cuda.empty_cache()
def seed_everything(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.set_num_threads(4)
def records_for(manifest, tiles, split):
return [r for r in manifest["records"] if r["tiles"] == tiles and r["split"] == split]
def auxiliary_records(manifest, targets, split):
return [r for r in records_for(manifest, 16, split) if r["target"] in targets]
def folder_records(folders, kind, split):
return (
[r for r in folders["records"] if r["kind"] == kind and r["split"] == split]
if folders
else []
)
def cell_logits(model, images):
"""Logits shaped (batch, classes, cells); a separate tile is one cell."""
output = model(images)
return output.unsqueeze(-1) if model.tiles == 9 else output.flatten(2)
def partial_loss(logits, labels, weights, any_labels, any_weights):
"""Weighted BCE on known cells, plus presence anywhere (max over cells) for image-level labels."""
cells = nn.functional.binary_cross_entropy_with_logits(
logits, labels, weight=weights, reduction="sum"
)
anywhere = nn.functional.binary_cross_entropy_with_logits(
logits.amax(-1), any_labels, weight=any_weights, reduction="sum"