app.py

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
    )