"""
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"""
๐ Processing Results
{len(table_results)}
Tables Found
{len(gdt_results)}
GD&T Symbols
{len(dimensions)}
Dimensions
{processing_time}s
Processing Time
"""
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"{error_msg}
"
# 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()