train.py

# /// script
# dependencies = [
#   "torch",
#   "transformers",
#   "datasets",
#   "torchvision",
#   "torchmetrics",
#   "pillow",
#   "numpy",
#   "safetensors",
#   "huggingface_hub",
#   "tqdm",
# ]
# ///
import os, io, json, math, argparse, random, time
import numpy as np
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from transformers import AutoImageProcessor, DetrForObjectDetection
from datasets import load_dataset
from torchmetrics.detection.mean_ap import MeanAveragePrecision
from huggingface_hub import HfApi
from tqdm import tqdm

ID2LABEL = {0:'Photograph',1:'Illustration',2:'Map',3:'Comics/Cartoon',4:'Editorial Cartoon',5:'Headline',6:'Advertisement'}
LABEL2ID = {v:k for k,v in ID2LABEL.items()}

def set_seed(seed=42):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

class DetrDataset(Dataset):
    def __init__(self, split, processor, is_train, size=480, max_size=800, limit=None):
        ds = load_dataset("biglam/loc_beyond_words", split=split)
        if limit:
            ds = ds.select(range(limit))
        self.ds = ds
        self.processor = processor
        self.is_train = is_train
        self.size = size
        self.max_size = max_size

    def __len__(self):
        return len(self.ds)

    def __getitem__(self, idx):
        item = self.ds[idx]
        image = item["image"].convert("RGB")  # PIL image (grayscale pages -> RGB)
        objects = item["objects"]
        annotations = {
            "image_id": item["image_id"],
            "annotations": [
                {"category_id": o["category_id"], "bbox": o["bbox"], "area": o["area"], "iscrowd": o["iscrowd"]} for o in objects
            ],
        }
        if self.is_train:
            enc = self.processor(images=image, annotations=annotations, return_tensors="pt")
            labels = enc["labels"][0]
            return {
                "pixel_values": enc["pixel_values"][0],
                "pixel_mask": enc["pixel_mask"][0],
                "class_labels": labels["class_labels"],
                "boxes": labels["boxes"],
            }
        else:
            enc = self.processor(images=image, return_tensors="pt")  # image already RGB
            # absolute xyxy target boxes from annotations
            boxes = torch.as_tensor([[o["bbox"][0], o["bbox"][1],
                                      o["bbox"][0]+o["bbox"][2], o["bbox"][1]+o["bbox"][3]]
                                     for o in objects], dtype=torch.float32)
            labels = torch.as_tensor([o["category_id"] for o in objects], dtype=torch.long)
            w, h = item["width"], item["height"]
            return {
                "pixel_values": enc["pixel_values"][0],
                "pixel_mask": enc["pixel_mask"][0],
                "orig_size": torch.tensor([h, w]),
                "tgt_boxes": boxes,
                "tgt_labels": labels,
            }

def collate_fn(batch):
    max_h = max(b["pixel_values"].shape[1] for b in batch)
    max_w = max(b["pixel_values"].shape[2] for b in batch)
    C = batch[0]["pixel_values"].shape[0]
    pixel_values = torch.zeros(len(batch), C, max_h, max_w)
    pixel_mask = torch.zeros(len(batch), max_h, max_w)
    for i, b in enumerate(batch):
        img, m = b["pixel_values"], b["pixel_mask"]
        pixel_values[i,:,:img.shape[1],:img.shape[2]] = img
        pixel_mask[i,:m.shape[0],:m.shape[1]] = m
    labels = [{"class_labels": b["class_labels"], "boxes": b["boxes"]} for b in batch]
    return {"pixel_values": pixel_values, "pixel_mask": pixel_mask, "labels": labels}

def eval_collate(batch):
    if len(batch) == 1:
        b = batch[0]
        return {
            "pixel_values": b["pixel_values"].unsqueeze(0),
            "pixel_mask": b["pixel_mask"].unsqueeze(0),
            "orig_sizes": [b["orig_size"]],
            "tgt_boxes": [b["tgt_boxes"]],
            "tgt_labels": [b["tgt_labels"]],
        }
    raise ValueError("eval batch size must be 1")

