Update app.py

app.py

@@ -1,34 +1,45 @@
-import gradio as gr
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import pytesseract
 from PIL import Image
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
-# โหลดโมเดล TrOCR และ Processor จาก Hugging Face
-processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
-model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+# โหลดโมเดล transformers ภาษาไทย สำหรับ text classification (ตัวอย่างใช้ WangchanBERT)
+model_name = "airesearch/wangchanberta-base-att-spm-uncased"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-# ฟังก์ชันหลักสำหรับแสดงข้อความที่ตรวจจับได้จากภาพ
-def ocr_text(image):
-    # ตรวจสอบว่าเป็น RGB
-    if image.mode != "RGB":
-        image = image.convert("RGB")
+# ฟังก์ชัน OCR + วิเคราะห์ข้อความ
+def ocr_and_classify(image):
+    # OCR อ่านข้อความจากภาพ (ตั้งภาษาไทย + อังกฤษ)
+    text = pytesseract.image_to_string(image, lang="tha+eng")
+    
+    if not text.strip():
+        return "❌ ไม่พบข้อความในภาพ"
     
-    # ทำ OCR ด้วย TrOCR
-    pixel_values = processor(images=image, return_tensors="pt").pixel_values
-    generated_ids = model.generate(pixel_values)
-    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    # เตรียม input สำหรับ transformers
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    outputs = model(**inputs)
+    
+    # แปลง logits เป็น probabilities และเลือกคลาสที่มีคะแนนสูงสุด
+    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    pred_class_idx = torch.argmax(probs).item()
+    confidence = probs[0][pred_class_idx].item()
+
+    # สำหรับ WangchanBERT ไม่มี label mapping โดยตรง ต้องมี mapping label เอง (สมมติมี 3 คลาส)
+    labels = ["คลาส 0", "คลาส 1", "คลาส 2"]  # แก้ตามโมเดลที่ใช้จริง
 
-    # แสดงข้อความที่ตรวจจับได้
-    return f"📄 ข้อความที่พบในภาพ:\n{text}"
+    result = f"ข้อความที่อ่านได้:\n{text}\n\nการจำแนกข้อความ:\n{labels[pred_class_idx]} (ความมั่นใจ {confidence:.2%})"
+    return result
 
-# สร้าง Gradio Interface
+# สร้าง Gradio UI
 iface = gr.Interface(
-    fn=ocr_text,
-    inputs=gr.Image(type="pil", label="อัปโหลดภาพที่มีข้อความ"),
-    outputs=gr.Textbox(lines=10, label="ข้อความที่ตรวจจับได้"),
-    title="🧾 อ่านข้อความจากภาพ (OCR)",
-    description="ระบบจะอ่านข้อความจากภาพและแสดงผลลัพธ์ทั้งหมดโดยไม่กรอง"
+    fn=ocr_and_classify,
+    inputs=gr.Image(type="pil", label="อัปโหลดภาพสลิปหรือใบเสร็จ"),
+    outputs=gr.Textbox(lines=15, label="ผลลัพธ์"),
+    title="OCR + วิเคราะห์ข้อความสลิป ด้วย pytesseract + Transformers",
+    description="อ่านข้อความด้วย pytesseract แล้วใช้โมเดล transformers วิเคราะห์ข้อความ"
 )
 
-# เปิดแอป
 if __name__ == "__main__":
     iface.launch()