Update app.py

app.py

@@ -1,33 +1,45 @@
 import gradio as gr
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
+import numpy as np
 
-name = "chanelcolgate/trocr-base-printed_captcha_ocr"
-model = VisionEncoderDecoderModel.from_pretrained(name)
+# Load model và processor
+model_name = "chanelcolgate/trocr-base-printed_captcha_ocr"
+model = VisionEncoderDecoderModel.from_pretrained(model_name)
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed")
 
-
+# Hàm xử lý ảnh captcha
 def process_image(image):
-    # Chuyển nền trong suốt thành trắng (nếu có alpha channel)
+    # Chuyển numpy -> PIL nếu cần
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+
+    # Thêm nền trắng nếu ảnh có alpha
     if image.mode in ("RGBA", "LA"):
         background = Image.new("RGB", image.size, (255, 255, 255))
-        image = Image.alpha_composite(background, image.convert("RGBA"))
+        background.paste(image, mask=image.split()[-1])
+        image = background
+    else:
+        image = image.convert("RGB")
 
-    # Xử lý ảnh
+    # Encode và predict
     pixel_values = processor(image, return_tensors="pt").pixel_values
     generated_ids = model.generate(pixel_values)
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     return generated_text
 
-
-title = "Interactive demo: Captcha OCR"
-description = "Demo tracuumasothue captcha"
+# Giao diện Gradio
+title = "Captcha OCR Demo"
+description = "Nhận diện captcha từ mã số thuế (MST) – Model TrOCR"
 
 interface = gr.Interface(
     fn=process_image,
     inputs="image",
-    examples=[f"examples/captcha-{i}.png" for i in range(10)],
     outputs="text",
+    examples=[f"examples/captcha-{i}.png" for i in range(10)],
     title=title,
-    description=description
-).launch()
+    description=description,
+)
+
+if __name__ == "__main__":
+    interface.launch()