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