Spaces:
Running
Running
Download ocr-training/scripts/ctc_lm_decoder.py from phucsd/cham-ocr-studio: direct link, hf CLI and curl.
- Browser
- Download file 4.61 kB
-
https://huggingface.co/spaces/phucsd/cham-ocr-studio/resolve/main/ocr-training/scripts/ctc_lm_decoder.py
- Command line
-
hf download hf://spaces/phucsd/cham-ocr-studio/ocr-training/scripts/ctc_lm_decoder.py
-
curl -L -o ctc_lm_decoder.py https://huggingface.co/spaces/phucsd/cham-ocr-studio/resolve/main/ocr-training/scripts/ctc_lm_decoder.py
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 | |