| |
| """Multi-task: classification + bbox-normalized position (reading order hint). |
| |
| Auxiliary head predicts (x_center, y_center) normalized to tablet bbox. |
| Forces backbone to encode spatial context → better discrimination for |
| signs that only differ in tablet position (numerals, logograms). |
| """ |
| import os, sys, json, argparse, time |
| from pathlib import Path |
| from collections import Counter |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler |
| from PIL import Image |
|
|
| ROOT = Path("/arf/scratch/stakan/hitit-proje") |
| sys.path.insert(0, str(ROOT / "hitit_ocr/src")) |
| from train_classification import build_backbone, get_arch_img_size |
|
|
| def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True) |
|
|
|
|
| class PosDS(Dataset): |
| def __init__(self, manifest, label_to_idx, img_size, is_train, val_fold, min_samples): |
| cls_count = Counter() |
| self.records = [] |
| for line in open(manifest): |
| r = json.loads(line) |
| if r.get('task') != 'classification' or not r.get('unified_label'): continue |
| cls_count[r['unified_label']] += 1 |
| for line in open(manifest): |
| r = json.loads(line) |
| if r.get('task') != 'classification': continue |
| if not r.get('unified_label') or r['unified_label'] not in label_to_idx: continue |
| if not r.get('path') or r.get('storage') != 'fs': continue |
| if r.get('integrity_ok') is False: continue |
| if cls_count[r['unified_label']] < min_samples: continue |
| fold = r.get('tablet_view_fold', 0) |
| if is_train and fold == val_fold: continue |
| if not is_train and fold != val_fold: continue |
| self.records.append(r) |
| self.l2i = label_to_idx |
| from torchvision import transforms |
| self.tf = transforms.Compose([ |
| transforms.Resize((img_size, img_size), antialias=True), |
| transforms.ToTensor(), |
| transforms.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]), |
| ]) |
| def __len__(self): return len(self.records) |
| def __getitem__(self, i): |
| r = self.records[i] |
| img = Image.open(r['path']).convert('RGB') |
| |
| bbox = r.get('bbox_normalized') or r.get('bbox') or [0.25, 0.25, 0.75, 0.75] |
| if isinstance(bbox, list) and len(bbox) == 4: |
| xc = 0.5 * (bbox[0] + bbox[2]); yc = 0.5 * (bbox[1] + bbox[3]) |
| else: |
| xc, yc = 0.5, 0.5 |
| pos = torch.tensor([xc, yc], dtype=torch.float32) |
| return self.tf(img), self.l2i[r['unified_label']], pos |
|
|
|
|
| class MultiTaskHead(nn.Module): |
| def __init__(self, feat_dim, n_classes): |
| super().__init__() |
| self.cls = nn.Linear(feat_dim, n_classes) |
| self.pos = nn.Sequential( |
| nn.Linear(feat_dim, 128), nn.GELU(), nn.Linear(128, 2), nn.Sigmoid() |
| ) |
| def forward(self, feats): |
| return self.cls(feats), self.pos(feats) |
|
|
|
|
| @torch.no_grad() |
| def extract(model, x, dtype): |
| raw = model.module if hasattr(model, 'module') else model |
| with torch.amp.autocast('cuda', dtype=dtype, enabled=True): |
| if hasattr(raw, 'forward_features'): |
| f = raw.forward_features(x) |
| if f.dim() == 3: f = f[:, 0] |
| elif f.dim() == 4: f = f.mean(dim=(2, 3)) |
| else: |
| f = raw(x) |
| return f.float() |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--ckpt', required=True) |
| ap.add_argument('--manifest', required=True) |
| ap.add_argument('--val-fold', type=int, default=0) |
| ap.add_argument('--min-samples', type=int, default=10) |
| ap.add_argument('--epochs', type=int, default=30) |
| ap.add_argument('--lr', type=float, default=5e-3) |
| ap.add_argument('--pos-weight', type=float, default=0.3) |
| ap.add_argument('--batch-size', type=int, default=128) |
| ap.add_argument('--output', required=True) |
| args = ap.parse_args() |
|
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float32 |
|
|
| arch, path = args.ckpt.split(':', 1) |
| ck = torch.load(path, map_location='cpu', weights_only=False) |
| label_to_idx = ck['label_to_idx'] |
| n_cls = len(label_to_idx) |
| img_size = get_arch_img_size(arch) |
| bb = build_backbone(arch, n_classes=n_cls, img_size_override=img_size).to(device) |
| sd = ck['model'] |
| sd = {k.replace('module.', '', 1): v for k, v in sd.items()} |
| sd = {k.replace('_orig_mod.', '', 1): v for k, v in sd.items()} |
| sd = {k: v for k, v in sd.items() if k != 'n_averaged'} |
| bb.load_state_dict(sd, strict=False) |
| for p in bb.parameters(): p.requires_grad = False |
| bb.eval() |
|
|
| tr_ds = PosDS(args.manifest, label_to_idx, img_size, True, args.val_fold, args.min_samples) |
| va_ds = PosDS(args.manifest, label_to_idx, img_size, False, args.val_fold, args.min_samples) |
| log(f"Train={len(tr_ds)} Val={len(va_ds)}") |
|
|
| cc = Counter(tr_ds.l2i[r['unified_label']] for r in tr_ds.records) |
| sample_w = np.array([1.0 / max(1, cc[tr_ds.l2i[r['unified_label']]])**0.5 |
| for r in tr_ds.records], dtype=np.float32) |
| tr_dl = DataLoader(tr_ds, batch_size=args.batch_size, |
| sampler=WeightedRandomSampler(sample_w, len(tr_ds), True), |
| num_workers=6, pin_memory=True, drop_last=True) |
| va_dl = DataLoader(va_ds, batch_size=args.batch_size, shuffle=False, |
| num_workers=4, pin_memory=True) |
|
|
| with torch.no_grad(): |
| feat_dim = extract(bb, torch.zeros(1, 3, img_size, img_size, device=device), dtype).size(-1) |
| head = MultiTaskHead(feat_dim, n_cls).to(device) |
| opt = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=1e-4) |
| sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs) |
|
|
| best, best_probs, best_y = 0, None, None |
| for ep in range(args.epochs): |
| head.train() |
| tl, nb = 0, 0 |
| for x, y, pos in tr_dl: |
| x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True) |
| pos = pos.to(device, non_blocking=True) |
| f = extract(bb, x, dtype) |
| logits_cls, pred_pos = head(f) |
| loss_cls = F.cross_entropy(logits_cls, y, label_smoothing=0.1) |
| loss_pos = F.mse_loss(pred_pos, pos) |
| loss = loss_cls + args.pos_weight * loss_pos |
| opt.zero_grad(); loss.backward(); opt.step() |
| tl += loss.item(); nb += 1 |
| sched.step() |
| head.eval() |
| cs, tot, p_all, y_all = 0, 0, [], [] |
| with torch.no_grad(): |
| for x, y, _ in va_dl: |
| x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True) |
| f = extract(bb, x, dtype) |
| logits, _ = head(f) |
| p = F.softmax(logits.float(), dim=-1) |
| p_all.append(p.cpu()); y_all.append(y.cpu()) |
| cs += (logits.argmax(-1) == y).sum().item(); tot += y.size(0) |
| acc = cs / max(1, tot) |
| if ep % 5 == 0 or ep == args.epochs - 1: |
| log(f"ep {ep}: loss={tl/max(1,nb):.4f} val_top1={acc:.4f}") |
| if acc > best: |
| best = acc |
| best_probs = torch.cat(p_all); best_y = torch.cat(y_all) |
| torch.save({'head_state': head.state_dict(), 'label_to_idx': label_to_idx, |
| 'probs': best_probs, 'targets': best_y, 'top1': best, |
| 'backbone_arch': arch}, |
| args.output) |
| log(f"=== MultiTask BEST: {best:.4f}") |
|
|
| if __name__ == '__main__': |
| main() |
|
|
