#!/usr/bin/env python3 """ Predict broken / lost signs from surrounding context using KN 5-gram LM. Strategy: For each position marked damage='broken' (lost) or damage='partial' with very low confidence, gather (left-4, right-4) context signs (non-broken). Score every vocab token by: p_fill(w) = p_LM(w | left context) × p_LM(right context | w, left context) Pick top-k; if top1 prob > threshold, use it (with ⸢⸣ half-brackets). Else keep "x" (lost). This uses only the LM — no image model. For stronger predictions you'd add the classifier's full prob distribution over the broken crop (if any crop exists), and blend with LM score. """ import argparse import json import math import pickle from collections import defaultdict from pathlib import Path def load_lm(path): from collections import Counter with open(path, 'rb') as f: d = pickle.load(f) counts = {k: Counter(v) for k, v in d['counts'].items()} ctypes = {k: defaultdict(set, {kk: set(vv) for kk, vv in v.items()}) for k, v in d['context_types'].items()} return { 'counts': counts, 'ctypes': ctypes, 'vocab': d['vocab'], 'order': d['order'], 'D': d['discount'], } class LM: def __init__(self, lm): self.__dict__.update(lm) self.ctx_total = {n: defaultdict(int) for n in range(1, self.order + 1)} for n, cnt in self.counts.items(): for ng, c in cnt.items(): if n > 1: self.ctx_total[n][ng[:-1]] += c self.cont_count = defaultdict(int) for ng in self.counts[2]: self.cont_count[ng[-1]] += 1 self.total_cont = sum(self.cont_count.values()) def cont_prob(self, w): return (self.cont_count.get(w, 0) + 1) / (self.total_cont + len(self.vocab) + 1) def prob(self, context, word): n = min(len(context) + 1, self.order) while n > 1: ctx = tuple(context[-(n - 1):]) if self.ctx_total[n].get(ctx, 0) > 0: break n -= 1 if n == 1: return self.cont_prob(word) ctx = tuple(context[-(n - 1):]) c_ng = self.counts[n].get(ctx + (word,), 0) c_ctx = self.ctx_total[n][ctx] types = len(self.ctypes[n - 1].get(ctx, ())) if (n - 1) in self.ctypes else 0 alpha = max(c_ng - self.D, 0) / c_ctx if c_ctx else 0 lam = (self.D * types) / c_ctx if c_ctx else 1.0 lower = self.prob(context[-(n - 2):] if n > 2 else (), word) return alpha + lam * lower def score_fill(self, left_ctx, right_ctx, word): """Joint prob of placing `word` between left_ctx and right_ctx.""" # p(word | left) × prod p(right_i | left + [word] + right[:i]) lp = math.log(max(self.prob(tuple(left_ctx), word), 1e-12)) seq = list(left_ctx) + [word] for rw in right_ctx: lp += math.log(max(self.prob(tuple(seq), rw), 1e-12)) seq.append(rw) return lp def topk_fill(self, left_ctx, right_ctx, k=5, candidates=None): """Return top-k (word, logprob) for filling a hole between contexts.""" if candidates is None: # Only consider words appearing at least twice (speed + quality) candidates = [w for w, c in self.vocab.items() if c >= 2] scores = [] for w in candidates: lp = self.score_fill(left_ctx, right_ctx, w) scores.append((w, lp)) scores.sort(key=lambda x: -x[1]) return scores[:k] def fill_tablet(tablet, lm, conf_thresh=math.log(0.15), ctx=4, topk=3): """Walk tablet lines; for each broken/x sign, predict fill from context. Modifies tablet in-place, adds 'predicted_fill' field to broken signs. """ # Flatten all signs with global index (respect line breaks as marker) flat = [] sign_refs = [] # parallel list of (line_idx, sign_idx) or None for EOL for li, L in enumerate(tablet['lines']): for si, s in enumerate(L['signs']): tok = s['sign'] if s.get('damage') != 'broken' else None flat.append(tok) sign_refs.append((li, si)) # Add end-of-line marker (None = boundary, not predicted) flat.append("") sign_refs.append(None) # For each hole, gather context filled = 0 for i, tok in enumerate(flat): if tok is not None and tok != "": continue if sign_refs[i] is None: continue # EOL, skip # Gather left-ctx (skip holes/EOLs) left = [] j = i - 1 while j >= 0 and len(left) < ctx: if flat[j] is not None and flat[j] != "": left.insert(0, flat[j]) j -= 1 right = [] j = i + 1 while j < len(flat) and len(right) < ctx: if flat[j] is not None and flat[j] != "": right.append(flat[j]) j += 1 topk_preds = lm.topk_fill(left, right, k=topk) li, si = sign_refs[i] sign = tablet['lines'][li]['signs'][si] sign['predicted_fill'] = [ {'word': w, 'logprob': round(lp, 3), 'prob': round(math.exp(lp / max(1, len(left) + len(right) + 1)), 4)} for w, lp in topk_preds ] if topk_preds and topk_preds[0][1] > conf_thresh * (len(left) + len(right) + 1): sign['fill_best'] = topk_preds[0][0] sign['sign_original'] = sign['sign'] sign['sign'] = topk_preds[0][0] sign['damage'] = 'predicted' # new category filled += 1 return filled def main(): ap = argparse.ArgumentParser() ap.add_argument('--lm', required=True, help='pickled KN LM') ap.add_argument('--structure', required=True, help='tablet_structure.jsonl') ap.add_argument('--output', required=True) ap.add_argument('--conf-thresh', type=float, default=0.15) ap.add_argument('--ctx', type=int, default=4) ap.add_argument('--topk', type=int, default=3) ap.add_argument('--limit', type=int, default=0) args = ap.parse_args() print(f"[load] LM {args.lm}") lm_data = load_lm(args.lm) lm = LM(lm_data) print(f" vocab={len(lm.vocab)}") # Pre-compute candidate list once candidates = [w for w, c in lm.vocab.items() if c >= 2] print(f" candidates={len(candidates)}") n = 0; total_filled = 0 with open(args.structure) as f, open(args.output, 'w') as out: for line in f: t = json.loads(line) filled = fill_tablet(t, lm, conf_thresh=math.log(args.conf_thresh), ctx=args.ctx, topk=args.topk) total_filled += filled out.write(json.dumps(t, ensure_ascii=False) + "\n") n += 1 if n % 10 == 0: print(f" [{n}] total_filled={total_filled}") if args.limit and n >= args.limit: break print(f"\nDONE: {n} tablets, {total_filled} broken signs predicted → {args.output}") if __name__ == '__main__': main()