train_loc.py

#!/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()