Upload train_loc.py with huggingface_hub

train_loc.py

@@ -0,0 +1,289 @@
+#!/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()