app.py

import os

# --- CẤU HÌNH HỆ THỐNG ---
os.environ["FLAGS_use_mkldnn"] = "0"
os.environ["FLAGS_enable_mkldnn"] = "0"
os.environ["DN_ENABLE_MKLDNN"] = "0"
os.environ["CPP_MIN_LOG_LEVEL"] = "3"

import logging
import re
import gradio as gr
from paddleocr import PaddleOCR
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import requests

# Tắt log thừa
logging.getLogger("ppocr").setLevel(logging.WARNING)

print("Đang khởi tạo PaddleOCR (Coordinate Sync Mode)...")

try:
    ocr = PaddleOCR(use_textline_orientation=False, use_doc_orientation_classify=False, 
        use_doc_unwarping=False, lang='ch')
except Exception as e:
    print(f"Lỗi khởi tạo: {e}. Chuyển về chế độ mặc định.")
    ocr = PaddleOCR(lang='ch')

print("Model đã sẵn sàng!")

# --- TẢI FONT ---
def check_and_download_font():
    font_path = "./simfang.ttf"
    if not os.path.exists(font_path):
        try:
            url = "https://github.com/StellarCN/scp_zh/raw/master/fonts/SimFang.ttf"
            r = requests.get(url, allow_redirects=True)
            with open(font_path, 'wb') as f:
                f.write(r.content)
        except:
            return None
    return font_path

FONT_PATH = check_and_download_font()

# --- HÀM VẼ ĐA NĂNG ---
def universal_draw(image, raw_data, font_path):
    if image is None: return image
    
    # Đảm bảo image là PIL
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    # Copy để vẽ
    canvas = image.copy()
    draw = ImageDraw.Draw(canvas)
    
    try:
        font_size = 24
        font = ImageFont.truetype(font_path, font_size) if font_path else ImageFont.load_default()
    except:
        font = ImageFont.load_default()

    # Hàm parse box
    def parse_box(b):
        try:
            if hasattr(b, 'tolist'): b = b.tolist()
            if len(b) > 0 and isinstance(b[0], list): return [tuple(p) for p in b]
            if len(b) == 4 and isinstance(b[0], (int, float)): 
                 return [(b[0], b[1]), (b[2], b[1]), (b[2], b[3]), (b[0], b[3])]
            return None
        except: return None

    items_to_draw = []

    # Logic tìm box/text
    # Ưu tiên cấu trúc PaddleX: rec_texts + dt_polys
    processed = False
    if isinstance(raw_data, list) and len(raw_data) > 0 and isinstance(raw_data[0], dict):
        data_dict = raw_data[0]
        texts = data_dict.get('rec_texts')
        boxes = data_dict.get('dt_polys', data_dict.get('rec_polys', data_dict.get('dt_boxes')))
        
        if texts and boxes and isinstance(texts, list) and isinstance(boxes, list):
            for i in range(min(len(texts), len(boxes))):
                txt = texts[i]
                box = parse_box(boxes[i])
                if box and txt: items_to_draw.append((box, txt))
            processed = True

    # Fallback Logic
    if not processed:
        def hunt(data):
            if isinstance(data, dict):
                box = None; text = None
                for k in ['points', 'box', 'dt_boxes', 'poly']:
                    if k in data: box = parse_box(data[k]); break
                for k in ['transcription', 'text', 'rec_text', 'label']:
                    if k in data: text = data[k]; break
                if box and text: items_to_draw.append((box, text)); return
                for v in data.values(): hunt(v)
            elif isinstance(data, (list, tuple)):
                if len(data) == 2 and isinstance(data[0], list) and len(data[0]) == 4:
                    box = parse_box(data[0])
                    txt_obj = data[1]
                    text = txt_obj[0] if isinstance(txt_obj, (list, tuple)) else txt_obj
                    if box and isinstance(text, str): items_to_draw.append((box, text)); return
                for item in data: hunt(item)
        hunt(raw_data)

    # Vẽ
    for box, txt in items_to_draw:
        try:
            # Vẽ khung đỏ
            draw.polygon(box, outline="red", width=3)
            # Vẽ chữ
            txt_x, txt_y = box[0]
            if hasattr(draw, "textbbox"):
                text_bbox = draw.textbbox((txt_x, txt_y), txt, font=font, anchor="lb")
                draw.rectangle(text_bbox, fill="red")
                draw.text((txt_x, txt_y), txt, fill="white", font=font, anchor="lb")
            else:
                draw.text((txt_x, txt_y - font_size), txt, fill="white", font=font)
        except: continue

    return canvas

