app.py

"""
eDOCr2 - Engineering Drawing OCR
Gradio Interface for Hugging Face Spaces

Extract dimensions, tables, and GD&T symbols from engineering drawings
"""

import gradio as gr
import cv2
import numpy as np
import json
import os
import time
from pathlib import Path
import zipfile
import tempfile
from PIL import Image

# Import eDOCr2 modules
from edocr2 import tools
from edocr2.keras_ocr.recognition import Recognizer
from edocr2.keras_ocr.detection import Detector
from pdf2image import convert_from_path

# Global variables for models
recognizer_gdt = None
recognizer_dim = None
detector = None
alphabet_dim = None
models_loaded = False

def load_models():
    """Load OCR models at startup"""
    global recognizer_gdt, recognizer_dim, detector, alphabet_dim, models_loaded
    
    if models_loaded:
        return True
    
    try:
        print("🔧 Loading OCR models...")
        start_time = time.time()
        
        # Model paths
        gdt_model = 'edocr2/models/recognizer_gdts.keras'
        dim_model = 'edocr2/models/recognizer_dimensions_2.keras'
        
        if not os.path.exists(gdt_model) or not os.path.exists(dim_model):
            print("❌ Model files not found!")
            return False
        
        # Load GD&T recognizer
        recognizer_gdt = Recognizer(alphabet=tools.ocr_pipelines.read_alphabet(gdt_model))
        recognizer_gdt.model.load_weights(gdt_model)
        
        # Load dimension recognizer
        alphabet_dim = tools.ocr_pipelines.read_alphabet(dim_model)
        recognizer_dim = Recognizer(alphabet=alphabet_dim)
        recognizer_dim.model.load_weights(dim_model)
        
        # Load detector
        detector = Detector()
        
        # Warm up models
        dummy_image = np.zeros((1, 1, 3), dtype=np.float32)
        _ = recognizer_gdt.recognize(dummy_image)
        _ = recognizer_dim.recognize(dummy_image)
        dummy_image = np.zeros((32, 32, 3), dtype=np.float32)
        _ = detector.detect([dummy_image])
        
        end_time = time.time()
        print(f"✅ Models loaded in {end_time - start_time:.2f} seconds")
        
        models_loaded = True
        return True
        
    except Exception as e:
        print(f"❌ Error loading models: {e}")
        import traceback
        traceback.print_exc()
        return False

