raidendg/vietnamese_ocr_bill

1import os2import io3import numpy as np4import gradio as gr5from PIL import Image6import google.generativeai as genai7from paddleocr import PaddleOCR8 9# Initialize PaddleOCR10print("🔍 Initializing PaddleOCR...")11ocr = PaddleOCR(use_angle_cls=True, lang='vi', show_log=False, use_gpu=False)12print("✅ PaddleOCR initialized")13 14def extract_text_from_image(image):15    """Extract text from image using PaddleOCR"""16    try:17        img_array = np.array(image)18        result = ocr.ocr(img_array, cls=True)19        20        if not result or not result[0]:21            return None22        23        extracted_lines = []24        for line in result[0]:25            if line and len(line) > 1:26                text = line[1][0]27                extracted_lines.append(text)28        29        extracted_text = ' '.join(extracted_lines)30        print(f"[PaddleOCR] Extracted: {extracted_text}")31        32        return extracted_text if extracted_text else None33    except Exception as e:34        print(f"Error in OCR: {e}")35        return None36 37def restore_vietnamese_diacritics(text, api_key):38    """Restore Vietnamese diacritics using Gemini API"""39    if not api_key or not api_key.strip():40        return None, "❌ Vui lòng nhập Gemini API Key"41    42    try:43        # Configure Gemini with provided API key44        genai.configure(api_key=api_key.strip())45        46        # Initialize model47        try:48            model = genai.GenerativeModel('gemini-2.5-flash')49        except:50            model = genai.GenerativeModel('gemini-1.5-flash')51        52        prompt = f"""53        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.54 55        Quy tắc:56        1. Phục hồi chính xác dấu thanh và dấu phụ cho tiếng Việt57        2. Hiểu ngữ cảnh về giao dịch chuyển tiền, chuyển khoản58        3. Các từ viết tắt phổ biến:59           - "ck" có thể là "chuyển khoản"60           - "chuyen khoan" → "chuyển khoản"61           - "du lich" → "du lịch"62           - "tien" → "tiền"63        4. Giữ nguyên số, ký hiệu đặc biệt64        5. Tên riêng cần viết hoa chữ cái đầu và có dấu chính xác65        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ác66 67        Văn bản cần phục hồi dấu:68        {text}69 70        Văn bản đã có dấu:71        """72        73        print(f"[Gemini] Restoring diacritics...")74        response = model.generate_content(prompt)75        76        # Try to get text from response77        try:78            restored_text = response.text.strip()79            print(f"[Gemini] Restored: {restored_text}")80            return restored_text, None81        except Exception as text_error:82            print(f"[Gemini] Method 1 failed, trying candidates...")83            try:84                if response.candidates and len(response.candidates) > 0:85                    candidate = response.candidates[0]86                    if candidate.content and candidate.content.parts:87                        restored_text = candidate.content.parts[0].text.strip()88                        print(f"[Gemini] Restored (via candidates): {restored_text}")89                        return restored_text, None90            except Exception as candidate_error:91                print(f"[Gemini] Both methods failed: {candidate_error}")92            return None, "❌ Không thể lấy response từ Gemini"93    except Exception as e:94        error_msg = str(e)95        print(f"[ERROR] Error restoring diacritics: {e}")96        if "API_KEY_INVALID" in error_msg or "invalid" in error_msg.lower():97            return None, "❌ API Key không hợp lệ"98        return None, f"❌ Lỗi: {error_msg}"99 100def process_image(image, api_key):101    """Process image: OCR + restore diacritics"""102    if image is None:103        return "❌ Vui lòng upload ảnh", "", ""104    105    if not api_key or not api_key.strip():106        return "❌ Vui lòng nhập Gemini API Key", "", ""107    108    # Extract text using OCR109    extracted_text = extract_text_from_image(image)110    111    if not extracted_text:112        return "❌ Không thể nhận diện text từ ảnh. Vui lòng dùng ảnh rõ ràng hơn.", "", ""113    114    # Restore diacritics115    restored_text, error = restore_vietnamese_diacritics(extracted_text, api_key)116    117    if error:118        return error, extracted_text, ""119    120    if not restored_text:121        return "❌ Không thể phục hồi dấu. Vui lòng thử lại.", extracted_text, ""122    123    return "✅ Xử lý thành công!", extracted_text, restored_text124 125def process_text(text, api_key):126    """Process text only: restore diacritics"""127    if not text or not text.strip():128        return "❌ Vui lòng nhập văn bản", ""129    130    if not api_key or not api_key.strip():131        return "❌ Vui lòng nhập Gemini API Key", ""132    133    # Restore diacritics134    restored_text, error = restore_vietnamese_diacritics(text.strip(), api_key)135    136    if error:137        return error, ""138    139    if not restored_text:140        return "❌ Không thể phục hồi dấu. Vui lòng thử lại.", ""141    142    return "✅ Xử lý thành công!", restored_text143 144# Create Gradio interface145with gr.Blocks(title="Vietnamese Bill OCR", theme=gr.themes.Soft()) as demo:146    gr.Markdown("""147    # 🇻🇳 Vietnamese Bill OCR148    149    Ứ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.150    151    **Công nghệ:**152    - 🔍 PaddleOCR - Nhận diện text từ ảnh153    - 🤖 Google Gemini AI - Phục hồi dấu tiếng Việt154    """)155    156    # API Key input (shared across tabs)157    with gr.Row():158        api_key_input = gr.Textbox(159            label="🔑 Gemini API Key",160            placeholder="Nhập API key của bạn (lấy tại: https://makersuite.google.com/app/apikey)",161            type="password",162            value=os.getenv('GEMINI_API_KEY', '')163        )164    165    gr.Markdown("""166    💡 **Lấy API Key miễn phí tại:** [Google AI Studio](https://makersuite.google.com/app/apikey)167    """)168    169    with gr.Tabs():170        # Tab 1: Upload Image171        with gr.Tab("📷 Upload Ảnh"):172            gr.Markdown("### Upload ảnh chứa văn bản tiếng Việt không dấu")173            174            with gr.Row():175                with gr.Column():176                    image_input = gr.Image(type="pil", label="Upload ảnh")177                    image_btn = gr.Button("🚀 Xử lý ảnh", variant="primary", size="lg")178                179                with gr.Column():180                    image_status = gr.Textbox(label="Trạng thái", interactive=False)181                    image_original = gr.Textbox(label="📄 Văn bản gốc (không dấu)", lines=3, interactive=False)182                    image_restored = gr.Textbox(label="✅ Văn bản đã có dấu", lines=3, interactive=False)183            184            gr.Markdown("""185            **Ví dụ:**186            - Upload ảnh chứa text: "Tung ck du lich"187            - Kết quả: "Tùng chuyển khoản du lịch"188            """)189        190        # Tab 2: Input Text191        with gr.Tab("✏️ Nhập Text"):192            gr.Markdown("### Nhập văn bản tiếng Việt không dấu")193            194            with gr.Row():195                with gr.Column():196                    text_input = gr.Textbox(197                        label="Nhập văn bản không dấu",198                        placeholder="Ví dụ: Tung ck du lich",199                        lines=5200                    )201                    text_btn = gr.Button("🚀 Phục hồi dấu", variant="primary", size="lg")202                203                with gr.Column():204                    text_status = gr.Textbox(label="Trạng thái", interactive=False)205                    text_restored = gr.Textbox(label="✅ Văn bản đã có dấu", lines=5, interactive=False)206            207            gr.Examples(208                examples=[209                    ["Tung ck du lich"],210                    ["Chuyen tien mua sam"],211                    ["Thanh toan hoa don internet thang 12"],212                    ["Nguyen Van A chuyen khoan tien dien"],213                ],214                inputs=text_input,215                label="📝 Ví dụ (click để thử)"216            )217    218    # Event handlers219    image_btn.click(220        fn=process_image,221        inputs=[image_input, api_key_input],222        outputs=[image_status, image_original, image_restored]223    )224    225    text_btn.click(226        fn=process_text,227        inputs=[text_input, api_key_input],228        outputs=[text_status, text_restored]229    )230    231    gr.Markdown("""232    ---233    ### 📚 Thông tin234    235    **Use cases:**236    - Xử lý tin nhắn chuyển khoản không dấu237    - OCR hóa đơn, bill chuyển tiền238    - Chuẩn hóa dữ liệu text tiếng Việt239    240    **Giới hạn:**241    - Hỗ trợ tốt nhất cho chữ in (printed text)242    - Ảnh cần rõ ràng, dễ đọc243    244    **Made with ❤️ using PaddleOCR & Google Gemini AI**245    """)246 247# Launch the app248if __name__ == "__main__":249    demo.launch()250