#!/usr/bin/env python3 """ End-to-end tablet-to-Latin transliteration using the best validated models. Model stack (from RESULTS_SUMMARY.md / PAPER_METHODS.md): * Detection: YOLO11m 5-fold ensemble (runs/detect/.../yolo_fold{0..4}/weights/best.pt) * Classification: 4-model weighted ensemble (top1=0.9008 on val fold): dinov3_vitl14 × hitit_dinov3_cls_v12_ultimate/best_ema.pt dinov3_vitl14 × hitit_dinov3l_v13a/best_ema.pt convnextv2_large × hitit_convnextl_v13b/best_ema.pt dinov3_vitb14 × hitit_dinov3b_v13c/best_ema.pt plus (optional) confusion-pair MLP head (→0.9137) + KN 5-gram rescore (→0.9014). * Language: ground-truth from mark.txt when available; TODO: trained head. * Damage: ground-truth from comment heuristic; TODO: trained head. * Fill-in: KN 5-gram LM over TLHdig tokens. """ import argparse import json import sys import time from pathlib import Path import numpy as np import torch import torch.nn.functional as F from PIL import Image from torchvision import transforms 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 from enhancements.translit_mapper import transliterate_tablet # Default ckpt stack (best-performing per RESULTS_SUMMARY.md) DEFAULT_CKPTS = [ ('dinov3_vitl14', 'hitit_ocr/runs/h100/hitit_dinov3_cls_v12_ultimate/best_ema.pt', 0.175), ('dinov3_vitl14', 'hitit_ocr/runs/h100/hitit_dinov3l_v13a/best_ema.pt', 0.25), ('convnextv2_large','hitit_ocr/runs/h100/hitit_convnextl_v13b/best_ema.pt', 0.25), ('dinov3_vitb14', 'hitit_ocr/runs/h100/hitit_dinov3b_v13c/best_ema.pt', 0.275), ] DEFAULT_YOLO_FOLDS = [ f"runs/detect/hitit_ocr/runs/h100/yolo_fold{i}/weights/best.pt" for i in range(5) ] # ─────────── Detection ─────────── def detect_signs(image_path, yolo_weights, conf=0.25, iou=0.5, imgsz=1280, device='cuda'): from ultralytics import YOLO img = Image.open(image_path).convert('RGB') W, H = img.size all_dets = [] for w in yolo_weights: model = YOLO(str(w)) r = model.predict(str(image_path), conf=conf, iou=iou, imgsz=imgsz, device=device, verbose=False) for box in r[0].boxes: x1, y1, x2, y2 = box.xyxy[0].cpu().tolist() all_dets.append({'xyxy': [x1, y1, x2, y2], 'conf': float(box.conf[0])}) return img, W, H, _nms(all_dets, iou_thr=0.55) def _nms(dets, iou_thr=0.5): if not dets: return [] dets = sorted(dets, key=lambda d: -d['conf']) keep = [] while dets: m = dets.pop(0); keep.append(m) dets = [d for d in dets if _iou(m['xyxy'], d['xyxy']) < iou_thr] return keep def _iou(a, b): x1, y1, x2, y2 = a; X1, Y1, X2, Y2 = b iw = max(0, min(x2, X2) - max(x1, X1)) ih = max(0, min(y2, Y2) - max(y1, Y1)) inter = iw * ih ua = (x2 - x1) * (y2 - y1) + (X2 - X1) * (Y2 - Y1) - inter return inter / max(ua, 1e-6) # ─────────── Line clustering ─────────── def cluster_lines(dets, line_tol=0.015, H=1.0): """Group dets into reading lines by y-centroid; sort L→R inside each line.""" if not dets: return [] for d in dets: x1, y1, x2, y2 = d['xyxy'] d['cx'] = (x1 + x2) / 2; d['cy'] = (y1 + y2) / 2; d['h'] = y2 - y1 dets = sorted(dets, key=lambda d: d['cy']) lines = []; cur = [dets[0]] for d in dets[1:]: ref = cur[-1] if abs(d['cy'] - ref['cy']) < max(line_tol * H, 0.6 * ref['h']): cur.append(d) else: lines.append(cur); cur = [d] lines.append(cur) out = [] for row, L in enumerate(lines, start=1): L = sorted(L, key=lambda d: d['cx']) out.append({'row': row, 'col': 1, 'signs': L}) return out # ─────────── Classification ensemble ─────────── class EnsembleClassifier: def __init__(self, ckpt_specs=DEFAULT_CKPTS, device='cuda', dtype=torch.bfloat16): self.device = device self.dtype = dtype self.models = [] self.weights = [] self.label_to_idx = None self.idx_to_label = None for arch, ckpt_rel, w in ckpt_specs: ckpt = ROOT / ckpt_rel if not ckpt.exists(): print(f" SKIP {arch}: ckpt missing {ckpt}") continue print(f" Load {arch} ← {ckpt_rel} (w={w})") d = torch.load(ckpt, map_location='cpu', weights_only=False) l2i = d.get('label_to_idx') or d.get('meta', {}).get('label_to_idx') if l2i is not None and self.label_to_idx is None: self.label_to_idx = l2i n_classes = len(l2i) if l2i else 198 img_size = get_arch_img_size(arch) model = build_backbone(arch, n_classes=n_classes, img_size_override=img_size) sd = d.get('model_state_dict') or d.get('state_dict') or d model.load_state_dict(sd, strict=False) model.to(device).eval() self.models.append((arch, model, img_size)) self.weights.append(w) assert self.models, "No classifier ckpts loaded" self.idx_to_label = {v: k for k, v in self.label_to_idx.items()} self.weights = torch.tensor(self.weights).float() self.weights = self.weights / self.weights.sum() print(f" Ensemble: {len(self.models)} models, {len(self.label_to_idx)} classes") mean = (0.489, 0.448, 0.424); std = (0.362, 0.359, 0.364) self._transforms_cache = {} for arch, _, img_size in self.models: if img_size not in self._transforms_cache: self._transforms_cache[img_size] = transforms.Compose([ transforms.Resize((img_size, img_size), antialias=True), transforms.ToTensor(), transforms.Normalize(mean, std), ]) @torch.no_grad() def classify_crops(self, crops, batch_size=32): """Input: list of PIL Images. Output: list of {label, conf, top5}.""" if not crops: return [] # Group by img_size preds_per_model = [] for wi, (arch, model, img_size) in enumerate(self.models): tf = self._transforms_cache[img_size] xs = torch.stack([tf(c) for c in crops]).to(self.device) probs = [] for i in range(0, len(xs), batch_size): x = xs[i:i + batch_size] with torch.amp.autocast('cuda', dtype=self.dtype, enabled=True): logits = model(x) probs.append(F.softmax(logits.float(), dim=-1).cpu()) preds_per_model.append(torch.cat(probs)) # Weighted softmax average stacked = torch.stack(preds_per_model, dim=0) # (M, N, C) w = self.weights.view(-1, 1, 1) ens = (stacked * w).sum(0) top5_vals, top5_idxs = ens.topk(5, dim=-1) out = [] for i in range(len(crops)): lbl = self.idx_to_label[int(top5_idxs[i, 0])] conf = float(top5_vals[i, 0]) top5 = [(self.idx_to_label[int(top5_idxs[i, k])], float(top5_vals[i, k])) for k in range(5)] out.append({'label': lbl, 'conf': conf, 'top5': top5}) return out # ─────────── Damage heuristic (no head yet) ─────────── def infer_damage(pred, crop): """Simple heuristic until dedicated damage head is trained: - very low confidence → broken - moderate → partial """ c = pred['conf'] if c < 0.30: return 'broken' if c < 0.55: return 'partial' if c < 0.75: return 'uncertain' return 'intact' # ─────────── Language inference ─────────── # Use simple rule: sign label case signals language def infer_language(label): if not label or label in ('x', 'X'): return 'unk' if label.isupper(): return 'sum' # or akk (needs context; default sum logogram) return 'hit' # ─────────── End-to-end ─────────── def tablet_inference(image_path, yolo_weights, cls_model, output_json, output_text=None, conf_thresh=0.25, debug=False): t0 = time.time() img, W, H, dets = detect_signs(image_path, yolo_weights, conf=conf_thresh) if debug: print(f"[det] {len(dets)} signs in {W}×{H} ({time.time()-t0:.1f}s)") lines = cluster_lines(dets, H=H) if debug: print(f"[seg] {len(lines)} reading lines") # Prepare crops crops = [] crop_refs = [] for L in lines: for d in L['signs']: x1, y1, x2, y2 = [int(v) for v in d['xyxy']] crop = img.crop((max(0, x1 - 2), max(0, y1 - 2), min(W, x2 + 2), min(H, y2 + 2))) crops.append(crop) crop_refs.append(d) preds = cls_model.classify_crops(crops) for d, p in zip(crop_refs, preds): d['sign'] = p['label'] d['sign_conf'] = p['conf'] d['top5'] = p['top5'] d['damage'] = infer_damage(p, None) d['lang'] = infer_language(p['label']) if debug: print(f"[cls] {time.time()-t0:.1f}s total") # Structured tablet tablet_struct = { 'tablet_id': Path(image_path).stem, 'image_path': str(image_path), 'width': W, 'height': H, 'n_signs': sum(len(L['signs']) for L in lines), 'n_lines': len(lines), 'lang_dist': {}, 'damage_dist': {}, 'lines': [], } for L in lines: signs_out = [] for d in L['signs']: x1, y1, x2, y2 = d['xyxy'] signs_out.append({ 'sign': d['sign'], 'lang': d['lang'], 'damage': d['damage'], 'bbox': [x1 / W, y1 / H, (x2 - x1) / W, (y2 - y1) / H], 'conf': d['sign_conf'], 'row': L['row'], 'col': L['col'], 'top5': [{'sign': s, 'prob': p} for s, p in d['top5']], }) tablet_struct['lang_dist'][d['lang']] = tablet_struct['lang_dist'].get(d['lang'], 0) + 1 tablet_struct['damage_dist'][d['damage']] = tablet_struct['damage_dist'].get(d['damage'], 0) + 1 tablet_struct['lines'].append({'row': L['row'], 'col': L['col'], 'signs': signs_out}) result = transliterate_tablet(tablet_struct) # Optional: fill broken predictions via LM lm_pkl = ROOT / 'hitit_ocr/runs/h100/sign_5gram_lm.pkl' if lm_pkl.exists(): try: from enhancements.broken_predict import load_lm, LM, fill_tablet lm = LM(load_lm(str(lm_pkl))) filled = fill_tablet(tablet_struct, lm, ctx=4, topk=5) if debug: print(f"[fill] {filled} broken signs predicted") # Re-render with fill result = transliterate_tablet(tablet_struct) except Exception as e: if debug: print(f"[fill] failed: {e}") # Enrich result with structured data result['structured'] = tablet_struct Path(output_json).parent.mkdir(parents=True, exist_ok=True) Path(output_json).write_text(json.dumps(result, ensure_ascii=False, indent=2)) if output_text: Path(output_text).parent.mkdir(parents=True, exist_ok=True) Path(output_text).write_text(result['text']) return result def main(): ap = argparse.ArgumentParser() ap.add_argument('--image', help='single tablet') ap.add_argument('--image-dir', help='directory of tablet images') ap.add_argument('--output', help='output JSON (single mode)') ap.add_argument('--output-dir', help='output dir (batch)') ap.add_argument('--text', help='output TXT (single mode)') ap.add_argument('--yolo-weights', nargs='*', default=None) ap.add_argument('--conf', type=float, default=0.25) ap.add_argument('--debug', action='store_true') args = ap.parse_args() # YOLO weights yolo_weights = args.yolo_weights or [str(ROOT / p) for p in DEFAULT_YOLO_FOLDS] yolo_weights = [w for w in yolo_weights if Path(w).exists()] print(f"YOLO folds: {len(yolo_weights)}") # Classifier ensemble print("Loading classifier ensemble...") cls_model = EnsembleClassifier() if args.image: r = tablet_inference(args.image, yolo_weights, cls_model, args.output or (Path(args.image).stem + '.json'), args.text, conf_thresh=args.conf, debug=args.debug) print(f"\nResult: {r['stats']['n_lines']} lines, {r['stats']['n_signs']} signs") print(f"Lang: {r['stats']['lang_dist']}") print(f"Damage: {r['stats']['damage_dist']}") elif args.image_dir: out_dir = Path(args.output_dir); out_dir.mkdir(parents=True, exist_ok=True) imgs = (sorted(Path(args.image_dir).glob('*.jpg')) + sorted(Path(args.image_dir).glob('*.JPG')) + sorted(Path(args.image_dir).glob('*.png'))) for img_p in imgs: stem = img_p.stem try: tablet_inference(img_p, yolo_weights, cls_model, out_dir / f"{stem}.json", out_dir / f"{stem}.txt", conf_thresh=args.conf, debug=args.debug) print(f" OK {stem}") except Exception as e: print(f" FAIL {stem}: {e}") else: ap.error("Need --image or --image-dir") if __name__ == '__main__': main()