Update app.py

app.py

@@ -1,93 +1,477 @@
-import os
 import gradio as gr
-import tempfile
-import shutil
 import pandas as pd
+import numpy as np
 from PIL import Image
-from preprocess import convert_pdf_to_images, preprocess_image
-from llm_utils import load_image, generate_response
-
-# Global temporary directory for image processing
-temp_dir = tempfile.mkdtemp()
-
-pdf_image_map = {}  # Map from PDF to its image list
-image_preview_map = {}  # Map from image name to full path
-
-
-def extract_images_from_pdfs(pdf_files):
-    global pdf_image_map, image_preview_map
-    pdf_image_map.clear()
-    image_preview_map.clear()
-
-    previews = []
-
-    for pdf in pdf_files:
-        image_paths = convert_pdf_to_images(pdf.name, temp_dir)
-        pdf_image_map[pdf.name] = image_paths
-
-        for img_path in image_paths:
-            img_name = os.path.basename(img_path)
-            image_preview_map[img_name] = img_path
-            previews.append((img_name, img_path))
-
-    # Return preview (tuples: filename, image_path)
-    return [img for _, img in previews]
-
-
-def process_selected_images(selected_images, output_excel_name):
-    results = []
-
-    for img_name in selected_images:
-        img_path = image_preview_map.get(img_name)
-        if img_path is None:
-            continue
-
-        processed = preprocess_image(img_path)
-        processed_path = os.path.join(temp_dir, f"processed_{img_name}")
-        Image.fromarray(processed).save(processed_path)
-
-        # Run LLM
-        pixel_values = load_image(processed_path)
-        response = generate_response(pixel_values)
-
-        # Clean filename to find source pdf
-        base_name = img_name.rsplit("_page_", 1)[0] + ".pdf"
-        results.append({"Source PDF": base_name, "Page Image": img_name, "LLM Output": response})
-
-    df = pd.DataFrame(results)
-    output_excel = os.path.join(temp_dir, output_excel_name)
-    df.to_excel(output_excel, index=False)
-    return output_excel
-
-
-def reset_all():
-    shutil.rmtree(temp_dir)
-    os.makedirs(temp_dir, exist_ok=True)
-
-
-with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 PDF → Image → LLM → Excel Output")
-
-    with gr.Row():
-        pdf_input = gr.File(file_types=[".pdf"], file_count="multiple", label="📎 Upload multiple PDF files")
-        extract_btn = gr.Button("🔍 Extract Images")
+import io
+import tempfile
+import os
+import zipfile
+from pathlib import Path
+import json
+from typing import List, Dict, Any, Tuple, Optional
+import logging
 
-    gallery = gr.Gallery(label="🖼️ Choose images to process", columns=4, allow_preview=True, interactive=True, show_label=True).style(grid=4)
-    selected_images = gr.CheckboxGroup(choices=[], label="Select image filenames for processing")
+# Import your custom modules
+from preprocess import PDFImageProcessor, process_uploaded_pdf
+from llm_utils import load_model, extract_info_from_image, cleanup_model
 
-    with gr.Row():
-        output_name = gr.Textbox(label="📄 Output Excel filename", value="output.xlsx")
-        generate_btn = gr.Button("🚀 Generate Excel")
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-    excel_output = gr.File(label="📥 Download Excel")
+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()
 
-    def update_gallery(pdf_files):
-        previews = extract_images_from_pdfs(pdf_files)
-        choices = list(image_preview_map.keys())
-        return gr.update(value=previews), gr.update(choices=choices, value=[])
+# Initialize the app
+app = PDFProcessingApp()
 
-    extract_btn.click(update_gallery, inputs=[pdf_input], outputs=[gallery, selected_images])
-    generate_btn.click(fn=process_selected_images, inputs=[selected_images, output_name], outputs=[excel_output])
+# 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__":
-    demo.launch()
+    # 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
+    )