import gradio as gr
import easyocr
from pdf2image import convert_from_path
import numpy as np
from PIL import Image
import os

# โหลดโมเดล EasyOCR เตรียมไว้เลย (โหลดครั้งเดียวตอน Start Space จะได้ทำงานเร็วขึ้น)
print("กำลังโหลดโมเดล EasyOCR...")
reader = easyocr.Reader(['th', 'en'])

def process_file(file_path):
    if file_path is None:
        return "กรุณาอัปโหลดไฟล์"
        
    # เช็กนามสกุลไฟล์
    ext = os.path.splitext(file_path)[1].lower()
    extracted_text = []

    try:
        if ext == '.pdf':
            # กรณีเป็นไฟล์ PDF แปลงเป็นรูปก่อน
            # แนะนำให้ใช้ dpi=200 เพื่อไม่ให้กิน RAM ของ Hugging Face Free Tier มากเกินไป
            pages = convert_from_path(file_path, dpi=200)
            
            for i, page in enumerate(pages):
                extracted_text.append(f"--- หน้าที่ {i+1} ---")
                # แปลงเป็น Numpy array ให้ EasyOCR อ่าน
                img_np = np.array(page)
                result = reader.readtext(img_np, detail=0)
                extracted_text.extend(result)
                extracted_text.append("\n")
                
        elif ext in ['.png', '.jpg', '.jpeg']:
            # กรณีเป็นรูปภาพทั่วไป
            result = reader.readtext(file_path, detail=0)
            extracted_text.extend(result)
            
        else:
            return "รองรับเฉพาะไฟล์ .pdf, .png, .jpg, .jpeg เท่านั้นครับ"

        # รวมข้อความทั้งหมดคั่นด้วยการขึ้นบรรทัดใหม่
        return "\n".join(extracted_text)

    except Exception as e:
        return f"เกิดข้อผิดพลาดในการอ่านไฟล์: {str(e)}"

# สร้างหน้าเว็บ UI ด้วย Gradio
interface = gr.Interface(
    fn=process_file,
    inputs=gr.File(label="อัปโหลดไฟล์รูปภาพ หรือ PDF"),
    outputs=gr.Textbox(label="ข้อความที่สกัดได้", lines=15),
    title="Thai OCR Web App (EasyOCR)",
    description="อัปโหลดไฟล์รูปภาพ หรือไฟล์ PDF เพื่อดึงข้อความภาษาไทยและอังกฤษ",
    theme="default"
)

# สั่งรันแอป
if __name__ == "__main__":
    interface.launch()