class LRWarmupCosine:
    def __init__(self, optimizer, warmup, total):
        self.opt = optimizer; self.warmup = warmup; self.total = total
        self.t = 0; self.base = [g["lr"] for g in optimizer.param_groups]
    def step(self):
        self.t += 1
        s = self.t
        if s <= self.warmup:
            f = (s+1)/max(1, self.warmup)
        else:
            p = (s - self.warmup)/max(1, (self.total - self.warmup))
            f = 0.5*(1+math.cos(math.pi*p))
        for base, g in zip(self.base, self.opt.param_groups):
            g["lr"] = base*f

def evaluate(model, loader, processor, device, num_eval=None):
    model.eval()
    metric = MeanAveragePrecision(class_metrics=True, iou_type="bbox", backend="faster_coco_eval")
    count = 0
    with torch.no_grad():
        for batch in tqdm(loader, desc="eval"):
            pv = batch["pixel_values"].to(device)
            pm = batch["pixel_mask"].to(device)
            outs = model(pixel_values=pv, pixel_mask=pm)
            preds = processor.post_process_object_detection(outs, target_sizes=batch["orig_sizes"], threshold=0.0)
            for pred, tb, tl in zip(preds, batch["tgt_boxes"], batch["tgt_labels"]):
                pb = pred["boxes"].cpu().double()
                ps = pred["scores"].cpu()
                pl = pred["labels"].cpu()
                metric.update(
                    [{"boxes": pb, "scores": ps, "labels": pl}],
                    [{"boxes": tb.double(), "labels": tl}],
                )
            count += 1
            if num_eval and count >= num_eval:
                break
    res = metric.compute()
    flat = {}
    for k, v in res.items():
        if isinstance(v, torch.Tensor):
            flat[k] = v.item() if v.dim() == 0 else v.tolist()
        else:
            flat[k] = v
    return res, flat

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--epochs", type=int, default=24)
    ap.add_argument("--batch-size", type=int, default=2)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--lr-backbone", type=float, default=1e-5)
    ap.add_argument("--size", type=int, default=480)
    ap.add_argument("--max-size", type=int, default=800)
    ap.add_argument("--limit", type=int, default=None, help="limit training samples (smoke test)")
    ap.add_argument("--limit-val", type=int, default=None)
    ap.add_argument("--push", type=str, default="harness-race/control-r2")
    ap.add_argument("--no-push", action="store_true")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--output", type=str, default="/workspace/out")
    args = ap.parse_args()

    set_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}  torch={torch.__version__}  cuda={torch.cuda.is_available()}", flush=True)
    if torch.cuda.is_available():
        print("GPU:", torch.cuda.get_device_name(0), "mem", torch.cuda.get_device_properties(0).total_memory/1e9, flush=True)

    print("Loading image processor...", flush=True)
    processor = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50")

    print("Loading dataset...", flush=True)
    train_ds = DetrDataset("train", processor, is_train=True, size=args.size, max_size=args.max_size, limit=args.limit)
    val_ds = DetrDataset("validation", processor, is_train=False, size=args.size, max_size=args.max_size, limit=args.limit_val)
    print(f"train={len(train_ds)} val={len(val_ds)}", flush=True)

    train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, collate_fn=collate_fn, num_workers=2, prefetch_factor=2)
    val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, collate_fn=eval_collate, num_workers=1)

    print("Loading base model facebook/detr-resnet-50...", flush=True)
    model = DetrForObjectDetection.from_pretrained(
        "facebook/detr-resnet-50",
        ignore_mismatched_sizes=True,
        num_labels=len(ID2LABEL),
        id2label=ID2LABEL, label2id=LABEL2ID,
    )
    model.to(device)

    param_dicts = [
        {"params": [p for n,p in model.named_parameters() if "backbone" not in n and "reference_points" not in n and p.requires_grad], "lr": args.lr},
        {"params": [p for n,p in model.named_parameters() if "backbone" in n and p.requires_grad], "lr": args.lr_backbone},
        {"params": [p for n,p in model.named_parameters() if "reference_points" in n and p.requires_grad], "lr": args.lr*5},
    ]
    optimizer = torch.optim.AdamW(param_dicts, lr=args.lr, weight_decay=1e-4)

    steps_per_epoch = math.ceil(len(train_ds)/args.batch_size)
    total_steps = steps_per_epoch * args.epochs
    warmup = min(1000, int(0.02*total_steps))
    sched = LRWarmupCosine(optimizer, warmup, total_steps)
    print(f"steps_per_epoch={steps_per_epoch} total_steps={total_steps} warmup={warmup}", flush=True)

    def train_step(batch):
        pv = batch["pixel_values"].to(device)
        pm = batch["pixel_mask"].to(device)
        labels = [{k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k,v in lb.items()} for lb in batch["labels"]]
        out = model(pixel_values=pv, pixel_mask=pm, labels=labels)
        loss = out.loss
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)
        optimizer.step(); sched.step()
        return loss.item()

    # measure timing on first steps
    model.train()
    global_step = 0
    t0 = time.time()
    first_iter = True
    best_epoch = -1
    best_map50 = -1.0
    best_state = None
    for epoch in range(1, args.epochs+1):
        model.train()
        ep_loss = 0.0; nb = 0
        ep_start = time.time()
        for batch in train_loader:
            loss = train_step(batch)
            ep_loss += loss; nb += 1; global_step += 1
            if global_step % 100 == 0:
                print(f"[e{epoch}] step {nb}/{steps_per_epoch} loss={loss:.4f} lr_head={optimizer.param_groups[0]['lr']:.2e}", flush=True)
        if first_iter:
            elapsed = time.time()-t0
            print(f"FIRST {nb} steps in {elapsed:.1f}s -> est epoch time {elapsed/args.batch_size*args.batch_size:.0f}s", flush=True)
            first_iter = False
        print(f"--- EPOCH {epoch} done. mean loss={ep_loss/max(1,nb):.4f} time={(time.time()-ep_start)/60:.2f} min ---", flush=True)

        if epoch % max(1, args.epochs//6) == 0 or epoch == args.epochs:
            res, flat = evaluate(model, val_loader, processor, device, num_eval=args.limit_val)
            if flat.get('map_50',0.0) > best_map50:
                best_map50 = flat['map_50']; best_epoch = epoch
                best_state = {k: v.detach().cpu().clone() for k,v in model.state_dict().items()}
                print(f"  -> new best map50={best_map50:.4f} at epoch {epoch}", flush=True)
            print(f"EVAL epoch {epoch}: map50={flat.get('map_50'):.4f} map={flat.get('map'):.4f}", flush=True)

    # restore best checkpoint
    if best_state is not None:
        model.load_state_dict(best_state)
        print(f"Loaded best checkpoint from epoch {best_epoch} (map50={best_map50:.4f})", flush=True)

    # final full eval
    print("Final evaluation on validation split...", flush=True)
    res, flat = evaluate(model, val_loader, processor, device, num_eval=args.limit_val)
    print("FINAL_METRICS " + json.dumps(flat, default=str), flush=True)
    for k,v in flat.items():
        print(f"  {k}: {v}", flush=True)

    if args.no_push:
        print("No push requested.", flush=True)
        return

    # save and push
    os.makedirs(args.output, exist_ok=True)
    model.save_pretrained(args.output)
    processor.save_pretrained(args.output)

    # metrics map with label names for model card
    metric_names = {
        "Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ]": flat.get("map", 0.0),
        "Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ]": flat.get("map_50", 0.0),
        "Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ]": flat.get("map_75", 0.0),
        "Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ]": flat.get("mar_100", 0.0),
    }
    metric_json = json.dumps([{"type":"map","value":float(flat.get("map",0.0))},
                              {"type":"map_50","value":float(flat.get("map_50",0.0))},
                              {"type":"map_75","value":float(flat.get("map_75",0.0))},
                              {"type":"mar_100","value":float(flat.get("mar_100",0.0))}])

    readme = f"""---
license: apache-2.0
base_model: facebook/detr-resnet-50
tags:
- object-detection
- detr
- document-layout-analysis
- historical-newspapers
- pytorch
datasets:
- biglam/loc_beyond_words
pipeline_tag: object-detection
model-index:
- name: control-r2
  results:
  - task:
      type: object-detection
      name: Object Detection
    dataset:
      name: biglam/loc_beyond_words
      type: biglam/loc_beyond_words
      split: validation
    metrics:
    - type: map
      value: {float(flat['map']):.4f}
      name: mean average precision (COCO @[IoU=0.50:0.95])
    - type: map_50
      value: {float(flat['map_50']):.4f}
      name: mean average precision @IoU=0.50
    - type: map_75
      value: {float(flat['map_75']):.4f}
      name: mean average precision @IoU=0.75
    - type: mar_100
      value: {float(flat['mar_100']):.4f}
      name: mean average recall @[IoU=0.50:0.95], maxDets=100
---

# control-r2

Fine-tuned object detection model for historical newspaper layout analysis, trained on the
**Beyond Words** dataset (`biglam/loc_beyond_words`), which contains crowdsourced bounding-box
annotations of visual content on World War I-era newspaper pages from the Library of Congress
Chronicling America collection.

This model is fine-tuned from [`facebook/detr-resnet-50`](https://huggingface.co/facebook/detr-resnet-50),
released under the Apache-2.0 license. The full model weights and checkpoints are therefore freely
shareable and usable.

## Classes (7)

| id | label |
|----|-------|
| 0 | Photograph |
| 1 | Illustration |
| 2 | Map |
| 3 | Comics/Cartoon |
| 4 | Editorial Cartoon |
| 5 | Headline |
| 6 | Advertisement |

## Intended use

Detecting and localizing the seven types of visual (non-text) content in historical newspaper page
images, as a building block for document layout analysis and digitized-archive navigation.

## Training procedure

- **Base model:** `facebook/detr-resnet-50` (COCO-pretrained)
- **Dataset:** `biglam/loc_beyond_words` (train split, {len(train_ds)} images)
- **Image size:** shortest side {args.size}px, max side {args.max_size}px (aspect-ratio preserved)
- **Optimizer:** AdamW, head LR {args.lr}, backbone LR {args.lr_backbone}, weight decay 1e-4
- **Scheduler:** linear warmup then cosine decay
- **Batch size:** {args.batch_size} (per GPU)
- **Epochs:** {args.epochs}
- **Losses:** DETR Hungarian matching CE + L1 bbox + GIoU
- Gradient clipping at 0.1

## Evaluation (validation split)

COCO-style metrics ({len(val_ds)} validation images):

| Metric | Value |
|--------|-------|
| AP @[IoU=0.50:0.95] | {float(flat['map']):.4f} |
| AP @[IoU=0.50] | {float(flat['map_50']):.4f} |
| AP @[IoU=0.75] | {float(flat['map_75']):.4f} |
| AR @[IoU=0.50:0.95] maxDets=100 | {float(flat['mar_100']):.4f} |

## Full metrics dictionary

```json
{json.dumps({k: (float(v) if isinstance(v,(int,float)) else str(v)) for k,v in flat.items()}, indent=2)}
```

## Usage

```python
from transformers import AutoImageProcessor, DetrForObjectDetection
from PIL import Image

processor = AutoImageProcessor.from_pretrained("harness-race/control-r2")
model = DetrForObjectDetection.from_pretrained("harness-race/control-r2")

img = Image.open("page.jpg").convert("RGB")
inputs = processor(images=img, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)
results = processor.post_process_object_detection(
    outputs, target_sizes=torch.tensor([img.size[::-1]]), threshold=0.7
)
for score, label, box in zip(results[0]["scores"], results[0]["labels"], results[0]["boxes"]):
    print(model.config.id2label[label.item()], round(score.item(), 3), box.tolist())
```

## License

Apache-2.0 (inherited from `facebook/detr-resnet-50`).
"""
    print("===== MODEL CARD preview (truncated) =====", flush=True)
    print(readme[:600], flush=True)
    with open(os.path.join(args.output, "README.md"), "w") as f:
        f.write(readme)

    print(f"Pushing to {args.push}...", flush=True)
    api = HfApi()
    api.create_repo(repo_id=args.push, repo_type="model", exist_ok=True)
    api.upload_folder(folder_path=args.output, repo_id=args.push, repo_type="model")
    print("PUSHED " + args.push, flush=True)

if __name__ == "__main__":
    main()