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