| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| 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") |
| 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") |
| |
| 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() |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
| |
| os.makedirs(args.output, exist_ok=True) |
| model.save_pretrained(args.output) |
| processor.save_pretrained(args.output) |
|
|
| |
| 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() |
|
|