# --- HÀM XỬ LÝ TEXT ---
def deep_extract_text(data):
    found_texts = []
    if isinstance(data, str):
        if len(data.strip()) > 0: return [data]
        return []
    if isinstance(data, (list, tuple)):
        for item in data: found_texts.extend(deep_extract_text(item))
    elif isinstance(data, dict):
        for val in data.values(): found_texts.extend(deep_extract_text(val))
    elif hasattr(data, '__dict__'): found_texts.extend(deep_extract_text(data.__dict__))
    return found_texts

def clean_text_result(text_list):
    cleaned = []
    block_list = ['min', 'max', 'general', 'header', 'footer', 'structure']
    for t in text_list:
        t = t.strip()
        if len(t) < 2 and not any(u'\u4e00' <= c <= u'\u9fff' for c in t): continue
        if t.lower().endswith(('.ttf', '.json', '.pdparams', '.yml', '.log')): continue
        if t.lower() in block_list: continue
        if not re.search(r'[\w\u4e00-\u9fff]', t): continue
        cleaned.append(t)
    return cleaned

# --- MAIN PREDICT ---
def predict(image):
    if image is None: return None, "Chưa có ảnh.", "No Data"

    try:
        # Chuẩn bị ảnh đầu vào
        original_pil = image.copy() if isinstance(image, Image.Image) else Image.fromarray(image).copy()
        image_np = np.array(image)
        
        # 1. OCR
        # Chuyển sang dùng predict() để sửa DeprecationWarning và tương thích tốt hơn với Model v5
        raw_result = ocr.predict(image_np)
        
        # Đảm bảo kết quả là list để tương thích với logic vẽ phía sau (nếu predict trả về generator/single obj)
        if not isinstance(raw_result, list):
            raw_result = [raw_result]
        
        # 2. XỬ LÝ ẢNH ĐỂ VẼ (KEY FIX: Lấy ảnh từ Preprocessor nếu có)
        target_image_for_drawing = original_pil
        
        # Kiểm tra xem Paddle có chỉnh sửa ảnh không (dựa vào key 'doc_preprocessor_res')
        if isinstance(raw_result, list) and len(raw_result) > 0 and isinstance(raw_result[0], dict):
            if 'doc_preprocessor_res' in raw_result[0]:
                proc_res = raw_result[0]['doc_preprocessor_res']
                # Nếu có ảnh đầu ra đã chỉnh sửa (output_img)
                if 'output_img' in proc_res:
                    print("Phát hiện ảnh đã qua xử lý hình học. Đang đồng bộ tọa độ...")
                    numpy_img = proc_res['output_img']
                    target_image_for_drawing = Image.fromarray(numpy_img)

        # 3. Vẽ lên ảnh ĐÚNG (Target Image)
        annotated_image = universal_draw(target_image_for_drawing, raw_result, FONT_PATH)

        # 4. Xử lý Text
        all_texts = deep_extract_text(raw_result)
        final_texts = clean_text_result(all_texts)
        text_output = "\n".join(final_texts) if final_texts else "Không tìm thấy văn bản."

        # Debug Info
        debug_str = str(raw_result)[:1000]
        debug_info = f"Used Image Source: {'Preprocessed' if target_image_for_drawing != original_pil else 'Original'}\nData Preview:\n{debug_str}..."

        return annotated_image, text_output, debug_info

    except Exception as e:
        import traceback
        return image, f"Lỗi: {str(e)}", traceback.format_exc()

# --- GIAO DIỆN ---
with gr.Blocks(title="PaddleOCR Perfect Overlay") as iface:
    gr.Markdown("## PaddleOCR Chinese - High Precision Overlay")
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Input Image")
            submit_btn = gr.Button("RUN OCR", variant="primary")
        
        with gr.Column():
            with gr.Tabs():
                with gr.TabItem("🖼️ Kết quả Khớp Tọa Độ"):
                    output_img = gr.Image(type="pil", label="Overlay Result")
                with gr.TabItem("📝 Văn bản"):
                    output_txt = gr.Textbox(label="Text Content", lines=15)
                with gr.TabItem("🐞 Debug"):
                    output_debug = gr.Textbox(label="Debug Info", lines=15)

    submit_btn.click(
        fn=predict,
        inputs=input_img,
        outputs=[output_img, output_txt, output_debug],
        api_name="predict"
    )

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)