Spaces:
Sleeping
Sleeping
| 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) |