| |
| """ |
| 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_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) |
| ] |
|
|
|
|
| |
| 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) |
|
|
|
|
| |
| 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 |
|
|
|
|
| |
| 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 [] |
| |
| 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)) |
| |
| stacked = torch.stack(preds_per_model, dim=0) |
| 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 |
|
|
|
|
| |
| 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' |
|
|
|
|
| |
| |
| def infer_language(label): |
| if not label or label in ('x', 'X'): |
| return 'unk' |
| if label.isupper(): |
| return 'sum' |
| return 'hit' |
|
|
|
|
| |
| 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") |
|
|
| |
| 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") |
|
|
| |
| 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) |
|
|
| |
| 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") |
| |
| result = transliterate_tablet(tablet_struct) |
| except Exception as e: |
| if debug: print(f"[fill] failed: {e}") |
|
|
| |
| 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 = 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)}") |
|
|
| |
| 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() |
|
|