def process_drawing(image_file):
    """
    Process an engineering drawing and extract information
    
    Args:
        image_file: Uploaded image file (PIL Image or file path)
    
    Returns:
        tuple: (annotated_image, json_data, csv_file, zip_file, stats_html)
    """
    
    if not models_loaded:
        return None, "❌ Models not loaded. Please check the logs.", None, None, "Error: Models not loaded"
    
    try:
        start_time = time.time()
        
        # Read image
        if isinstance(image_file, str):
            # File path
            if image_file.lower().endswith('.pdf'):
                img = convert_from_path(image_file)
                img = np.array(img[0])
                gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                _, img = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
                img = cv2.merge([img, img, img])
            else:
                img = cv2.imread(image_file)
        else:
            # PIL Image
            img = np.array(image_file)
            if len(img.shape) == 2:  # Grayscale
                img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
            elif img.shape[2] == 4:  # RGBA
                img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
            else:  # RGB
                img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        
        if img is None:
            return None, "❌ Could not read image file", None, None, "Error: Invalid image"
        
        # Create temporary output directory
        with tempfile.TemporaryDirectory() as temp_dir:
            output_dir = temp_dir
            
            # Segmentation
            print("🔍 Segmenting layers...")
            img_boxes, frame, gdt_boxes, tables, dim_boxes = tools.layer_segm.segment_img(
                img, autoframe=True, frame_thres=0.7, GDT_thres=0.02, binary_thres=127
            )
            
            # OCR Tables
            print("📋 Processing tables...")
            process_img = img.copy()
            table_results, updated_tables, process_img = tools.ocr_pipelines.ocr_tables(
                tables, process_img, language='eng'
            )
            
            # OCR GD&T
            print("🎯 Processing GD&T symbols...")
            gdt_results, updated_gdt_boxes, process_img = tools.ocr_pipelines.ocr_gdt(
                process_img, gdt_boxes, recognizer_gdt
            )
            
            # OCR Dimensions
            print("📏 Processing dimensions...")
            if frame:
                process_img = process_img[frame.y : frame.y + frame.h, frame.x : frame.x + frame.w]
            
            dimensions, other_info, process_img, dim_tess = tools.ocr_pipelines.ocr_dimensions(
                process_img, detector, recognizer_dim, alphabet_dim, frame, dim_boxes,
                cluster_thres=20, max_img_size=1048, language='eng', backg_save=False
            )
            
            # Generate mask image
            print("🎨 Generating visualization...")
            mask_img = tools.output_tools.mask_img(
                img, updated_gdt_boxes, updated_tables, dimensions, frame, other_info
            )
            
            # Convert to RGB for display
            mask_img_rgb = cv2.cvtColor(mask_img, cv2.COLOR_BGR2RGB)
            
            # Process and save results
            print("💾 Saving results...")
            table_results, gdt_results, dimensions, other_info = tools.output_tools.process_raw_output(
                output_dir, table_results, gdt_results, dimensions, other_info, save=True
            )
            
            # Prepare JSON data
            json_data = {
                'tables': table_results,
                'gdts': gdt_results,
                'dimensions': dimensions,
                'other_info': other_info
            }
            json_str = json.dumps(json_data, indent=2)
            
            # Save JSON file
            json_path = os.path.join(output_dir, 'results.json')
            with open(json_path, 'w') as f:
                f.write(json_str)
            
            # Create CSV file (if exists)
            csv_files = list(Path(output_dir).glob('*.csv'))
            csv_path = csv_files[0] if csv_files else None
            
            # Create ZIP file with all results
            zip_path = os.path.join(output_dir, 'edocr2_results.zip')
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                # Add mask image
                mask_img_path = os.path.join(output_dir, 'annotated_drawing.png')
                cv2.imwrite(mask_img_path, mask_img)
                zipf.write(mask_img_path, 'annotated_drawing.png')
                
                # Add JSON
                zipf.write(json_path, 'results.json')
                
                # Add CSV if exists
                if csv_path:
                    zipf.write(csv_path, os.path.basename(csv_path))
            
            # Copy ZIP to a permanent location
            permanent_zip = os.path.join(tempfile.gettempdir(), f'edocr2_results_{int(time.time())}.zip')
            import shutil
            shutil.copy(zip_path, permanent_zip)
            
            end_time = time.time()
            processing_time = round(end_time - start_time, 2)
            
            # Create statistics HTML
            stats_html = f"""
            <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
                        padding: 20px; border-radius: 10px; color: white; margin: 10px 0;">
                <h3 style="margin-top: 0;">📊 Processing Results</h3>
                <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(150px, 1fr)); gap: 15px; margin-top: 15px;">
                    <div style="background: rgba(255,255,255,0.2); padding: 15px; border-radius: 8px; text-align: center;">
                        <div style="font-size: 2em; font-weight: bold;">{len(table_results)}</div>
                        <div style="font-size: 0.9em; opacity: 0.9;">Tables Found</div>
                    </div>
                    <div style="background: rgba(255,255,255,0.2); padding: 15px; border-radius: 8px; text-align: center;">
                        <div style="font-size: 2em; font-weight: bold;">{len(gdt_results)}</div>
                        <div style="font-size: 0.9em; opacity: 0.9;">GD&T Symbols</div>
                    </div>
                    <div style="background: rgba(255,255,255,0.2); padding: 15px; border-radius: 8px; text-align: center;">
                        <div style="font-size: 2em; font-weight: bold;">{len(dimensions)}</div>
                        <div style="font-size: 0.9em; opacity: 0.9;">Dimensions</div>
                    </div>
                    <div style="background: rgba(255,255,255,0.2); padding: 15px; border-radius: 8px; text-align: center;">
                        <div style="font-size: 2em; font-weight: bold;">{processing_time}s</div>
                        <div style="font-size: 0.9em; opacity: 0.9;">Processing Time</div>
                    </div>
                </div>
            </div>
            """
            
            print(f"✅ Processing complete in {processing_time}s")
            
            return (
                Image.fromarray(mask_img_rgb),
                json_str,
                csv_path if csv_path else None,
                permanent_zip,
                stats_html
            )
            
    except Exception as e:
        error_msg = f"❌ Error processing drawing: {str(e)}"
        print(error_msg)
        import traceback
        traceback.print_exc()
        return None, error_msg, None, None, f"<div style='color: red;'>{error_msg}</div>"

