Upload train_fasterrcnn.py with huggingface_hub

train_fasterrcnn.py

@@ -0,0 +1,363 @@
+# /// script
+# dependencies = [
+#   "torch",
+#   "torchvision",
+#   "datasets",
+#   "pycocotools",
+#   "Pillow",
+#   "numpy",
+#   "huggingface_hub",
+# ]
+# ///
+"""Fine-tune torchvision Faster R-CNN (ResNet50-FPN, COCO-pretrained, BSD-3-Clause)
+on biglam/loc_beyond_words and push to harness-race/opencode-r1.
+"""
+import argparse
+import json
+import os
+import random
+import time
+
+import numpy as np
+import torch
+from PIL import Image
+
+import torchvision
+from torchvision.models.detection import fasterrcnn_resnet50_fpn, FasterRCNN_ResNet50_FPN_Weights
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+
+class_names = [
+    "Photograph", "Illustration", "Map", "Comics/Cartoon",
+    "Editorial Cartoon", "Headline", "Advertisement",
+]
+NUM_CLASSES = len(class_names) + 1  # + background for torchvision
+
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def process_row(row, max_dim):
+    img = row["image"]
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+    w, h = img.size
+    scale = min(1.0, max_dim / max(h, w))
+    nw, nh = max(1, round(w * scale)), max(1, round(h * scale))
+    if (nw, nh) != (w, h):
+        img = img.resize((nw, nh), Image.BILINEAR)
+    arr = np.asarray(img, dtype=np.uint8)  # H,W,C
+
+    boxes, labels, areas, ids = [], [], [], []
+    for obj in row["objects"]:
+        x, y, bw, bh = obj["bbox"]
+        x1, y1 = x * scale, y * scale
+        x2, y2 = (x + bw) * scale, (y + bh) * scale
+        if x2 <= x1 or y2 <= y1:
+            continue
+        boxes.append([x1, y1, x2, y2])
+        labels.append(int(obj["category_id"]))
+        areas.append((x2 - x1) * (y2 - y1))
+        ids.append(int(obj["id"]))
+    target = {
+        "boxes": torch.as_tensor(boxes, dtype=torch.float32) if boxes else torch.zeros((0, 4), dtype=torch.float32),
+        "labels": torch.as_tensor(labels, dtype=torch.int64) if labels else torch.zeros(0, dtype=torch.int64),
+        "image_id": torch.tensor([int(row["image_id"])]),
+        "area": torch.as_tensor(areas, dtype=torch.float32) if areas else torch.zeros(0, dtype=torch.float32),
+        "iscrowd": torch.zeros((len(boxes),), dtype=torch.int64),
+    }
+    return arr, target
+
+
+def collate(batch):
+    images, targets = [], []
+    for item in batch:
+        arr, target = item
+        t = torch.as_tensor(arr, dtype=torch.float32).permute(2, 0, 1) / 255.0
+        images.append(t)
+        targets.append(target)
+    return images, targets
+
+
+def train_one_epoch(model, optimizer, loader, device, epoch, log_every=25):
+    model.train()
+    tot, cnt = 0.0, 0
+    t0 = time.time()
+    for i, (images, targets) in enumerate(loader):
+        images = [im.to(device) for im in images]
+        targets = [{k: (v.to(device) if k != "image_id" else v) for k, v in t.items()} for t in targets]
+        loss_dict = model(images, targets)
+        loss = sum(v for v in loss_dict.values())
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        tot += loss.item()
+        cnt += 1
+        if i % log_every == 0:
+            names = {k: round(float(v.item()), 3) for k, v in loss_dict.items()}
+            print(f"[epoch {epoch}] step {i}/{len(loader)} loss={loss.item():.4f} {names} "
+                  f"elapsed={time.time()-t0:.0f}s", flush=True)
+    return tot / max(cnt, 1)
+
+
+@torch.no_grad()
+def evaluate(model, loader, device, images_per_run=0):
+    model.eval()
+    preds = []
+    for images, targets in loader:
+        images = [im.to(device) for im in images]
+        out = model(images)
+        for img_id, t, dets in zip([int(t["image_id"][0]) for t in targets], targets, out):
+            boxes = dets["boxes"].cpu().numpy()
+            scores = dets["scores"].cpu().numpy()
+            labels = dets["labels"].cpu().numpy()
+            for box, sc, lab in zip(boxes, scores, labels):
+                if sc < 0.5:
+                    continue
+                x1, y1, x2, y2 = box
+                preds.append({
+                    "image_id": img_id,
+                    "category_id": int(lab),
+                    "bbox": [float(x1), float(y1), float(x2 - x1), float(y2 - y1)],
+                    "score": float(sc),
+                })
+    return preds
+
+
+def build_gt(items):
+    anns, img_infos = [], {}
+    for arr, target in items:
+        iid = int(target["image_id"][0])
+        img_infos[iid] = {"id": iid, "width": arr.shape[1], "height": arr.shape[0]}
+        for bx, lab, ar in zip(target["boxes"], target["labels"], target["area"]):
+            x1, y1, x2, y2 = bx.tolist()
+            anns.append({
+                "id": len(anns) + 1, "image_id": iid, "category_id": int(lab),
+                "bbox": [x1, y1, max(x2 - x1, 1), max(y2 - y1, 1)],
+                "area": float(ar), "iscrowd": 0,
+            })
+    gt = {"images": list(img_infos.values()), "annotations": anns,
+          "categories": [{"id": i, "name": n} for i, n in enumerate(class_names)]}
+    return gt
+
+
+def coco_eval(gt, preds):
+    from pycocotools.coco import COCO
+    from pycocotools.cocoeval import COCOeval
+    coco_gt = COCO()
+    coco_gt.dataset = gt
+    coco_gt.createIndex()
+    coco_dt = coco_gt.loadRes(preds)
+    ev = COCOeval(coco_gt, coco_dt, "bbox")
+    ev.evaluate()
+    ev.accumulate()
+    ev.summarize()
+    stats = ev.stats  # [mAP .5:.95, mAP50, mAP75, mAP small, med, large, AR...]
+    out = {
+        "mAP_050_095": float(stats[0]),
+        "mAP_050": float(stats[1]),
+        "mAP_075": float(stats[2]),
+    }
+    print(json.dumps(out), flush=True)
+    return out
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--epochs", type=int, default=12)
+    ap.add_argument("--batch", type=int, default=8)
+    ap.add_argument("--max-dim", type=int, default=520)
+    ap.add_argument("--lr", type=float, default=2e-3)
+    ap.add_argument("--base-lr", type=float, default=2e-4)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--push", default="1")
+    args = ap.parse_args()
+
+    set_seed(args.seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("device:", device, torch.cuda.get_device_name(0) if torch.cuda.is_available() else "", flush=True)
+
+    from datasets import load_dataset
+    from torch.utils.data import DataLoader
+
+    print("loading dataset ...", flush=True)
+    ds = load_dataset("biglam/loc_beyond_words")
+    train_src = ds["train"]
+    val_src = ds["validation"]
+    print("train rows:", len(train_src), "val rows:", len(val_src), flush=True)
+
+    print("preprocessing ...", flush=True)
+    t0 = time.time()
+    train_items = [process_row(r, args.max_dim) for r in train_src]
+    val_items = [process_row(r, args.max_dim) for r in val_src]
+    print(f"preprocess done in {time.time()-t0:.0f}s", flush=True)
+    gt = build_gt(val_items)
+
+    from torch.utils.data import Dataset
+
+    class Wrap(Dataset):
+        def __init__(self, items):
+            self.items = items
+
+        def __len__(self):
+            return len(self.items)
+
+        def __getitem__(self, i):
+            return self.items[i]
+
+    train_loader = DataLoader(Wrap(train_items), batch_size=args.batch, shuffle=True,
+                              num_workers=4, collate_fn=collate, drop_last=False)
+    val_loader = DataLoader(Wrap(val_items), batch_size=4, shuffle=False,
+                            num_workers=4, collate_fn=collate)
+
+    model = fasterrcnn_resnet50_fpn(weights=FasterRCNN_ResNet50_FPN_Weights.COCO_V1)
+    in_features = model.roi_heads.box_predictor.cls_score.in_features
+    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, NUM_CLASSES)
+    model.transform.min_size = (args.max_dim,)
+    model.transform.max_size = int(args.max_dim * 1.5)
+    model.to(device)
+
+    params = [
+        {"params": [p for n, p in model.backbone.named_parameters() if p.requires_grad],
+         "lr": args.base_lr},
+        {"params": [p for n, p in model.named_parameters()
+                    if not n.startswith("backbone") and p.requires_grad],
+         "lr": args.lr},
+    ]
+    optimizer = torch.optim.SGD(params, momentum=0.9, weight_decay=1e-4)
+    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=8, gamma=0.3)
+
+    best = -1.0
+    for epoch in range(1, args.epochs + 1):
+        avg = train_one_epoch(model, optimizer, train_loader, device, epoch)
+        lr_scheduler.step()
+        print(f"=== epoch {epoch} done, avg_train_loss={avg:.4f}, lr={optimizer.param_groups[-1]['lr']:.2e} ===", flush=True)
+        if epoch % 4 == 0 or epoch == args.epochs:
+            preds = evaluate(model, val_loader, device)
+            results = coco_eval(gt, preds)
+            if epoch == args.epochs or results["mAP_050_095"] > best:
+                best = results["mAP_050_095"]
+                os.makedirs("output", exist_ok=True)
+                torch.save(model.state_dict(), "output/model.pth")
+                torch.save(model, "output/full_model.pth")
+                with open("output/results.json", "w") as f:
+                    json.dump(results, f)
+                with open("output/args.json", "w") as f:
+                    json.dump(vars(args), f)
+                with open("output/class_names.json", "w") as f:
+                    json.dump(class_names, f)
+                print(f"[saved checkpoint] mAP={results['mAP_050_095']:.4f}", flush=True)
+
+    print("training complete.", flush=True)
+    if args.push == "1" and os.path.exists("output/results.json"):
+        push_model(args, "output")
+    else:
+        print("skipping push.", flush=True)
+
+
+def push_model(args, out_dir):
+    from huggingface_hub import HfApi, upload_folder
+    import shutil
+
+    results = json.load(open(os.path.join(out_dir, "results.json")))
+    params = json.load(open(os.path.join(out_dir, "args.json")))
+
+    repo_id = "harness-race/opencode-r1"
+    token = os.environ.get("HF_TOKEN")
+    api = HfApi(token=token)
+    api.create_repo(repo_id, repo_type="model", exist_ok=True)
+    # put README / metadata inside out_dir (upload_folder pushes everything there)
+    readme = f"""---
+license: bsd-3-clause
+language:
+  - en
+tags:
+  - object-detection
+  - faster-rcnn
+  - resnet50
+  - document-layout
+  - historical-newspapers
+pipeline_tag: object-detection
+metrics:
+  - {float(results['mAP_050_095']):.4f}
+widget:
+  - src: https://datasets-server.huggingface.co/cached-assets/biglam/loc_beyond_words/--/6c7f5fb3c60f02d9fe925cfc14aa7008f6c89099/--/default/train/0/image/image.jpg
+---
+
+# opencode-r1
+
+Object detection model fine-tuned from **torchvision Faster R-CNN (ResNet-50-FPN)**
+pre-trained on COCO (base model license: BSD-3-Clause, open and shareable) on the
+[`biglam/loc_beyond_words`](https://huggingface.co/datasets/biglam/loc_beyond_words) dataset
+(Library of Congress "Beyond Words", data license CC0-1.0).
+
+## Classes (7) + background
+
+{", ".join(class_names)}
+
+## Validation results (COCO-style, biglam/loc_beyond_words validation set)
+
+| Metric | Value |
+|---|---|
+| mAP @[IoU=0.50:0.95] | {results['mAP_050_095']:.4f} |
+| mAP @ IoU=0.50 | {results['mAP_050']:.4f} |
+| mAP @ IoU=0.75 | {results['mAP_075']:.4f} |
+
+## Training
+
+| Setting | Value |
+|---|---|
+| Base model | Faster R-CNN ResNet50-FPN (COCO, BSD-3-Clause) |
+| Epochs | {params['epochs']} |
+| Batch size | {params['batch']} |
+| Max image dim | {params['max_dim']} |
+| Optimizer | SGD (momentum 0.9), StepLR x0.3/8 epochs |
+| Head LR / Backbone LR | {params['lr']} / {params['base_lr']} |
+| Hardware | NVIDIA GPU (Hugging Face jobs) |
+
+Images are downscaled so the largest dimension is {params['max_dim']}px (aspect preserved);
+boxes scaled accordingly. Predictions below score 0.5 are discarded.
+
+## To load and run
+
+```python
+import torch
+from torchvision.models.detection import fasterrcnn_resnet50_fpn
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+from huggingface_hub import hf_hub_download
+from PIL import Image
+import numpy as np
+
+state = torch.load(hf_hub_download("harness-race/opencode-r1", "model.pth"),
+                   map_location="cpu")
+model = fasterrcnn_resnet50_fpn(weights=None)
+in_features = model.roi_heads.box_predictor.cls_score.in_features
+model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 1 + 7)  # 7 + background
+model.load_state_dict(state)
+model.eval()
+
+img = Image.open("page.jpg").convert("RGB")
+# resize to max-dim {params['max_dim']} like training, then:
+x = torch.as_tensor(np.asarray(img), dtype=torch.float32).permute(2, 0, 1) / 255.0
+with torch.no_grad():
+    dets = model([x])[0]
+```
+"""
+    with open(os.path.join(out_dir, "README.md"), "w") as f:
+        f.write(readme)
+
+    upload_folder(
+        repo_id=repo_id,
+        folder_path=out_dir,
+        repo_type="model",
+        token=token,
+        commit_message="Fine-tuned Faster R-CNN ResNet50-FPN on loc_beyond_words",
+    )
+    print(f"pushed model to {repo_id}", flush=True)
+
+
+if __name__ == "__main__":
+    main()