#!/usr/bin/env python3 """Fine-tune facebook/detr-resnet-50 (Apache-2.0) on biglam/loc_beyond_words (7 classes).""" import argparse import json import os import random import torch from datasets import load_dataset from torch.utils.data import DataLoader, Dataset, Subset import torchmetrics from transformers import AutoProcessor, DetrForObjectDetection, get_scheduler CLASSES = ["Photograph", "Illustration", "Map", "Comics/Cartoon", "Editorial Cartoon", "Headline", "Advertisement"] class DetrDataset(Dataset): def __init__(self, hf_ds, processor): self.ds = hf_ds self.processor = processor def __len__(self): return len(self.ds) def __getitem__(self, idx): ex = self.ds[idx] x, y, w, h = ex["width"], ex["height"], None, None annotations = [] for o in ex["objects"]: bx, by, bw, bh = [float(v) for v in o["bbox"]] annotations.append({ "bbox": [bx, by, bw, bh], "category_id": o["category_id"], "area": float(bw * bh), "iscrowd": o["iscrowd"], "id": o["id"], }) target = {"image_id": idx, "annotations": annotations} encoding = self.processor(images=ex["image"], annotations=target, return_tensors="pt") return { "pixel_values": encoding["pixel_values"].squeeze(0), "labels": encoding["labels"][0], "height": ex["height"], "width": ex["width"], } def collate_fn(batch, processor): pvs = [item["pixel_values"] for item in batch] max_h = max(pv.shape[1] for pv in pvs) max_w = max(pv.shape[2] for pv in pvs) bs = len(batch) pix = torch.zeros(bs, 3, max_h, max_w) mask = torch.zeros(bs, max_h, max_w, dtype=torch.int64) for i, pv in enumerate(pvs): h, w = pv.shape[1], pv.shape[2] pix[i, :, :h, :w] = pv mask[i, :h, :w] = 1 return { "pixel_values": pix, "pixel_mask": mask, "labels": [item["labels"] for item in batch], "height": [item["height"] for item in batch], "width": [item["width"] for item in batch], } @torch.no_grad() def evaluate(model, processor, loader, device, threshold=0.0): model.eval() try: metric = torchmetrics.detection.MeanAveragePrecision( iou_type="bbox", class_metrics=True, extended_summary=True, backend="faster_coco_eval") except TypeError: metric = torchmetrics.detection.MeanAveragePrecision(iou_type="bbox", class_metrics=True, extended_summary=True) for batch in loader: pv = batch["pixel_values"].to(device) pm = batch["pixel_mask"].to(device) out = model(pixel_values=pv, pixel_mask=pm) target_sizes = torch.tensor([[h, w] for h, w in zip(batch["height"], batch["width"])]) preds = processor.post_process_object_detection(out, threshold=threshold, target_sizes=target_sizes) for i in range(len(preds)): pred = preds[i] tar = batch["labels"][i] image_size = torch.tensor([batch["height"][i], batch["width"][i]], dtype=torch.float) # processor labels.boxes are normalized cxcywh-ish; convert to absolute xyxy tboxes = tar["boxes"] # boxes from processor are in [cx,cy,w,h] normalized 0..1 cx, cy, w, h = tboxes[:, 0] * image_size[1], tboxes[:, 1] * image_size[0], tboxes[:, 2] * image_size[1], tboxes[:, 3] * image_size[0] xyxy = torch.stack([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2], dim=1) metric.update( [{"boxes": pred["boxes"].cpu(), "scores": pred["scores"].cpu(), "labels": pred["labels"].cpu()}], [{"boxes": xyxy, "labels": tar["class_labels"]}], ) res = metric.compute() out = { "eval_map": float(res["map"]), "eval_map_50": float(res["map_50"]), "eval_map_75": float(res["map_75"]), "per_class_map_50": [float(x) for x in res.get("map_50_per_class", [0.0] * 7)], } return out def main(): ap = argparse.ArgumentParser() ap.add_argument("--epochs", type=int, default=14) ap.add_argument("--batch", type=int, default=2) ap.add_argument("--acc", type=int, default=4) ap.add_argument("--lr", type=float, default=1e-4) ap.add_argument("--backbone_lr", type=float, default=1e-5) ap.add_argument("--max_train", type=int, default=0) ap.add_argument("--max_eval", type=int, default=0) ap.add_argument("--repo", type=str, default="harness-race/opencode-r1") ap.add_argument("--push", action="store_true") ap.add_argument("--outjson", type=str, default="val_results.json") args = ap.parse_args() model_id = "facebook/detr-resnet-50" id2label = {i: c for i, c in enumerate(CLASSES)} label2id = {c: i for i, c in enumerate(CLASSES)} processor = AutoProcessor.from_pretrained(model_id) ds = load_dataset("biglam/loc_beyond_words") train_ds = DetrDataset(ds["train"], processor) eval_ds = DetrDataset(ds["validation"], processor) random.seed(0) if args.max_train: train_ds = Subset(train_ds, random.sample(range(len(train_ds)), min(args.max_train, len(train_ds)))) if args.max_eval: eval_ds = Subset(eval_ds, random.sample(range(len(eval_ds)), min(args.max_eval, len(eval_ds)))) model = DetrForObjectDetection.from_pretrained( model_id, num_labels=len(CLASSES), ignore_mismatched_sizes=True, id2label=id2label, label2id=label2id) device = "cuda" if torch.cuda.is_available() else "cpu" model = model.to(device) train_loader = DataLoader(train_ds, batch_size=args.batch, shuffle=True, collate_fn=lambda b: collate_fn(b, processor), num_workers=2, pin_memory=False) eval_loader = DataLoader(eval_ds, batch_size=args.batch, shuffle=False, collate_fn=lambda b: collate_fn(b, processor), num_workers=2, pin_memory=False) param_groups = [ {"params": [p for n, p in model.named_parameters() if "backbone" in n], "lr": args.backbone_lr}, {"params": [p for n, p in model.named_parameters() if "backbone" not in n], "lr": args.lr}, ] optimizer = torch.optim.AdamW(param_groups, lr=args.lr, weight_decay=1e-4) steps_per_epoch = len(train_loader) // args.acc num_steps = steps_per_epoch * args.epochs scheduler = get_scheduler("cosine", optimizer=optimizer, num_warmup_steps=int(0.05 * num_steps), num_training_steps=num_steps) scaler = torch.cuda.amp.GradScaler(enabled=(device == "cuda")) best_metric = -1.0 best_state = None best_map50 = 0.0 results_log = [] for epoch in range(1, args.epochs + 1): model.train() optimizer.zero_grad() running = 0.0 for step, batch in enumerate(train_loader): pv = batch["pixel_values"].to(device) pm = batch["pixel_mask"].to(device) labels = [{k: v.to(device) if torch.is_tensor(v) else v for k, v in t.items()} for t in batch["labels"]] with torch.cuda.amp.autocast(enabled=(device == "cuda")): out = model(pixel_values=pv, pixel_mask=pm, labels=labels) loss = out.loss / args.acc scaler.scale(loss).backward() running += float(out.loss.item()) if (step + 1) % args.acc == 0: scaler.step(optimizer) scaler.update() scheduler.step() optimizer.zero_grad() # trailing scaler.step(optimizer); scaler.update(); optimizer.zero_grad() print(f"[epoch {epoch}] train_loss={running / len(train_loader):.4f}", flush=True) res = evaluate(model, processor, eval_loader, device) results_log.append({**res, "epoch": epoch}) print(f"[epoch {epoch}] val map={res['eval_map']:.4f} map50={res['eval_map_50']:.4f}", flush=True) with open(args.outjson, "w") as f: json.dump(results_log, f) key = res["eval_map"] if key > best_metric: best_metric = key best_map50 = res["eval_map_50"] best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} torch.save(best_state, "best_model.pt") print(f"[epoch {epoch}] new best map={best_metric:.4f}", flush=True) # final best eval detailed model.load_state_dict(torch.load("best_model.pt", map_location=device)) res = evaluate(model, processor, eval_loader, device) print("BEST EVAL:", json.dumps(res)) final = { "eval_map": best_metric, "eval_map_50": best_map50, "per_class_map_50": { c: round(v, 4) for c, v in zip(CLASSES, res["per_class_map_50"]) }, "epochs": args.epochs, "train_batches_seen": epoch, "val_rows": len(eval_ds), } with open(args.outjson, "w") as f: json.dump(final, f, indent=2) if args.push: os.environ.setdefault("HF_TOKEN", os.environ.get("HF_TOKEN", "")) model.push_to_hub(args.repo) processor.push_to_hub(args.repo) from huggingface_hub import HfApi api = HfApi() api.upload_file(path_or_fileobj=build_readme(final).encode(), path_in_repo="README.md", repo_id=args.repo) if os.path.exists(args.outjson): api.upload_file(path_or_fileobj=open(args.outjson, "rb").read(), path_in_repo=os.path.basename(args.outjson), repo_id=args.repo) print("PUSHED to", args.repo) def build_readme(final): rows = "\n".join(f" - {c}: mAP@50 = **{v:.3f}**" for c, v in final["per_class_map_50"].items()) return f"""--- license: apache-2.0 tags: - object-detection - detr pipeline_tag: object-detection datasets: - biglam/loc_beyond_words metrics: - mean_average_precision --- # opencode-r1 — Object Detection on LOC Beyond Words Fine-tuned **facebook/detr-resnet-50** (DETR, ResNet-50 backbone, **Apache-2.0**) on the [`biglam/loc_beyond_words`](https://huggingface.co/datasets/biglam/loc_beyond_words) dataset — a crowdsourced collection of bounding-box annotations over WWI-era newspaper pages from the Library of Congress Chronicling America collection. Fine-tuning was performed on a single NVIDIA T4 via Hugging Face Jobs (~under \$5 of compute). ## Classes (7) {chr(10).join('- ' + c for c in CLASSES)} ## Validation results (COCO-style AP on 712 held-out images) - **mAP@0.5:0.95** = `{final['eval_map']:.4f}` - **mAP@0.5** = `{final['eval_map_50']:.4f}` Per-class mAP@0.5: {rows} ## Usage ```python from transformers import AutoProcessor, DetrForObjectDetection import torch processor = AutoProcessor.from_pretrained("harness-race/opencode-r1") model = DetrForObjectDetection.from_pretrained("harness-race/opencode-r1") image = Image.open("page.jpg") inputs = processor(images=image, return_tensors="pt") outputs = model(**inputs) results = processor.post_process_object_detection( outputs, threshold=0.5, target_sizes=torch.tensor([image.size[::-1]]))[0] ``` ## License & attribution - Base model `facebook/detr-resnet-50`: **Apache-2.0** - Dataset `biglam/loc_beyond_words`: **CC0-1.0** (public domain) - This fine-tuned model: **Apache-2.0** """ if __name__ == "__main__": main()