app.py

import gradio as gr
from docling.document_converter import DocumentConverter
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions, TableFormerMode
from docling.document_converter import PdfFormatOption
from PIL import Image, ImageDraw, ImageFont
import json
import fitz  # PyMuPDF

# Color mapping for different layout elements
COLORS = {
    "title": "#FF6B6B",
    "text": "#4ECDC4",
    "section_header": "#95E1D3",
    "table": "#F38181",
    "list": "#AA96DA",
    "figure": "#FCBAD3",
    "caption": "#A8D8EA",
    "formula": "#FFD93D",
    "footnote": "#6BCB77",
    "page_header": "#4D96FF",
    "page_footer": "#9D84B7",
    "picture": "#FF8C42",
    # Picture classifications
    "signature": "#9D4EDD",
    "qr_code": "#06FFA5",
    "bar_code": "#06FFA5",
    "logo": "#FFB627",
    "stamp": "#E63946",
    "icon": "#F4A261",
    "bar_chart": "#2A9D8F",
    "pie_chart": "#E76F51",
    "line_chart": "#264653",
    "flow_chart": "#8338EC",
    "map": "#3A86FF",
    "screenshot": "#FB5607",
    "other": "#CCCCCC",
}

def draw_layout_boxes(image_path, layout_data, scale_x=1.0, scale_y=1.0):
    """Draw bounding boxes on the image based on layout predictions"""
    # Open the image
    if isinstance(image_path, str):
        img = Image.open(image_path).convert("RGB")
    else:
        img = image_path.convert("RGB")
    
    draw = ImageDraw.Draw(img)
    
    # Try to load a font, fallback to default if not available
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
        small_font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 14)
    except:
        font = ImageFont.load_default()
        small_font = ImageFont.load_default()
    
    # Draw each cluster
    for cluster in layout_data:
        label = cluster.get("label", "unknown")
        bbox = cluster.get("bbox")
        classification = cluster.get("classification")
        
        if bbox:
            # bbox format: [x0, y0, x1, y1] from PDF coordinates
            # Scale to match rendered image dimensions
            x0, y0, x1, y1 = bbox
            x0 = x0 * scale_x
            y0 = y0 * scale_y
            x1 = x1 * scale_x
            y1 = y1 * scale_y
            
            # Get color for this label
            color = COLORS.get(label, "#999999")
            
            # Draw rectangle
            draw.rectangle([x0, y0, x1, y1], outline=color, width=3)
            
            # Draw label with classification confidence if available
            if classification:
                confidence_pct = classification['confidence'] * 100
                label_text = f"{label.replace('_', ' ').title()} ({confidence_pct:.0f}%)"
            else:
                label_text = label.replace("_", " ").title()
            
            bbox_text = draw.textbbox((x0, y0 - 25), label_text, font=small_font)
            draw.rectangle([bbox_text[0] - 2, bbox_text[1] - 2, bbox_text[2] + 2, bbox_text[3] + 2], 
                         fill=color)
            
            # Draw label text
            draw.text((x0, y0 - 25), label_text, fill="white", font=small_font)
    
    return img

