import gradio as gr import pandas as pd import numpy as np from PIL import Image import io import tempfile import os import zipfile from pathlib import Path import json from typing import List, Dict, Any, Tuple, Optional import logging # Import your custom modules from preprocess import PDFImageProcessor, process_uploaded_pdf from llm_utils import load_model, extract_info_from_image, cleanup_model # Setup logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class PDFProcessingApp: def __init__(self): self.processor = PDFImageProcessor() self.model = None self.tokenizer = None self.current_images = {} # Store images by PDF filename self.processed_results = {} # Store OCR results def load_models(self): """Load ML models on startup""" try: if self.model is None: logger.info("Loading models...") self.model, self.tokenizer = load_model() logger.info("Models loaded successfully!") except Exception as e: logger.error(f"Error loading models: {e}") raise def process_multiple_pdfs(self, pdf_files: List[Any]) -> Tuple[Dict, str]: """Process multiple PDF files and extract images""" if not pdf_files: return {}, "❌ Vui lòng upload ít nhất một file PDF" self.current_images = {} pdf_info = {} try: for pdf_file in pdf_files: filename = Path(pdf_file.name).stem logger.info(f"Processing PDF: {filename}") # Read PDF bytes pdf_bytes = pdf_file.read() # Convert PDF to images images = self.processor.pdf_to_images(pdf_bytes) # Store images self.current_images[filename] = images pdf_info[filename] = { 'total_pages': len(images), 'filename': filename } logger.info(f"Extracted {len(images)} pages from {filename}") # Create gallery data for display gallery_data = self.create_gallery_data() info_text = f"✅ Đã xử lý {len(pdf_files)} file PDF, tổng {sum(info['total_pages'] for info in pdf_info.values())} trang" return gallery_data, info_text except Exception as e: logger.error(f"Error processing PDFs: {e}") return {}, f"❌ Lỗi xử lý PDF: {str(e)}" def create_gallery_data(self) -> Dict: """Create data structure for image gallery""" gallery_data = {} for pdf_name, images in self.current_images.items(): gallery_images = [] for i, img_array in enumerate(images): # Convert numpy array to PIL Image if img_array.dtype != np.uint8: img_array = (img_array * 255).astype(np.uint8) if len(img_array.shape) == 2: # Grayscale pil_img = Image.fromarray(img_array, mode='L') else: # Color pil_img = Image.fromarray(img_array) # Resize for display (keeping aspect ratio) pil_img.thumbnail((300, 400), Image.Resampling.LANCZOS) gallery_images.append(pil_img) gallery_data[pdf_name] = gallery_images return gallery_data def update_image_gallery(self, pdf_files): """Update the image gallery when PDFs are uploaded""" if not pdf_files: return gr.update(value=[], visible=False), gr.update(visible=False), "" gallery_data, info_text = self.process_multiple_pdfs(pdf_files) if not gallery_data: return gr.update(value=[], visible=False), gr.update(visible=False), info_text # Flatten all images for gallery display all_images = [] image_metadata = [] for pdf_name, images in gallery_data.items(): for i, img in enumerate(images): all_images.append(img) image_metadata.append({ 'pdf_name': pdf_name, 'page_num': i + 1, 'display_name': f"{pdf_name} - Trang {i + 1}" }) # Store metadata for later use self.image_metadata = image_metadata return ( gr.update(value=all_images, visible=True), gr.update(visible=True), info_text ) def parse_range_input(self, range_text: str, total_images: int) -> List[int]: """Parse range input like '1-5, 8, 10-12' to list of indices""" if not range_text.strip(): return [] indices = [] parts = range_text.split(',') for part in parts: part = part.strip() if '-' in part: # Range like "1-5" try: start, end = map(int, part.split('-')) indices.extend(range(max(1, start) - 1, min(total_images, end))) except: continue else: # Single number try: num = int(part) if 1 <= num <= total_images: indices.append(num - 1) except: continue return sorted(list(set(indices))) # Remove duplicates and sort def process_selected_images(self, selected_indices: List[int], range_input: str, custom_prompt: str) -> Tuple[str, Any]: """Process selected images with OCR""" if not hasattr(self, 'image_metadata'): return "❌ Vui lòng upload PDF trước", None # Load models if not loaded if self.model is None: try: self.load_models() except Exception as e: return f"❌ Lỗi tải model: {str(e)}", None # Combine selected indices from gallery and range input total_images = len(self.image_metadata) range_indices = self.parse_range_input(range_input, total_images) all_selected = sorted(list(set(selected_indices + range_indices))) if not all_selected: return "❌ Vui lòng chọn ít nhất một ảnh để xử lý", None logger.info(f"Processing {len(all_selected)} selected images") # Group by PDF for organized results pdf_results = {} processed_count = 0 try: for idx in all_selected: if idx >= len(self.image_metadata): continue metadata = self.image_metadata[idx] pdf_name = metadata['pdf_name'] page_num = metadata['page_num'] # Get the image img_array = self.current_images[pdf_name][page_num - 1] # Preprocess image processed_img = self.processor.preprocess_image(img_array) # Save to temp file for OCR with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_file: pil_img = Image.fromarray(processed_img) pil_img.save(tmp_file.name) # Extract text using OCR try: result = extract_info_from_image( tmp_file.name, self.model, self.tokenizer, custom_prompt=custom_prompt if custom_prompt.strip() else None ) # Store result if pdf_name not in pdf_results: pdf_results[pdf_name] = [] pdf_results[pdf_name].append({ 'page': page_num, 'content': result }) processed_count += 1 logger.info(f"Processed page {page_num} of {pdf_name}") except Exception as e: logger.error(f"OCR error for {pdf_name} page {page_num}: {e}") if pdf_name not in pdf_results: pdf_results[pdf_name] = [] pdf_results[pdf_name].append({ 'page': page_num, 'content': f"Lỗi xử lý: {str(e)}" }) # Clean up temp file try: os.unlink(tmp_file.name) except: pass # Create Excel files excel_files = self.create_excel_outputs(pdf_results) status_msg = f"✅ Đã xử lý {processed_count} ảnh từ {len(pdf_results)} file PDF" return status_msg, excel_files except Exception as e: logger.error(f"Error in OCR processing: {e}") return f"❌ Lỗi xử lý: {str(e)}", None def create_excel_outputs(self, pdf_results: Dict) -> Any: """Create Excel files for each PDF with OCR results""" if not pdf_results: return None # Create a zip file containing all Excel files zip_buffer = io.BytesIO() with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file: for pdf_name, results in pdf_results.items(): # Create DataFrame for this PDF df_data = [] for result in results: page_num = result['page'] content = result['content'] # Try to parse the content into structured data # This is a simple example - you might want to enhance this df_data.append({ 'Trang': page_num, 'Nội dung': content, 'Thời gian xử lý': pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S') }) # Create DataFrame df = pd.DataFrame(df_data) # Save to Excel in memory excel_buffer = io.BytesIO() with pd.ExcelWriter(excel_buffer, engine='openpyxl') as writer: df.to_excel(writer, sheet_name='OCR_Results', index=False) # Add to zip excel_filename = f"{pdf_name}_OCR_Results.xlsx" zip_file.writestr(excel_filename, excel_buffer.getvalue()) zip_buffer.seek(0) return zip_buffer.getvalue() # Initialize the app app = PDFProcessingApp() # Create Gradio interface def create_interface(): with gr.Blocks( title="PDF OCR Processor", theme=gr.themes.Soft(), css=""" .gradio-container { max-width: 1200px !important; } .gallery-container { max-height: 600px; overflow-y: auto; } """ ) as demo: gr.Markdown(""" # 📄 PDF OCR Processor **Hướng dẫn sử dụng:** 1. Upload nhiều file PDF cùng lúc 2. Xem preview các trang và chọn trang cần xử lý OCR 3. Tùy chỉnh prompt nếu cần 4. Tải xuống kết quả Excel theo tên file gốc """) with gr.Tab("📤 Upload & Xử lý"): # File upload section with gr.Row(): pdf_files = gr.Files( label="📁 Upload PDF Files (Có thể chọn nhiều file)", file_types=[".pdf"], file_count="multiple" ) # Status and info upload_status = gr.Textbox( label="📊 Trạng thái", interactive=False, value="Chưa có file nào được upload" ) # Image gallery and selection with gr.Row(): with gr.Column(scale=2): image_gallery = gr.Gallery( label="📸 Preview các trang PDF (Click để chọn)", show_label=True, elem_classes=["gallery-container"], columns=3, rows=2, height=400, allow_preview=True, selected_index=None, visible=False ) with gr.Column(scale=1): selection_options = gr.Group(visible=False) with selection_options: gr.Markdown("### 🎯 Tùy chọn chọn ảnh") range_input = gr.Textbox( label="📝 Chọn theo range (VD: 1-5, 8, 10-12)", placeholder="1-5, 8, 10-12", info="Có thể kết hợp với việc click chọn ở gallery" ) custom_prompt = gr.Textbox( label="✏️ Custom Prompt (Tùy chọn)", placeholder="Nhập prompt tùy chỉnh cho OCR...", lines=3, value="Trích xuất dữ liệu các cột: STT, Mã số thuế, Tên người nộp thuế, Địa chỉ, Số tiền thuế nợ, Biện pháp cưỡng chế. Hãy cố gắng đọc rõ những con số hoặc chữ bị đóng dấu và trả về dạng markdown." ) process_btn = gr.Button( "🚀 Bắt đầu xử lý OCR", variant="primary", size="lg" ) # Results section with gr.Row(): processing_status = gr.Textbox( label="⚡ Kết quả xử lý", interactive=False ) with gr.Row(): download_files = gr.File( label="📥 Tải xuống kết quả Excel", visible=False ) with gr.Tab("ℹ️ Hướng dẫn chi tiết"): gr.Markdown(""" ## 📋 Hướng dẫn sử dụng chi tiết ### 1. Upload PDF Files - Click vào "📁 Upload PDF Files" để chọn nhiều file PDF - Hệ thống sẽ tự động chuyển đổi tất cả các trang thành ảnh - Xem trạng thái upload trong mục "📊 Trạng thái" ### 2. Chọn trang cần xử lý **Cách 1: Click chọn trong Gallery** - Xem preview tất cả các trang trong "📸 Preview các trang PDF" - Click vào các trang muốn xử lý (có thể chọn nhiều) - Trang được chọn sẽ có viền xanh **Cách 2: Nhập range** - Sử dụng ô "📝 Chọn theo range" - Định dạng: `1-5, 8, 10-12` (có thể kết hợp range và số đơn) - Ví dụ: `1-3, 7, 10-15` sẽ chọn trang 1,2,3,7,10,11,12,13,14,15 **Cách 3: Kết hợp cả hai** - Có thể vừa click trong gallery vừa nhập range - Hệ thống sẽ gộp tất cả lựa chọn lại ### 3. Tùy chỉnh Prompt (Tùy chọn) - Mặc định: Trích xuất thông tin thuế - Có thể thay đổi để phù hợp với nội dung PDF khác - Ví dụ: "Trích xuất tất cả văn bản trong ảnh" ### 4. Xử lý và Tải xuống - Click "🚀 Bắt đầu xử lý OCR" - Xem tiến trình trong "⚡ Kết quả xử lý" - Tải file Excel kết quả (mỗi PDF sẽ có file Excel riêng) ### 5. Kết quả - File Excel sẽ có tên giống file PDF gốc - Mỗi trang được xử lý sẽ là một dòng trong Excel - Có thông tin trang số và nội dung OCR ## 🔧 Lưu ý kỹ thuật - Hỗ trợ PDF scan và PDF có ảnh - Tự động tiền xử lý ảnh để tăng độ chính xác OCR - Sử dụng AI model Vintern-1B-v3_5 cho OCR tiếng Việt - Kết quả trả về dạng Markdown để dễ đọc """) # Event handlers pdf_files.change( fn=app.update_image_gallery, inputs=[pdf_files], outputs=[image_gallery, selection_options, upload_status] ) process_btn.click( fn=app.process_selected_images, inputs=[image_gallery, range_input, custom_prompt], outputs=[processing_status, download_files] ).then( fn=lambda: gr.update(visible=True), outputs=[download_files] ) return demo # Launch the app if __name__ == "__main__": # Pre-load models for faster processing try: logger.info("Pre-loading models...") app.load_models() except Exception as e: logger.warning(f"Could not pre-load models: {e}") # Create and launch interface demo = create_interface() demo.launch( server_name="0.0.0.0", server_port=7860, share=True, show_error=True, debug=True )