from itertools import groupby import os from pathlib import Path import asyncio import warnings import albumentations as A from albumentations.pytorch import ToTensorV2 import cv2 import gradio as gr import kenlm # 🌟 ไลบรารีสำหรับโมเดลสถิติภาษา import numpy as np import torch import torch.nn as nn from ultralytics import YOLO # --- 1. การตั้งค่าและโหลดโมเดล (Settings & Model Loading) --- print("⏳ กำลังเริ่มต้นระบบโครงข่ายประสาทเทียม... กรุณารอสักครู่") # ปิด warning ของ asyncio ที่ไม่มีผลต่อการทำงาน warnings.filterwarnings("ignore", category=ResourceWarning) asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy()) DETECTOR_MODEL_PATH = "bestV4.pt" # โมเดลตัดบรรทัด (YOLO) OCR_MODEL_PATH = "best_ocr_model_cer_0.0561.pth" # โมเดลอ่านข้อความ (CRNN) VOCAB_FILE_PATH = "vocab1.txt" # ไฟล์พจนานุกรม LM_MODEL_PATH = "shan_lm_v2.binary" # โมเดลภาษา KenLM ภาษาไทใหญ่ V2 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"🖥️ ระบบกำลังประมวลผลผ่าน: {DEVICE}") # --- สถาปัตยกรรมโมเดลและพจนานุกรม (Model Classes) --- class ShanVocab: def __init__(self, file_path): with open(file_path, "r", encoding="utf-8") as f: char_list = list(f.read()) self.vocab = ["[BLANK]"] + sorted(list(set(char_list))) if "[UNK]" not in self.vocab: self.vocab.insert(1, "[UNK]") self.char_to_idx = {char: i for i, char in enumerate(self.vocab)} self.idx_to_char = {i: char for i, char in enumerate(self.vocab)} self.blank_idx = self.char_to_idx["[BLANK]"] def __len__(self): return len(self.vocab) def ctc_decode(predictions, idx_to_char, blank_idx=0): """แปลงค่า Logits/Predictions จากโมเดล CRNN/CTC ให้เป็นข้อความภาษาไทใหญ่ โดยข้ามตัวอักษรซ้ำติดต่อกันและข้าม blank token (blank_idx) """ argmax_preds = torch.argmax(predictions, dim=2).permute( 1, 0 ) # [batch, seq_len] decoded_texts = [] for pred in argmax_preds: decoded_chars = [] prev_idx = None for idx in pred: idx = idx.item() if idx != blank_idx and idx != prev_idx: decoded_chars.append(idx_to_char.get(idx, "")) prev_idx = idx decoded_texts.append("".join(decoded_chars)) return decoded_texts class ShanCRNN(nn.Module): def __init__(self, num_chars): super().__init__() self.cnn_p1 = nn.Sequential( nn.Conv2d(3, 64, 3, 1, 1), nn.ReLU(True), nn.MaxPool2d(2, 2) ) self.cnn_p2 = nn.Sequential( nn.Conv2d(64, 128, 3, 1, 1), nn.ReLU(True), nn.MaxPool2d(2, 2) ) self.cnn_p3 = nn.Sequential( nn.Conv2d(128, 256, 3, 1, 1), nn.BatchNorm2d(256), nn.ReLU(True), nn.Conv2d(256, 256, 3, 1, 1), nn.ReLU(True), nn.MaxPool2d((2, 2), (2, 1), (0, 1)), ) self.cnn_p4 = nn.Sequential( nn.Conv2d(256, 512, 3, 1, 1), nn.BatchNorm2d(512), nn.ReLU(True), nn.Conv2d(512, 512, 3, 1, 1), nn.ReLU(True), nn.MaxPool2d((2, 2), (2, 1), (0, 1)), ) self.cnn_p5 = nn.Sequential( nn.Conv2d(512, 512, kernel_size=(4, 1), stride=1, padding=0), nn.BatchNorm2d(512), nn.ReLU(True), ) self.lstm = nn.LSTM( 512, 256, bidirectional=True, num_layers=2, batch_first=True ) self.output = nn.Linear(512, num_chars) def forward(self, x): x = self.cnn_p1(x) x = self.cnn_p2(x) x = self.cnn_p3(x) x = self.cnn_p4(x) x = self.cnn_p5(x) # [batch, 512, 1, width] x = x.squeeze(2) # [batch, 512, width] x = x.permute(0, 2, 1) # [batch, width, 512] x, _ = self.lstm(x) # [batch, width, 512] x = self.output(x) # [batch, width, num_chars] x = x.permute(1, 0, 2) # [width, batch, num_chars] ให้ตรงตามมาตรฐาน CTC return x def kenlm_correct_sentence(text_line): """ปรับปรุงการแก้ไขวรรณยุกต์ไทใหญ่ด้วย KenLM ให้มีความเป็นธรรมชาติ ป้องกันการยัดเยียดวรรณยุกต์เกินความจำเป็น และทำความสะอาดวรรณยุกต์ซ้อน """ if lm_model is None or not text_line.strip(): return text_line shan_tones = ["ႇ", "ႈ", "း", "ႉ", "ႊ", "ႆ"] # 1. ทำความสะอาดข้อความเบื้องต้น: ลบวรรณยุกต์ที่ติดซ้ำสองตัวออกก่อน (เช่น းး -> း, ႈႈ -> ႈ) for tone in shan_tones: double_tone = tone + tone while double_tone in text_line: text_line = text_line.replace(double_tone, tone) # 2. แยก Token (หากไม่มีช่องว่าง จะใช้วิธีคงประโยคไว้เพื่อประมวลผลบริบท) tokens = text_line.split() if " " in text_line else list(text_line) corrected_tokens = [] for token in tokens: if not token.strip(): corrected_tokens.append(token) continue # หากคำนั้นมีวรรณยุกต์ไทใหญ่อยู่แล้ว ให้ข้าม ไม่ต้องไปเติมเพิ่ม has_tone = any(char in shan_tones for char in token) if has_tone: corrected_tokens.append(token) continue # ถ้าเป็นตัวอักษรเดี่ยวๆ หรือเครื่องหมาย ไม่ต้องเติมวรรณยุกต์ if len(token) == 1 and token not in [ "ၵ", "ၶ", "င", "ၸ", "သ", "ၺ", "တ", "ထ", "ၼ", "ပ", "ၽ", "ၾ", "မ", "ယ", "ရ", "ႁ", "လ", "ဝ", "ဢ", ]: corrected_tokens.append(token) continue # คำนวณคะแนนคำเดิม best_token = token base_score = lm_model.score(" ".join(corrected_tokens + [token])) best_score = base_score # ทดลองเติมวรรณยุกต์ทีละตัว แล้วเช็คว่าคะแนนดีขึ้นอย่างมีนัยสำคัญหรือไม่ (Margin Threshold) for tone in shan_tones: candidate = token + tone test_sentence = " ".join(corrected_tokens + [candidate]) score = lm_model.score(test_sentence) # ตั้งค่า Threshold: คะแนนต้องดีกว่าคำเดิมอย่างน้อย 2.0 log-probability ถึงจะยอมเปลี่ยน if score > (best_score + 2.0): best_score = score best_token = candidate corrected_tokens.append(best_token) join_char = " " if " " in text_line else "" result_text = join_char.join(corrected_tokens) # ทำความสะอาดซ้ำอีกรอบเพื่อความชัวร์ for tone in shan_tones: for tone2 in shan_tones: if (tone + tone2) in result_text: result_text = result_text.replace(tone + tone2, tone) return result_text # 🛠️ ฟังก์ชันจัดระเบียบพิกัดกล่องข้อความ (Box Optimization Algorithm) def optimize_and_pad_boxes( boxes, img_shape, max_gap=80, pad_left=15, pad_right=65, pad_y=12 ): if len(boxes) == 0: return [] h_img, w_img = img_shape[:2] boxes = sorted(boxes, key=lambda b: b[1]) rows = [] current_row = [boxes[0]] for i in range(1, len(boxes)): prev_box = current_row[-1] current_box = boxes[i] current_y_center = (current_box[1] + current_box[3]) / 2.0 if current_y_center < prev_box[3]: current_row.append(current_box) else: rows.append(current_row) current_row = [current_box] rows.append(current_row) final_processed_boxes = [] for row in rows: row = sorted(row, key=lambda b: b[0]) merged_row = [] curr_box = row[0] for next_box in row[1:]: gap = next_box[0] - curr_box[2] if gap <= max_gap: curr_box = [ min(curr_box[0], next_box[0]), min(curr_box[1], next_box[1]), max(curr_box[2], next_box[2]), max(curr_box[3], next_box[3]), ] else: merged_row.append(curr_box) curr_box = next_box merged_row.append(curr_box) for box in merged_row: x1, y1, x2, y2 = box x1 = max(0, int(x1 - pad_left)) y1 = max(0, int(y1 - pad_y)) x2 = min(w_img, int(x2 + pad_right)) y2 = min(h_img, int(y2 + pad_y)) if (w_img - x2) < (w_img * 0.12): x2 = w_img final_processed_boxes.append([x1, y1, x2, y2]) return final_processed_boxes # --- โหลดน้ำหนักโมเดลเข้าสู่หน่วยความจำ --- try: print(f" - กำลังโหลด YOLO Detector จาก: {DETECTOR_MODEL_PATH}") line_detector_model = YOLO(DETECTOR_MODEL_PATH) print(f" - กำลังโหลด OCR Model จาก: {OCR_MODEL_PATH}") vocab_handler = ShanVocab(VOCAB_FILE_PATH) ocr_model = ShanCRNN(num_chars=len(vocab_handler)) ocr_model.load_state_dict(torch.load(OCR_MODEL_PATH, map_location=DEVICE)) ocr_model.to(DEVICE) ocr_model.eval() # โหลดโมเดลภาษา KenLM if os.path.exists(LM_MODEL_PATH): print(f" - กำลังโหลด KenLM Language Model จาก: {LM_MODEL_PATH}") lm_model = kenlm.Model(LM_MODEL_PATH) print("✅ โหลดโมเดลทั้งหมดรวมถึง KenLM สำเร็จ! ระบบพร้อมทำงานเต็มรูปแบบ") else: print( f"⚠️ คำเตือน: ไม่พบไฟล์ {LM_MODEL_PATH}" " ระบบจะรัน OCR โดยไม่มีตัวกรองคำผิด" ) lm_model = None except FileNotFoundError as e: print(f"❌ ไม่พบไฟล์โมเดล: {e}") line_detector_model = None ocr_model = None lm_model = None except Exception as e: print(f"❌ เกิดข้อผิดพลาดร้ายแรง: {e}") line_detector_model = None ocr_model = None lm_model = None # --- 2. ฟังก์ชันประมวลผลหลัก (Main Pipeline) --- def process_image(image_input, progress=gr.Progress()): if line_detector_model is None or ocr_model is None: raise gr.Error( "❌ ไม่พบไฟล์โมเดล! กรุณาตรวจสอบไฟล์โมเดลและพจนานุกรมในไดเรกทอรี" ) if image_input is None: raise gr.Error("⚠️ กรุณาอัปโหลดรูปภาพก่อน (Please upload an image)") try: full_image_bgr = cv2.cvtColor(image_input, cv2.COLOR_RGB2BGR) annotated_image = image_input.copy() # --- STEP 1: ตรวจจับพิกัดบรรทัดด้วย YOLO --- progress(0.1, desc="🔍 กำลังตรวจจับโครงสร้างหน้ากระดาษ...") line_results = line_detector_model( full_image_bgr, verbose=False, iou=0.5, conf=0.3 ) raw_boxes = line_results[0].boxes.xyxy.cpu().numpy() if len(raw_boxes) == 0: return "⚠️ ไม่พบข้อความในภาพ (No text detected)", image_input # --- STEP 2: จัดระเบียบกล่องข้อความและเย็บประโยคยาว --- progress(0.2, desc="🛠️ กำลังผสานกล่องบรรทัดยาวและกู้คืนสระ...") grouped_lines = optimize_and_pad_boxes( raw_boxes, full_image_bgr.shape, max_gap=80, pad_left=15, pad_right=65, # ขยายขอบขวากันวรรณยุกต์ท้ายคำหลุด pad_y=12, ) for idx, box in enumerate(grouped_lines): x1, y1, x2, y2 = map(int, box) cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.putText( annotated_image, f"L{idx+1}", (x1 + 5, y1 + 22), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2, ) # --- STEP 3: ถอดรหัสอักษรไทใหญ่ด้วย OCR พร้อมการกรองด้วย KenLM V2 --- progress(0.4, desc="📖 กำลังประมวลผลตัวอักษรไทใหญ่ด้วย AI...") final_text_lines = [] total_lines = len(grouped_lines) with torch.inference_mode(): for i, box in enumerate(grouped_lines): progress( 0.4 + (0.6 * (i / total_lines)), desc=f"ถอดรหัสข้อความบรรทัดที่ {i+1}/{total_lines}", ) x1, y1, x2, y2 = map(int, box) cropped_img = full_image_bgr[y1:y2, x1:x2] if cropped_img.size == 0: continue h_orig, w_orig = cropped_img.shape[:2] if h_orig == 0 or w_orig == 0: continue new_width = int(w_orig * (64 / h_orig)) new_width = max(32, (new_width // 4) * 4) if new_width > 1600: new_width = 1600 cropped_rgb = cv2.cvtColor(cropped_img, cv2.COLOR_BGR2RGB) transform = A.Compose([ A.Resize(height=64, width=new_width), A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)), ToTensorV2(), ]) image_tensor = ( transform(image=cropped_rgb)["image"].unsqueeze(0).to(DEVICE) ) # ถอดพิกเซลภาพเป็นข้อความ outputs = ocr_model(image_tensor) raw_text = ctc_decode( outputs.cpu(), vocab_handler.idx_to_char, vocab_handler.blank_idx )[0] # ตรวจสอบและปรับปรุงด้วย KenLM V2 corrected_text = kenlm_correct_sentence(raw_text) final_text_lines.append(corrected_text) progress(1.0, desc="✅ ถอดความเสร็จสิ้น!") return "\n".join(final_text_lines), annotated_image except Exception as e: raise gr.Error(f"❌ เกิดข้อผิดพลาดภายในระบบ: {str(e)}") # --- 3. ส่วนสร้างหน้าต่างการแสดงผล (Gradio Interface) --- custom_css = """ .gradio-container {font-family: 'Noto Sans Myanmar', sans-serif;} h1 {color: #2c3e50; text-align: center;} footer {display: none !important;} """ with gr.Blocks(title="Panya Tai OCR API") as demo: gr.Markdown("# 🦉 Panya Tai OCR API (v3.2 - KenLM Language Model Integrated)") gr.Markdown( "### ระบบแปลงรูปภาพเป็นข้อความภาษาไทใหญ่" " พร้อมตัวดักจับวรรณยุกต์เพี้ยนอัตโนมัติ" ) with gr.Row(): with gr.Column(scale=1): image_input = gr.Image( type="numpy", label="1. သႂ်ႇၶႅပ်းႁၢင်ႈ (อัปโหลดรูปภาพต้นฉบับ)", height=400, ) run_btn = gr.Button( "🚀 ဢၢၼ်ႇလိၵ်ႈ (เริ่มอ่านข้อความ)", variant="primary", size="lg" ) with gr.Column(scale=1): output_text = gr.Textbox( label=( "2. ၽွၼ်းလႆႈ (ผลลัพธ์ข้อความภาษาไทใหญ่ - ผ่านการแก้วรรณยุกต์แล้ว)" ), placeholder="ข้อความภาษาไทใหญ่จะปรากฏที่นี่หลังจากประมวลผล...", lines=10, interactive=False, ) output_image = gr.Image( label=( "3. ႁၢင်ႈတတ်းထႅဝ်" " (ผลการวิเคราะห์การตีกรอบตัดบรรทัดหลังการเพิ่มประสิทธิภาพ)" ), type="numpy", interactive=False, ) run_btn.click( fn=process_image, inputs=image_input, outputs=[output_text, output_image], api_name="predict", ) if __name__ == "__main__": demo.launch(theme=gr.themes.Soft(), css=custom_css)