def process_document(file_path, mode, enable_ocr, enable_tables):
    """Process document with Docling and return results"""
    try:
        # Configure pipeline options
        pipeline_options = PdfPipelineOptions()
        pipeline_options.do_table_structure = enable_tables
        
        if enable_tables:
            if mode == "Accurate":
                pipeline_options.table_structure_options.mode = TableFormerMode.ACCURATE
            else:
                pipeline_options.table_structure_options.mode = TableFormerMode.FAST
        
        pipeline_options.do_ocr = enable_ocr
        pipeline_options.generate_page_images = True
        pipeline_options.generate_picture_images = True
        pipeline_options.do_picture_classification = True  # Enable classification
        pipeline_options.images_scale = 3.0  # Higher resolution for better accuracy
        
        # Create converter
        converter = DocumentConverter(
            format_options={
                InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options),
                InputFormat.IMAGE: PdfFormatOption(pipeline_options=pipeline_options),
            }
        )
        
        # Convert document
        result = converter.convert(file_path)
        
        # Extract layout information
        layout_info = []
        total_clusters = 0
        table_count = 0
        
        # Get picture classifications for enrichment
        # We need to store by page number and use a more flexible matching
        picture_classifications_by_page = {}
        print(f"DEBUG: Total pictures found: {len(result.document.pictures)}")
        for picture in result.document.pictures:
            page_num = picture.prov[0].page_no
            bbox = picture.prov[0].bbox
            
            if page_num not in picture_classifications_by_page:
                picture_classifications_by_page[page_num] = []
            
            # Get classification if available
            for annotation in picture.annotations:
                if hasattr(annotation, 'predicted_classes') and annotation.predicted_classes:
                    top_pred = annotation.predicted_classes[0]
                    picture_classifications_by_page[page_num].append({
                        'bbox': bbox,
                        'class': top_pred.class_name,
                        'confidence': top_pred.confidence
                    })
                    print(f"DEBUG: Found classification - page: {page_num}, bbox: ({bbox.l:.2f}, {bbox.t:.2f}, {bbox.r:.2f}, {bbox.b:.2f}), class: {top_pred.class_name}")
                    break
        
        for page_no, page in enumerate(result.pages, 1):
            if page.predictions.layout:
                clusters = page.predictions.layout.clusters
                total_clusters += len(clusters)
                
                for cluster in clusters:
                    # Check if this is a picture with classification
                    label = cluster.label
                    classification = None
                    if cluster.label == "picture" and page_no in picture_classifications_by_page:
                        print(f"DEBUG: Picture cluster at page {page_no}: ({cluster.bbox.l:.2f}, {cluster.bbox.t:.2f}, {cluster.bbox.r:.2f}, {cluster.bbox.b:.2f})")
                        
                        # Find matching classification by comparing bounding boxes with tolerance
                        for pic_class in picture_classifications_by_page[page_no]:
                            pic_bbox = pic_class['bbox']
                            # Check if bboxes match with small tolerance (allowing for floating point differences)
                            # Compare left and right which should match exactly
                            if (abs(cluster.bbox.l - pic_bbox.l) < 1.0 and 
                                abs(cluster.bbox.r - pic_bbox.r) < 1.0):
                                # X coordinates match, this is likely the same picture
                                classification = {
                                    'class': pic_class['class'],
                                    'confidence': pic_class['confidence']
                                }
                                label = f"{classification['class']}"
                                print(f"DEBUG: Matched classification: {label} (conf: {classification['confidence']:.2%})")
                                break
                        
                        if not classification:
                            print(f"DEBUG: No classification match found")
                    
                    layout_info.append({
                        "page": page_no,
                        "label": label,
                        "bbox": [cluster.bbox.l, cluster.bbox.t, cluster.bbox.r, cluster.bbox.b],
                        "confidence": getattr(cluster, "confidence", None),
                        "classification": classification
                    })
            
            # Count tables
            if page.predictions.tablestructure and page.predictions.tablestructure.table_map:
                table_count += len(page.predictions.tablestructure.table_map)
        
        # Get markdown output
        markdown_output = result.document.export_to_markdown()
        
        # Create visualization for first page
        visualization = None
        if result.pages and layout_info:
            # Draw boxes on first page only
            first_page_layout = [item for item in layout_info if item["page"] == 1]
            
            try:
                # Check if input is an image or PDF
                file_ext = file_path.lower().split('.')[-1]
                
                if file_ext in ['jpg', 'jpeg', 'png', 'tiff', 'bmp']:
                    # For images: Open directly, coordinates should match 1:1
                    first_page_image = Image.open(file_path).convert("RGB")
                    # No scaling needed for images - coordinates are already in pixels
                    visualization = draw_layout_boxes(first_page_image, first_page_layout, 
                                                     scale_x=1.0, scale_y=1.0)
                else:
                    # For PDFs: Render and calculate scale
                    doc = fitz.open(file_path)
                    page = doc[0]
                    
                    # Get page dimensions in PDF points
                    page_rect = page.rect
                    pdf_width = page_rect.width
                    pdf_height = page_rect.height
                    
                    # Render at 2x for better quality
                    zoom = 2.0
                    mat = fitz.Matrix(zoom, zoom)
                    pix = page.get_pixmap(matrix=mat)
                    first_page_image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                    
                    # Calculate scale: rendered_pixels / pdf_points
                    scale_x = pix.width / pdf_width
                    scale_y = pix.height / pdf_height
                    
                    doc.close()
                    
                    # Draw boxes with calculated scale
                    visualization = draw_layout_boxes(first_page_image, first_page_layout, 
                                                     scale_x=scale_x, scale_y=scale_y)
            except Exception as e:
                print(f"Could not create visualization: {e}")
                import traceback
                traceback.print_exc()
        
        # Create summary
        summary = f"""## Document Analysis Summary

📄 **Total Pages:** {len(result.document.pages)}
🏷️ **Layout Elements Detected:** {total_clusters}
📊 **Tables Found:** {table_count}

### Layout Elements by Type:
"""
        # Count elements by type
        element_counts = {}
        for item in layout_info:
            label = item["label"]
            element_counts[label] = element_counts.get(label, 0) + 1
        
        for label, count in sorted(element_counts.items()):
            summary += f"- **{label.replace('_', ' ').title()}**: {count}\n"
        
        # JSON output
        json_output = json.dumps(layout_info, indent=2)
        
        return visualization, summary, markdown_output, json_output
        
    except Exception as e:
        error_msg = f"Error processing document: {str(e)}"
        return None, error_msg, error_msg, error_msg

