code/src/tablet_to_latin.py

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