raidendg/vietnamese_bill_ocr

1import os2import io3import numpy as np4import gradio as gr5from PIL import Image6import google.generativeai as genai7from paddleocr import PaddleOCR8 9# Configure Gemini API10GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', '')11 12if not GEMINI_API_KEY:13    print("⚠️ WARNING: GEMINI_API_KEY not set!")14    print("Please set it in Hugging Face Space settings under 'Settings > Repository secrets'")15else:16    print(f"✅ Gemini API Key configured")17    genai.configure(api_key=GEMINI_API_KEY)18 19# Initialize Gemini model20try:21    text_model = genai.GenerativeModel('gemini-2.0-flash-exp')22    print("✅ Gemini model initialized: gemini-2.0-flash-exp")23except Exception as e:24    print(f"⚠️ Error initializing gemini-2.0-flash-exp: {e}")25    print("Falling back to gemini-1.5-flash")26    text_model = genai.GenerativeModel('gemini-1.5-flash')27 28# Initialize PaddleOCR29print("🔍 Initializing PaddleOCR...")30ocr = PaddleOCR(use_angle_cls=True, lang='vi', show_log=False, use_gpu=False)31print("✅ PaddleOCR initialized")32 33def extract_text_from_image(image):34    """Extract text from image using PaddleOCR"""35    try:36        # Convert PIL Image to numpy array37        img_array = np.array(image)38        39        # Perform OCR40        result = ocr.ocr(img_array, cls=True)41        42        if not result or not result[0]:43            return None44        45        # Extract text from result46        extracted_lines = []47        for line in result[0]:48            if line and len(line) > 1:49                text = line[1][0]50                extracted_lines.append(text)51        52        extracted_text = ' '.join(extracted_lines)53        print(f"[PaddleOCR] Extracted: {extracted_text}")54        55        return extracted_text if extracted_text else None56    except Exception as e:57        print(f"Error in OCR: {e}")58        return None59 60def restore_vietnamese_diacritics(text):61    """Restore Vietnamese diacritics using Gemini API"""62    try:63        prompt = f"""64        Bạn là chuyên gia về tiếng Việt. Nhiệm vụ của bạn là phục hồi dấu tiếng Việt cho văn bản không dấu.65 66        Quy tắc:67        1. Phục hồi chính xác dấu thanh và dấu phụ cho tiếng Việt68        2. Hiểu ngữ cảnh về giao dịch chuyển tiền, chuyển khoản69        3. Các từ viết tắt phổ biến:70           - "ck" có thể là "chuyển khoản"71           - "chuyen khoan" → "chuyển khoản"72           - "du lich" → "du lịch"73           - "tien" → "tiền"74        4. Giữ nguyên số, ký hiệu đặc biệt75        5. Tên riêng cần viết hoa chữ cái đầu và có dấu chính xác76        6. CHỈ trả về văn bản đã được phục hồi dấu, KHÔNG thêm giải thích hay văn bản khác77 78        Văn bản cần phục hồi dấu:79        {text}80 81        Văn bản đã có dấu:82        """83        84        print(f"[Gemini] Restoring diacritics...")85        response = text_model.generate_content(prompt)86        87        # Try to get text from response88        try:89            restored_text = response.text.strip()90            print(f"[Gemini] Restored: {restored_text}")91            return restored_text92        except Exception as text_error:93            print(f"[Gemini] Method 1 failed, trying candidates...")94            try:95                if response.candidates and len(response.candidates) > 0:96                    candidate = response.candidates[0]97                    if candidate.content and candidate.content.parts:98                        restored_text = candidate.content.parts[0].text.strip()99                        print(f"[Gemini] Restored (via candidates): {restored_text}")100                        return restored_text101            except Exception as candidate_error:102                print(f"[Gemini] Both methods failed: {candidate_error}")103            return None104    except Exception as e:105        print(f"[ERROR] Error restoring diacritics: {e}")106        return None107 108def process_image(image):109    """Process image: OCR + restore diacritics"""110    if image is None:111        return "❌ Vui lòng upload ảnh", "", ""112    113    # Extract text using OCR114    extracted_text = extract_text_from_image(image)115    116    if not extracted_text:117        return "❌ Không thể nhận diện text từ ảnh. Vui lòng dùng ảnh rõ ràng hơn.", "", ""118    119    # Restore diacritics120    restored_text = restore_vietnamese_diacritics(extracted_text)121    122    if not restored_text:123        return "❌ Không thể phục hồi dấu. Vui lòng thử lại.", extracted_text, ""124    125    return "✅ Xử lý thành công!", extracted_text, restored_text126 127def process_text(text):128    """Process text only: restore diacritics"""129    if not text or not text.strip():130        return "❌ Vui lòng nhập văn bản", ""131    132    # Restore diacritics133    restored_text = restore_vietnamese_diacritics(text.strip())134    135    if not restored_text:136        return "❌ Không thể phục hồi dấu. Vui lòng thử lại.", ""137    138    return "✅ Xử lý thành công!", restored_text139 140# Create Gradio interface141with gr.Blocks(title="Vietnamese OCR - Phục hồi dấu tiếng Việt", theme=gr.themes.Soft()) as demo:142    gr.Markdown("""143    # 🇻🇳 Vietnamese OCR - Phục hồi dấu tiếng Việt144    145    Ứng dụng nhận diện và phục hồi dấu tiếng Việt từ văn bản không dấu.146    147    **Công nghệ:**148    - 🔍 PaddleOCR - Nhận diện text từ ảnh149    - 🤖 Google Gemini AI - Phục hồi dấu tiếng Việt150    """)151    152    with gr.Tabs():153        # Tab 1: Upload Image154        with gr.Tab("📷 Upload Ảnh"):155            gr.Markdown("### Upload ảnh chứa văn bản tiếng Việt không dấu")156            157            with gr.Row():158                with gr.Column():159                    image_input = gr.Image(type="pil", label="Upload ảnh")160                    image_btn = gr.Button("🚀 Xử lý ảnh", variant="primary", size="lg")161                162                with gr.Column():163                    image_status = gr.Textbox(label="Trạng thái", interactive=False)164                    image_original = gr.Textbox(label="📄 Văn bản gốc (không dấu)", lines=3, interactive=False)165                    image_restored = gr.Textbox(label="✅ Văn bản đã có dấu", lines=3, interactive=False)166            167            gr.Markdown("""168            **Ví dụ:**169            - Upload ảnh chứa text: "Tung ck du lich"170            - Kết quả: "Tùng chuyển khoản du lịch"171            """)172        173        # Tab 2: Input Text174        with gr.Tab("✏️ Nhập Text"):175            gr.Markdown("### Nhập văn bản tiếng Việt không dấu")176            177            with gr.Row():178                with gr.Column():179                    text_input = gr.Textbox(180                        label="Nhập văn bản không dấu",181                        placeholder="Ví dụ: Tung ck du lich",182                        lines=5183                    )184                    text_btn = gr.Button("🚀 Phục hồi dấu", variant="primary", size="lg")185                186                with gr.Column():187                    text_status = gr.Textbox(label="Trạng thái", interactive=False)188                    text_restored = gr.Textbox(label="✅ Văn bản đã có dấu", lines=5, interactive=False)189            190            gr.Examples(191                examples=[192                    ["Tung ck du lich"],193                    ["Chuyen tien mua sam"],194                    ["Thanh toan hoa don internet thang 12"],195                    ["Nguyen Van A chuyen khoan tien dien"],196                ],197                inputs=text_input,198                label="📝 Ví dụ (click để thử)"199            )200    201    # Event handlers202    image_btn.click(203        fn=process_image,204        inputs=[image_input],205        outputs=[image_status, image_original, image_restored]206    )207    208    text_btn.click(209        fn=process_text,210        inputs=[text_input],211        outputs=[text_status, text_restored]212    )213    214    gr.Markdown("""215    ---216    ### 📚 Thông tin217    218    **Use cases:**219    - Xử lý tin nhắn chuyển khoản không dấu220    - OCR tài liệu văn bản không dấu221    - Chuẩn hóa dữ liệu text tiếng Việt222    223    **Giới hạn:**224    - Hỗ trợ tốt nhất cho chữ in (printed text)225    - Ảnh cần rõ ràng, dễ đọc226    227    **Made with ❤️ using PaddleOCR & Google Gemini AI**228    """)229 230# Launch the app231if __name__ == "__main__":232    demo.launch()233 234