# Load models at startup
print("="*60)
print("🔧 eDOCr2 - Engineering Drawing OCR")
print("="*60)
load_models()

# Create Gradio interface
with gr.Blocks(theme=gr.themes.Soft(), title="eDOCr2 - Engineering Drawing OCR") as demo:
    
    gr.Markdown("""
    # 🔧 eDOCr2 - Engineering Drawing OCR
    
    Extract **dimensions**, **tables**, and **GD&T symbols** from engineering drawings automatically.
    
    Upload your engineering drawing (JPG, PNG, or PDF) and get structured data extracted instantly!
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Upload Drawing")
            image_input = gr.Image(
                type="pil",
                label="Engineering Drawing",
                sources=["upload", "clipboard"],
                height=400
            )
            
            process_btn = gr.Button(
                "🚀 Process Drawing",
                variant="primary",
                size="lg"
            )
            
            gr.Markdown("""
            **Supported formats:** JPG, PNG, PDF  
            **Best results:** High-resolution scans with clear text
            """)
        
        with gr.Column(scale=1):
            gr.Markdown("### 📊 Statistics")
            stats_output = gr.HTML()
            
            gr.Markdown("### 🎨 Annotated Drawing")
            image_output = gr.Image(
                label="Processed Drawing",
                type="pil",
                height=400
            )
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📋 Extracted Data (JSON)")
            json_output = gr.Textbox(
                label="Structured Data",
                lines=15,
                max_lines=20,
                show_copy_button=True
            )
        
        with gr.Column():
            gr.Markdown("### 💾 Download Results")
            
            csv_output = gr.File(
                label="📊 CSV File (if available)",
                type="filepath"
            )
            
            zip_output = gr.File(
                label="📦 Complete Results (ZIP)",
                type="filepath"
            )
            
            gr.Markdown("""
            **ZIP contains:**
            - Annotated drawing image
            - Structured data (JSON)
            - Tabular data (CSV, if applicable)
            """)
    
    # Examples section
    gr.Markdown("### 🧪 Try Example Drawings")
    
    example_files = []
    examples_dir = "tests/test_samples"
    if os.path.exists(examples_dir):
        for file in os.listdir(examples_dir):
            if file.endswith(('.jpg', '.png')):
                example_files.append([os.path.join(examples_dir, file)])
    
    if example_files:
        gr.Examples(
            examples=example_files[:3],  # Show first 3 examples
            inputs=image_input,
            label="Click to load example"
        )
    
    # Footer
    gr.Markdown("""
    ---
    
    ### 📚 About eDOCr2
    
    eDOCr2 is a specialized OCR tool for engineering drawings that can extract:
    - **Tables**: Title blocks, revision tables, bill of materials
    - **GD&T Symbols**: Geometric dimensioning and tolerancing
    - **Dimensions**: Measurements with tolerances
    - **Other Information**: Additional text and annotations
    
    **Research Paper:** [http://dx.doi.org/10.2139/ssrn.5045921](http://dx.doi.org/10.2139/ssrn.5045921)  
    **GitHub:** [github.com/javvi51/edocr2](https://github.com/javvi51/edocr2)  
    **Created by:** Javier Villena Toro | **Deployed by:** Jeyanthan GJ
    """)
    
    # Connect the process button
    process_btn.click(
        fn=process_drawing,
        inputs=image_input,
        outputs=[image_output, json_output, csv_output, zip_output, stats_output]
    )

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