def gradio_interface(file, mode, enable_ocr, enable_tables):
    """Gradio interface function"""
    if file is None:
        return None, "Please upload a document", "", ""
    
    return process_document(file.name, mode, enable_ocr, enable_tables)

# Create Gradio interface
with gr.Blocks(title="Document Layout Detection", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 📄 Document Layout & Structure Detection
    
    Upload a document (PDF, image, etc.) to automatically detect its layout structure including text, tables, figures, and more!
    
    **Features:**
    - **AI-Powered Layout Detection**: Automatically identifies document elements
    - **Table Structure Extraction**: Recognizes and extracts table data
    - **OCR Support**: Reads text from scanned documents and images
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(
                label="Upload Document", 
                file_types=[".pdf", ".jpg", ".jpeg", ".png", ".tiff", ".bmp"]
            )
            
            mode_dropdown = gr.Dropdown(
                choices=["Fast", "Accurate"],
                value="Fast",
                label="Processing Mode",
                info="Accurate mode is slower but better for complex tables"
            )
            
            ocr_checkbox = gr.Checkbox(
                label="Enable OCR",
                value=True,
                info="Use OCR for scanned documents and images"
            )
            
            tables_checkbox = gr.Checkbox(
                label="Enable Table Detection",
                value=True,
                info="Detect and extract table structures"
            )
            
            process_btn = gr.Button("🚀 Process Document", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            visualization_output = gr.Image(label="Layout Visualization (First Page)")
            summary_output = gr.Markdown(label="Summary")
    
    with gr.Tabs():
        with gr.Tab("📝 Markdown Output"):
            markdown_output = gr.Textbox(
                label="Extracted Content (Markdown)",
                lines=20,
                max_lines=30
            )
        
        with gr.Tab("🔧 JSON Layout Data"):
            json_output = gr.Code(
                label="Layout Predictions (JSON)",
                language="json",
                lines=20
            )
    
    gr.Markdown("""
    ### Legend
    Different colors represent different document elements:
    
    **Layout Elements:**
    - 🔴 Title • 🔵 Text • 🟢 Section Header • 🟠 Table • 🟣 List/Figure/Formula
    
    **Picture Classifications (AI-detected):**
    - 🟣 Signature • 🟢 QR Code • 🟢 Barcode • 🟡 Logo • 🔴 Stamp
    - 🟦 Charts (Bar/Pie/Line) • 🟣 Flow Chart • 🟠 Screenshot • ⚪ Other
    
    ### How to Use
    1. Upload your document (PDF or image of ID card, invoice, report, etc.)
    2. Choose processing options (Fast mode recommended for quick results)
    3. Click "Process Document"
    4. View the visualization with bounding boxes and explore the outputs
    
    ### 💡 Try Examples Below!
    Click on any example document to see instant results on different document types.
    """)
    
    # Add examples with image previews
    with gr.Row():
        gr.Examples(
            examples=[
                ["sample/Screenshot 2025-10-13 114010.png", "Fast", True, True],
                ["sample/Screenshot 2025-10-13 114606.png", "Fast", True, True],
                ["sample/Screenshot 2025-10-15 191615.png", "Fast", True, True],
            ],
            inputs=[file_input, mode_dropdown, ocr_checkbox, tables_checkbox],
            outputs=[visualization_output, summary_output, markdown_output, json_output],
            fn=gradio_interface,
            cache_examples=False,
            label="📚 Example Documents",
            examples_per_page=3
        )
    
    # Connect the button
    process_btn.click(
        fn=gradio_interface,
        inputs=[file_input, mode_dropdown, ocr_checkbox, tables_checkbox],
        outputs=[visualization_output, summary_output, markdown_output, json_output]
    )
    
    # Auto-process on file upload (optional)
    file_input.change(
        fn=gradio_interface,
        inputs=[file_input, mode_dropdown, ocr_checkbox, tables_checkbox],
        outputs=[visualization_output, summary_output, markdown_output, json_output]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()