cham-ocr-studio / ocr-training /scripts /ctc_lm_decoder.py
phucsd's picture
Upload folder using huggingface_hub
a183ff9 verified
Raw History Blame Contribute Delete
4.61 kB
import numpy as np
import re
def greedy_decode(logits, vocab, blank_idx):
"""
Giải mã Greedy CTC từ ma trận xác suất logits [T, C].
"""
preds = np.argmax(logits, axis=-1)
prev = -1
decoded = []
for p in preds:
if p != prev:
if p != blank_idx and p < len(vocab):
decoded.append(vocab[p])
prev = p
return "".join(decoded)
def log_add(x, y):
"""
Log-add exp helper: returns log(exp(x) + exp(y))
"""
if x == -float('inf'):
return y
if y == -float('inf'):
return x
return max(x, y) + np.log1p(np.exp(-abs(x - y)))
def prefix_beam_search_decode(logits, vocab, blank_idx, lm_model=None, beam_size=10, alpha=0.5, beta=0.1):
"""
Giải mã CTC Prefix Beam Search trong log-domain tránh tràn số dưới (underflow) trên chuỗi dài.
logits: ma trận xác suất [T, C] hoặc log-probs.
"""
T, C = logits.shape
# Đảm bảo logits là log-probabilities
if np.max(logits) <= 1.0 + 1e-5:
log_probs = np.log(np.clip(logits, 1e-20, 1.0))
else:
log_probs = logits
# Khởi tạo path: {text: (log_p_blank, log_p_non_blank)}
paths = {"": (0.0, -float('inf'))}
for t in range(T):
new_paths = {}
step_log_probs = log_probs[t]
for text, (log_pb, log_pnb) in paths.items():
for c_idx in range(C):
log_p_char = float(step_log_probs[c_idx])
# Xử lý blank token
if c_idx == blank_idx:
log_pb_new, log_pnb_new = new_paths.get(text, (-float('inf'), -float('inf')))
log_pb_new = log_add(log_pb_new, log_p_char + log_add(log_pb, log_pnb))
new_paths[text] = (log_pb_new, log_pnb_new)
else:
char = vocab[c_idx] if c_idx < len(vocab) else ""
# Xử lý lặp ký tự ở biên
if len(text) > 0 and text[-1] == char:
# 1. Kéo dài chuỗi (nếu trước đó có blank)
log_pb_new_ext, log_pnb_new_ext = new_paths.get(text + char, (-float('inf'), -float('inf')))
log_pnb_new_ext = log_add(log_pnb_new_ext, log_p_char + log_pb)
new_paths[text + char] = (log_pb_new_ext, log_pnb_new_ext)
# 2. Giữ nguyên chuỗi (nếu trước đó không có blank)
log_pb_new, log_pnb_new = new_paths.get(text, (-float('inf'), -float('inf')))
log_pnb_new = log_add(log_pnb_new, log_p_char + log_pnb)
new_paths[text] = (log_pb_new, log_pnb_new)
else:
# Ký tự mới: luôn kéo dài chuỗi
log_pb_new_ext, log_pnb_new_ext = new_paths.get(text + char, (-float('inf'), -float('inf')))
log_pnb_new_ext = log_add(log_pnb_new_ext, log_p_char + log_add(log_pb, log_pnb))
new_paths[text + char] = (log_pb_new_ext, log_pnb_new_ext)
# Thu hẹp chùm (Keep top beam_size * 2)
sorted_paths = sorted(new_paths.items(), key=lambda x: log_add(x[1][0], x[1][1]), reverse=True)
paths = dict(sorted_paths[:beam_size * 2])
# Tính điểm rescoring kết hợp LM
final_results = []
for text, (log_pb, log_pnb) in paths.items():
log_ctc_prob = log_add(log_pb, log_pnb)
if log_ctc_prob == -float('inf'):
continue
lm_score = 0.0
if lm_model is not None and len(text) > 0:
tokens = [('_' if c == ' ' else c) for c in text]
lm_query = " ".join(tokens)
try:
lm_score = lm_model.score(lm_query)
except Exception:
lm_score = -10.0
# Score = log_ctc_prob + alpha * lm_score + beta * len(text)
score = log_ctc_prob + alpha * lm_score + beta * len(text)
final_results.append((text, score))
final_results = sorted(final_results, key=lambda x: x[1], reverse=True)
return final_results[:beam_size]
def calculate_oracle_accuracy(predictions_top_n, ground_truth):
"""
Kiểm tra xem ground_truth có nằm trong danh sách Top-N hay không.
predictions_top_n: list của text giải mã.
"""
gt_clean = re.sub(r'\s+', ' ', ground_truth.strip())
for pred in predictions_top_n:
pred_clean = re.sub(r'\s+', ' ', pred.strip())
if pred_clean == gt_clean:
return 1.0
return 0.0