INTRODUCTION
The transformation of hand-drawn graphics into clean digital representations represents a significant challenge in computer vision and artificial intelligence. This article presents a comprehensive system that accepts photographs of hand-drawn diagrams, analyzes their semantic content, recognizes diagram types, and produces optimized digital equivalents while preserving the original intent and meaning.
The system addresses several complex problems simultaneously. First, it must handle the inherent variability and imperfection of hand-drawn input, including inconsistent line weights, irregular shapes, and varying handwriting styles. Second, it must recognize high-level diagram semantics, distinguishing between flowcharts, UML diagrams, circuit schematics, and other specialized notations. Third, it must beautify and optimize the output while maintaining semantic equivalence to the original drawing. Finally, it must support diverse hardware configurations and both local and remote language model backends.
This implementation leverages open-source tools exclusively and provides production-ready code that supports Intel GPUs, AMD ROCm, Apple Metal Performance Shaders, and NVIDIA CUDA architectures. The system architecture follows clean code principles with clear separation of concerns and extensible design patterns.
SYSTEM ARCHITECTURE OVERVIEW
The digitization system consists of five primary components working in a pipeline architecture. The first component handles image preprocessing and normalization. The second performs optical character recognition and text extraction. The third detects and classifies geometric shapes and symbols. The fourth component recognizes diagram types and semantic structures. The fifth generates the final digital representation.
Each component operates independently with well-defined interfaces, allowing for parallel processing where appropriate and easy replacement of individual modules. The system uses a factory pattern for instantiating backend-specific implementations and a strategy pattern for selecting appropriate processing algorithms based on detected diagram types.
HARDWARE ABSTRACTION AND MULTI-GPU SUPPORT
Supporting multiple GPU architectures requires careful abstraction of device-specific operations. The system must detect available hardware at runtime and configure PyTorch accordingly. This involves checking for CUDA availability, ROCm support, Metal Performance Shaders on Apple Silicon, and Intel extension for PyTorch.
The device manager component handles all hardware detection and configuration. It probes the system for available accelerators and selects the most appropriate device based on a priority hierarchy. NVIDIA CUDA receives highest priority when available, followed by AMD ROCm, Apple MPS, and Intel GPU support. The system falls back to CPU execution when no GPU acceleration is available.
Here is the device manager implementation:
import torch
import platform
import subprocess
import logging
from typing import Optional, Dict, Any
from enum import Enum
class DeviceType(Enum):
"""Enumeration of supported device types."""
CUDA = "cuda"
ROCM = "rocm"
MPS = "mps"
INTEL = "xpu"
CPU = "cpu"
class DeviceManager:
"""
Manages hardware detection and device configuration for multi-GPU support.
Supports NVIDIA CUDA, AMD ROCm, Apple MPS, Intel XPU, and CPU fallback.
"""
def __init__(self):
"""Initialize device manager and detect available hardware."""
self.logger = logging.getLogger(__name__)
self.device_type = None
self.device = None
self.device_properties = {}
self._detect_and_configure()
def _check_cuda_availability(self) -> bool:
"""Check if NVIDIA CUDA is available and functional."""
try:
if torch.cuda.is_available():
# Verify CUDA actually works by attempting a simple operation
test_tensor = torch.zeros(1).cuda()
del test_tensor
torch.cuda.empty_cache()
return True
except Exception as e:
self.logger.warning(f"CUDA detected but not functional: {e}")
return False
def _check_rocm_availability(self) -> bool:
"""Check if AMD ROCm is available."""
try:
# ROCm uses the same CUDA API in PyTorch
if torch.cuda.is_available():
# Check if this is actually ROCm
device_name = torch.cuda.get_device_name(0).lower()
if 'amd' in device_name or 'radeon' in device_name:
return True
except Exception as e:
self.logger.warning(f"ROCm check failed: {e}")
return False
def _check_mps_availability(self) -> bool:
"""Check if Apple Metal Performance Shaders is available."""
try:
if platform.system() == 'Darwin':
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
# Verify MPS actually works
test_tensor = torch.zeros(1).to('mps')
del test_tensor
return True
except Exception as e:
self.logger.warning(f"MPS detected but not functional: {e}")
return False
def _check_intel_availability(self) -> bool:
"""Check if Intel GPU extension is available."""
try:
import intel_extension_for_pytorch as ipex
if hasattr(torch, 'xpu') and torch.xpu.is_available():
test_tensor = torch.zeros(1).to('xpu')
del test_tensor
return True
except ImportError:
self.logger.info("Intel Extension for PyTorch not installed")
except Exception as e:
self.logger.warning(f"Intel XPU check failed: {e}")
return False
def _detect_and_configure(self):
"""Detect available hardware and configure the appropriate device."""
self.logger.info("Detecting available hardware accelerators...")
# Check in priority order: CUDA > ROCm > MPS > Intel > CPU
if self._check_cuda_availability():
if self._check_rocm_availability():
self.device_type = DeviceType.ROCM
self.device = torch.device('cuda')
self.device_properties = {
'name': torch.cuda.get_device_name(0),
'compute_capability': torch.cuda.get_device_capability(0),
'total_memory': torch.cuda.get_device_properties(0).total_memory,
'backend': 'ROCm'
}
self.logger.info(f"Using AMD ROCm: {self.device_properties['name']}")
else:
self.device_type = DeviceType.CUDA
self.device = torch.device('cuda')
self.device_properties = {
'name': torch.cuda.get_device_name(0),
'compute_capability': torch.cuda.get_device_capability(0),
'total_memory': torch.cuda.get_device_properties(0).total_memory,
'backend': 'CUDA'
}
self.logger.info(f"Using NVIDIA CUDA: {self.device_properties['name']}")
elif self._check_mps_availability():
self.device_type = DeviceType.MPS
self.device = torch.device('mps')
self.device_properties = {
'name': 'Apple Silicon GPU',
'backend': 'MPS'
}
self.logger.info("Using Apple Metal Performance Shaders")
elif self._check_intel_availability():
self.device_type = DeviceType.INTEL
self.device = torch.device('xpu')
self.device_properties = {
'name': 'Intel GPU',
'backend': 'Intel Extension for PyTorch'
}
self.logger.info("Using Intel GPU acceleration")
else:
self.device_type = DeviceType.CPU
self.device = torch.device('cpu')
self.device_properties = {
'name': 'CPU',
'backend': 'CPU'
}
self.logger.warning("No GPU acceleration available, using CPU")
def get_device(self) -> torch.device:
"""Return the configured PyTorch device."""
return self.device
def get_device_type(self) -> DeviceType:
"""Return the device type enumeration."""
return self.device_type
def get_device_info(self) -> Dict[str, Any]:
"""Return detailed device properties."""
return self.device_properties.copy()
def optimize_for_inference(self):
"""Apply device-specific optimizations for inference."""
if self.device_type == DeviceType.CUDA or self.device_type == DeviceType.ROCM:
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
elif self.device_type == DeviceType.MPS:
# MPS-specific optimizations
pass
elif self.device_type == DeviceType.INTEL:
try:
import intel_extension_for_pytorch as ipex
ipex.optimize(optimizer=None)
except Exception as e:
self.logger.warning(f"Intel optimization failed: {e}")
The device manager encapsulates all hardware-specific logic in a single component. It performs actual device testing rather than relying solely on availability flags, ensuring that selected devices are truly functional. The optimization method applies backend-specific performance tuning, such as enabling cuDNN benchmarking for CUDA devices.
LANGUAGE MODEL BACKEND ABSTRACTION
Supporting both local and remote language models requires a flexible backend architecture. The system must accommodate different API interfaces, authentication mechanisms, and response formats while presenting a unified interface to higher-level components.
The language model backend uses an abstract base class defining the interface that all implementations must satisfy. Concrete implementations handle specific backends such as local Llama models, OpenAI GPT-4 Vision, and other vision-language models. The factory pattern instantiates appropriate backends based on configuration.
from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional, Union
import base64
import io
from PIL import Image
import torch
from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer
import requests
import json
class VisionLanguageBackend(ABC):
"""Abstract base class for vision-language model backends."""
@abstractmethod
def analyze_image(self, image: Image.Image, prompt: str,
max_tokens: int = 1000) -> str:
"""
Analyze an image with a text prompt and return the model's response.
Args:
image: PIL Image object to analyze
prompt: Text prompt describing the analysis task
max_tokens: Maximum tokens in the response
Returns:
String response from the model
"""
pass
@abstractmethod
def get_backend_info(self) -> Dict[str, Any]:
"""Return information about the backend configuration."""
pass
class LocalLlamaVisionBackend(VisionLanguageBackend):
"""
Backend for local Llama vision models (e.g., LLaVA, Llama 3.2 Vision).
Supports multi-GPU architectures through the device manager.
"""
def __init__(self, model_name: str, device_manager: DeviceManager):
"""
Initialize local Llama vision backend.
Args:
model_name: HuggingFace model identifier
device_manager: Configured device manager instance
"""
self.model_name = model_name
self.device_manager = device_manager
self.device = device_manager.get_device()
self.logger = logging.getLogger(__name__)
self.logger.info(f"Loading local model: {model_name}")
# Load model and processor
self.processor = AutoProcessor.from_pretrained(model_name)
# Configure model loading based on device type
load_kwargs = {'torch_dtype': torch.float16}
if device_manager.get_device_type() == DeviceType.CUDA:
load_kwargs['device_map'] = 'auto'
elif device_manager.get_device_type() == DeviceType.MPS:
# MPS doesn't support float16 for all operations
load_kwargs['torch_dtype'] = torch.float32
self.model = AutoModelForCausalLM.from_pretrained(
model_name,
**load_kwargs
)
if device_manager.get_device_type() != DeviceType.CUDA:
self.model = self.model.to(self.device)
self.model.eval()
self.logger.info(f"Model loaded successfully on {self.device}")
def analyze_image(self, image: Image.Image, prompt: str,
max_tokens: int = 1000) -> str:
"""Analyze image using local vision-language model."""
try:
# Prepare inputs
messages = [
{
"role": "user",
"content": [
{"type": "image"},
{"type": "text", "text": prompt}
]
}
]
# Process inputs
inputs = self.processor(
text=self.processor.apply_chat_template(messages, add_generation_prompt=True),
images=image,
return_tensors="pt"
).to(self.device)
# Generate response
with torch.no_grad():
output = self.model.generate(
**inputs,
max_new_tokens=max_tokens,
do_sample=False
)
# Decode response
response = self.processor.decode(output[0], skip_special_tokens=True)
# Extract assistant response
if "assistant" in response:
response = response.split("assistant")[-1].strip()
return response
except Exception as e:
self.logger.error(f"Error during image analysis: {e}")
raise
def get_backend_info(self) -> Dict[str, Any]:
"""Return backend configuration information."""
return {
'backend_type': 'local_llama_vision',
'model_name': self.model_name,
'device': str(self.device),
'device_type': self.device_manager.get_device_type().value
}
class OpenAIVisionBackend(VisionLanguageBackend):
"""Backend for OpenAI GPT-4 Vision API."""
def __init__(self, api_key: str, model: str = "gpt-4-vision-preview"):
"""
Initialize OpenAI vision backend.
Args:
api_key: OpenAI API key
model: Model identifier (default: gpt-4-vision-preview)
"""
self.api_key = api_key
self.model = model
self.api_url = "https://api.openai.com/v1/chat/completions"
self.logger = logging.getLogger(__name__)
def _encode_image(self, image: Image.Image) -> str:
"""Encode PIL Image to base64 string."""
buffered = io.BytesIO()
image.save(buffered, format="PNG")
return base64.b64encode(buffered.getvalue()).decode('utf-8')
def analyze_image(self, image: Image.Image, prompt: str,
max_tokens: int = 1000) -> str:
"""Analyze image using OpenAI GPT-4 Vision API."""
try:
base64_image = self._encode_image(image)
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}"
}
payload = {
"model": self.model,
"messages": [
{
"role": "user",
"content": [
{
"type": "text",
"text": prompt
},
{
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{base64_image}"
}
}
]
}
],
"max_tokens": max_tokens
}
response = requests.post(self.api_url, headers=headers, json=payload)
response.raise_for_status()
result = response.json()
return result['choices'][0]['message']['content']
except Exception as e:
self.logger.error(f"Error calling OpenAI API: {e}")
raise
def get_backend_info(self) -> Dict[str, Any]:
"""Return backend configuration information."""
return {
'backend_type': 'openai_vision',
'model': self.model,
'api_url': self.api_url
}
class VisionBackendFactory:
"""Factory for creating vision-language model backends."""
@staticmethod
def create_backend(backend_type: str, device_manager: Optional[DeviceManager] = None,
**kwargs) -> VisionLanguageBackend:
"""
Create a vision-language backend based on type.
Args:
backend_type: Type of backend ('local_llama', 'openai', etc.)
device_manager: Device manager for local models
**kwargs: Additional backend-specific arguments
Returns:
Configured VisionLanguageBackend instance
"""
if backend_type == 'local_llama':
if device_manager is None:
device_manager = DeviceManager()
model_name = kwargs.get('model_name', 'llava-hf/llava-1.5-7b-hf')
return LocalLlamaVisionBackend(model_name, device_manager)
elif backend_type == 'openai':
api_key = kwargs.get('api_key')
if not api_key:
raise ValueError("OpenAI backend requires 'api_key' parameter")
model = kwargs.get('model', 'gpt-4-vision-preview')
return OpenAIVisionBackend(api_key, model)
else:
raise ValueError(f"Unknown backend type: {backend_type}")
This backend architecture provides complete flexibility in model selection. The abstract base class ensures consistent interfaces across implementations. Local models leverage the device manager for hardware acceleration, while remote backends handle API communication and authentication. The factory pattern simplifies backend instantiation and configuration management.
IMAGE PREPROCESSING AND NORMALIZATION
Raw photographs of hand-drawn diagrams require substantial preprocessing before analysis. Images may suffer from perspective distortion, uneven lighting, shadows, background noise, and varying resolutions. The preprocessing pipeline must correct these issues while preserving the essential features of the drawing.
The preprocessing component performs several operations in sequence. First, it converts the image to grayscale and applies adaptive thresholding to separate foreground content from background. Second, it detects and corrects perspective distortion by identifying the document boundaries. Third, it applies noise reduction through morphological operations. Fourth, it normalizes the image resolution and aspect ratio. Finally, it enhances contrast to improve feature detection in subsequent stages.
import cv2
import numpy as np
from typing import Tuple, Optional
from PIL import Image
class ImagePreprocessor:
"""
Handles preprocessing of hand-drawn diagram photographs.
Corrects perspective, removes noise, and normalizes images.
"""
def __init__(self, target_size: Tuple[int, int] = (1024, 1024)):
"""
Initialize image preprocessor.
Args:
target_size: Target dimensions for normalized output (width, height)
"""
self.target_size = target_size
self.logger = logging.getLogger(__name__)
def preprocess(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:
"""
Complete preprocessing pipeline for hand-drawn diagram images.
Args:
image: Input image as PIL Image or numpy array
Returns:
Preprocessed PIL Image
"""
# Convert to numpy array if needed
if isinstance(image, Image.Image):
img_array = np.array(image)
else:
img_array = image
# Convert to grayscale if needed
if len(img_array.shape) == 3:
gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
else:
gray = img_array
# Detect and correct perspective
corrected = self._correct_perspective(gray)
# Apply adaptive thresholding
binary = self._adaptive_threshold(corrected)
# Remove noise
denoised = self._remove_noise(binary)
# Normalize size
normalized = self._normalize_size(denoised)
# Enhance contrast
enhanced = self._enhance_contrast(normalized)
# Convert back to PIL Image
return Image.fromarray(enhanced)
def _correct_perspective(self, image: np.ndarray) -> np.ndarray:
"""
Detect document boundaries and correct perspective distortion.
Args:
image: Grayscale input image
Returns:
Perspective-corrected image
"""
# Apply edge detection
edges = cv2.Canny(image, 50, 150, apertureSize=3)
# Dilate edges to close gaps
kernel = np.ones((5, 5), np.uint8)
dilated = cv2.dilate(edges, kernel, iterations=1)
# Find contours
contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
if not contours:
self.logger.warning("No contours found for perspective correction")
return image
# Find largest contour (assumed to be document boundary)
largest_contour = max(contours, key=cv2.contourArea)
# Approximate contour to polygon
epsilon = 0.02 * cv2.arcLength(largest_contour, True)
approx = cv2.approxPolyDP(largest_contour, epsilon, True)
# If we found a quadrilateral, apply perspective transform
if len(approx) == 4:
return self._apply_perspective_transform(image, approx)
else:
self.logger.info("Document boundary not quadrilateral, skipping perspective correction")
return image
def _apply_perspective_transform(self, image: np.ndarray,
corners: np.ndarray) -> np.ndarray:
"""
Apply perspective transformation to correct distortion.
Args:
image: Input image
corners: Four corner points of the document
Returns:
Transformed image
"""
# Reshape corners
corners = corners.reshape(4, 2)
# Order corners: top-left, top-right, bottom-right, bottom-left
rect = self._order_points(corners)
# Calculate dimensions of corrected image
width_a = np.linalg.norm(rect[2] - rect[3])
width_b = np.linalg.norm(rect[1] - rect[0])
max_width = max(int(width_a), int(width_b))
height_a = np.linalg.norm(rect[1] - rect[2])
height_b = np.linalg.norm(rect[0] - rect[3])
max_height = max(int(height_a), int(height_b))
# Define destination points
dst = np.array([
[0, 0],
[max_width - 1, 0],
[max_width - 1, max_height - 1],
[0, max_height - 1]
], dtype=np.float32)
# Calculate perspective transform matrix
matrix = cv2.getPerspectiveTransform(rect.astype(np.float32), dst)
# Apply transformation
warped = cv2.warpPerspective(image, matrix, (max_width, max_height))
return warped
def _order_points(self, points: np.ndarray) -> np.ndarray:
"""
Order points in clockwise order starting from top-left.
Args:
points: Array of 4 points
Returns:
Ordered points array
"""
rect = np.zeros((4, 2), dtype=np.float32)
# Sum and difference to find corners
s = points.sum(axis=1)
diff = np.diff(points, axis=1)
rect[0] = points[np.argmin(s)] # Top-left has smallest sum
rect[2] = points[np.argmax(s)] # Bottom-right has largest sum
rect[1] = points[np.argmin(diff)] # Top-right has smallest difference
rect[3] = points[np.argmax(diff)] # Bottom-left has largest difference
return rect
def _adaptive_threshold(self, image: np.ndarray) -> np.ndarray:
"""
Apply adaptive thresholding to separate foreground from background.
Args:
image: Grayscale input image
Returns:
Binary image
"""
# Apply Gaussian blur to reduce noise before thresholding
blurred = cv2.GaussianBlur(image, (5, 5), 0)
# Apply adaptive threshold
binary = cv2.adaptiveThreshold(
blurred,
255,
cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
cv2.THRESH_BINARY,
11,
2
)
return binary
def _remove_noise(self, image: np.ndarray) -> np.ndarray:
"""
Remove noise using morphological operations.
Args:
image: Binary input image
Returns:
Denoised image
"""
# Define morphological kernels
small_kernel = np.ones((3, 3), np.uint8)
# Remove small noise with opening
opened = cv2.morphologyEx(image, cv2.MORPH_OPEN, small_kernel, iterations=1)
# Close small gaps with closing
closed = cv2.morphologyEx(opened, cv2.MORPH_CLOSE, small_kernel, iterations=1)
return closed
def _normalize_size(self, image: np.ndarray) -> np.ndarray:
"""
Normalize image to target size while preserving aspect ratio.
Args:
image: Input image
Returns:
Resized image
"""
height, width = image.shape[:2]
target_width, target_height = self.target_size
# Calculate scaling factor to fit within target size
scale = min(target_width / width, target_height / height)
# Calculate new dimensions
new_width = int(width * scale)
new_height = int(height * scale)
# Resize image
resized = cv2.resize(image, (new_width, new_height),
interpolation=cv2.INTER_AREA)
# Create canvas with target size
canvas = np.ones((target_height, target_width), dtype=np.uint8) * 255
# Calculate position to center the image
y_offset = (target_height - new_height) // 2
x_offset = (target_width - new_width) // 2
# Place resized image on canvas
canvas[y_offset:y_offset+new_height, x_offset:x_offset+new_width] = resized
return canvas
def _enhance_contrast(self, image: np.ndarray) -> np.ndarray:
"""
Enhance image contrast using CLAHE.
Args:
image: Input image
Returns:
Contrast-enhanced image
"""
# Apply CLAHE (Contrast Limited Adaptive Histogram Equalization)
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
enhanced = clahe.apply(image)
return enhanced
The preprocessing pipeline transforms raw photographs into clean, normalized images suitable for analysis. Perspective correction ensures that diagrams are viewed from a frontal orientation, eliminating distortion caused by camera angle. Adaptive thresholding handles varying lighting conditions better than global thresholding. Morphological operations remove isolated noise pixels while preserving the structure of drawn elements. Size normalization ensures consistent processing regardless of input resolution.
OPTICAL CHARACTER RECOGNITION AND TEXT EXTRACTION
Hand-drawn diagrams frequently contain text labels, annotations, and descriptions. Accurate text extraction is essential for understanding diagram semantics and preserving information in the digital output. The system must handle various handwriting styles, text orientations, and font sizes.
The text extraction component combines multiple OCR engines to maximize accuracy. It uses both Tesseract OCR for printed-style text and EasyOCR for handwritten text recognition. The component detects text regions, extracts text content, and associates text with spatial locations for later use in diagram reconstruction.
import pytesseract
import easyocr
from typing import List, Dict, Tuple
import numpy as np
from dataclasses import dataclass
@dataclass
class TextRegion:
"""Represents a detected text region with content and location."""
text: str
confidence: float
bbox: Tuple[int, int, int, int] # (x, y, width, height)
orientation: float # Angle in degrees
class TextExtractor:
"""
Extracts text from hand-drawn diagrams using multiple OCR engines.
Combines Tesseract and EasyOCR for robust text detection.
"""
def __init__(self, languages: List[str] = ['en']):
"""
Initialize text extractor.
Args:
languages: List of language codes for OCR
"""
self.languages = languages
self.logger = logging.getLogger(__name__)
# Initialize EasyOCR reader
self.logger.info("Initializing EasyOCR reader...")
self.easyocr_reader = easyocr.Reader(languages, gpu=True)
self.logger.info("EasyOCR reader initialized")
def extract_text(self, image: Union[Image.Image, np.ndarray]) -> List[TextRegion]:
"""
Extract all text regions from an image.
Args:
image: Input image as PIL Image or numpy array
Returns:
List of TextRegion objects containing detected text
"""
# Convert to numpy array if needed
if isinstance(image, Image.Image):
img_array = np.array(image)
else:
img_array = image
# Extract text using both engines
tesseract_results = self._extract_with_tesseract(img_array)
easyocr_results = self._extract_with_easyocr(img_array)
# Merge results, preferring higher confidence detections
merged_results = self._merge_text_regions(tesseract_results, easyocr_results)
self.logger.info(f"Extracted {len(merged_results)} text regions")
return merged_results
def _extract_with_tesseract(self, image: np.ndarray) -> List[TextRegion]:
"""
Extract text using Tesseract OCR.
Args:
image: Input image as numpy array
Returns:
List of TextRegion objects
"""
results = []
try:
# Get detailed OCR data
data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
# Process each detected word
n_boxes = len(data['text'])
for i in range(n_boxes):
text = data['text'][i].strip()
if text: # Only process non-empty text
confidence = float(data['conf'][i])
if confidence > 0: # Filter out low-confidence detections
x = data['left'][i]
y = data['top'][i]
w = data['width'][i]
h = data['height'][i]
results.append(TextRegion(
text=text,
confidence=confidence / 100.0, # Normalize to 0-1
bbox=(x, y, w, h),
orientation=0.0
))
except Exception as e:
self.logger.error(f"Tesseract OCR error: {e}")
return results
def _extract_with_easyocr(self, image: np.ndarray) -> List[TextRegion]:
"""
Extract text using EasyOCR.
Args:
image: Input image as numpy array
Returns:
List of TextRegion objects
"""
results = []
try:
# Perform OCR
detections = self.easyocr_reader.readtext(image)
# Process each detection
for detection in detections:
bbox_points, text, confidence = detection
# Convert bbox points to x, y, w, h format
bbox_array = np.array(bbox_points)
x = int(bbox_array[:, 0].min())
y = int(bbox_array[:, 1].min())
w = int(bbox_array[:, 0].max() - x)
h = int(bbox_array[:, 1].max() - y)
# Calculate orientation from bbox points
orientation = self._calculate_text_orientation(bbox_points)
results.append(TextRegion(
text=text,
confidence=confidence,
bbox=(x, y, w, h),
orientation=orientation
))
except Exception as e:
self.logger.error(f"EasyOCR error: {e}")
return results
def _calculate_text_orientation(self, bbox_points: List[List[float]]) -> float:
"""
Calculate text orientation angle from bounding box points.
Args:
bbox_points: List of 4 corner points
Returns:
Orientation angle in degrees
"""
# Calculate angle from top-left to top-right point
p1 = np.array(bbox_points[0])
p2 = np.array(bbox_points[1])
dx = p2[0] - p1[0]
dy = p2[1] - p1[1]
angle = np.degrees(np.arctan2(dy, dx))
return angle
def _merge_text_regions(self, regions1: List[TextRegion],
regions2: List[TextRegion]) -> List[TextRegion]:
"""
Merge text regions from multiple OCR engines, removing duplicates.
Args:
regions1: First list of text regions
regions2: Second list of text regions
Returns:
Merged list with duplicates removed
"""
merged = []
used_indices = set()
# Start with all regions from first list
for r1 in regions1:
merged.append(r1)
# Add regions from second list that don't overlap significantly
for r2 in regions2:
overlaps = False
for r1 in regions1:
if self._regions_overlap(r1, r2):
overlaps = True
break
if not overlaps:
merged.append(r2)
# Sort by position (top to bottom, left to right)
merged.sort(key=lambda r: (r.bbox[1], r.bbox[0]))
return merged
def _regions_overlap(self, r1: TextRegion, r2: TextRegion,
threshold: float = 0.5) -> bool:
"""
Check if two text regions overlap significantly.
Args:
r1: First text region
r2: Second text region
threshold: Minimum IoU to consider regions overlapping
Returns:
True if regions overlap above threshold
"""
x1, y1, w1, h1 = r1.bbox
x2, y2, w2, h2 = r2.bbox
# Calculate intersection
x_left = max(x1, x2)
y_top = max(y1, y2)
x_right = min(x1 + w1, x2 + w2)
y_bottom = min(y1 + h1, y2 + h2)
if x_right < x_left or y_bottom < y_top:
return False
intersection_area = (x_right - x_left) * (y_bottom - y_top)
# Calculate union
area1 = w1 * h1
area2 = w2 * h2
union_area = area1 + area2 - intersection_area
# Calculate IoU
iou = intersection_area / union_area if union_area > 0 else 0
return iou > threshold
The text extraction component provides robust text detection by combining complementary OCR engines. Tesseract excels at printed-style text with clear letterforms, while EasyOCR handles handwritten text more effectively. The merging algorithm eliminates duplicate detections by calculating intersection-over-union between bounding boxes. Orientation detection enables proper handling of rotated text labels.
SHAPE DETECTION AND CLASSIFICATION
Geometric shapes form the fundamental building blocks of most diagrams. Circles represent states or nodes, rectangles indicate processes or components, arrows show relationships and flow, and various specialized shapes convey domain-specific semantics. Accurate shape detection and classification is critical for understanding diagram structure.
The shape detection component uses computer vision techniques to identify geometric primitives in the preprocessed image. It detects lines, circles, rectangles, polygons, and curves using contour analysis and Hough transforms. Each detected shape is classified and parameterized for later use in digital reconstruction.
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from enum import Enum
import cv2
import numpy as np
class ShapeType(Enum):
"""Enumeration of detectable shape types."""
LINE = "line"
ARROW = "arrow"
RECTANGLE = "rectangle"
CIRCLE = "circle"
ELLIPSE = "ellipse"
TRIANGLE = "triangle"
DIAMOND = "diamond"
POLYGON = "polygon"
CURVE = "curve"
UNKNOWN = "unknown"
@dataclass
class DetectedShape:
"""Represents a detected geometric shape."""
shape_type: ShapeType
confidence: float
contour: np.ndarray
parameters: Dict[str, Any] # Shape-specific parameters
bbox: Tuple[int, int, int, int] # (x, y, width, height)
class ShapeDetector:
"""
Detects and classifies geometric shapes in hand-drawn diagrams.
Uses contour analysis and Hough transforms for robust detection.
"""
def __init__(self, min_area: int = 100, max_area: Optional[int] = None):
"""
Initialize shape detector.
Args:
min_area: Minimum contour area to consider
max_area: Maximum contour area to consider (None for no limit)
"""
self.min_area = min_area
self.max_area = max_area
self.logger = logging.getLogger(__name__)
def detect_shapes(self, image: Union[Image.Image, np.ndarray]) -> List[DetectedShape]:
"""
Detect all shapes in an image.
Args:
image: Input image as PIL Image or numpy array
Returns:
List of DetectedShape objects
"""
# Convert to numpy array if needed
if isinstance(image, Image.Image):
img_array = np.array(image)
else:
img_array = image
# Ensure binary image
if len(img_array.shape) == 3:
gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
else:
gray = img_array
_, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)
# Find contours
contours, hierarchy = cv2.findContours(
binary,
cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE
)
shapes = []
# Process each contour
for i, contour in enumerate(contours):
area = cv2.contourArea(contour)
# Filter by area
if area < self.min_area:
continue
if self.max_area and area > self.max_area:
continue
# Classify shape
shape = self._classify_shape(contour, hierarchy[0][i] if hierarchy is not None else None)
if shape:
shapes.append(shape)
# Detect lines and arrows separately using Hough transform
lines = self._detect_lines(binary)
shapes.extend(lines)
self.logger.info(f"Detected {len(shapes)} shapes")
return shapes
def _classify_shape(self, contour: np.ndarray,
hierarchy_info: Optional[np.ndarray]) -> Optional[DetectedShape]:
"""
Classify a contour as a specific shape type.
Args:
contour: Contour points
hierarchy_info: Hierarchy information for this contour
Returns:
DetectedShape object or None if classification fails
"""
# Calculate contour properties
area = cv2.contourArea(contour)
perimeter = cv2.arcLength(contour, True)
if perimeter == 0:
return None
# Approximate contour to polygon
epsilon = 0.04 * perimeter
approx = cv2.approxPolyDP(contour, epsilon, True)
num_vertices = len(approx)
# Get bounding box
x, y, w, h = cv2.boundingRect(contour)
bbox = (x, y, w, h)
# Calculate circularity
circularity = 4 * np.pi * area / (perimeter * perimeter)
# Classify based on properties
if num_vertices == 3:
return self._create_triangle(contour, approx, bbox)
elif num_vertices == 4:
return self._classify_quadrilateral(contour, approx, bbox)
elif circularity > 0.85:
return self._classify_circular(contour, bbox)
elif num_vertices > 4:
if num_vertices < 8 and circularity < 0.7:
return DetectedShape(
shape_type=ShapeType.POLYGON,
confidence=0.8,
contour=contour,
parameters={'vertices': approx.reshape(-1, 2).tolist()},
bbox=bbox
)
else:
return self._classify_curve(contour, bbox)
return None
def _classify_quadrilateral(self, contour: np.ndarray, approx: np.ndarray,
bbox: Tuple[int, int, int, int]) -> DetectedShape:
"""
Classify a 4-sided polygon as rectangle or diamond.
Args:
contour: Original contour
approx: Approximated polygon
bbox: Bounding box
Returns:
DetectedShape object
"""
x, y, w, h = bbox
aspect_ratio = float(w) / h if h > 0 else 0
# Calculate angles between edges
vertices = approx.reshape(-1, 2)
angles = []
for i in range(4):
p1 = vertices[i]
p2 = vertices[(i + 1) % 4]
p3 = vertices[(i + 2) % 4]
v1 = p1 - p2
v2 = p3 - p2
angle = np.abs(np.arccos(
np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-6)
))
angles.append(np.degrees(angle))
# Check if angles are close to 90 degrees (rectangle)
right_angles = sum(1 for angle in angles if 80 < angle < 100)
if right_angles >= 3:
return DetectedShape(
shape_type=ShapeType.RECTANGLE,
confidence=0.9,
contour=contour,
parameters={
'x': x, 'y': y, 'width': w, 'height': h,
'aspect_ratio': aspect_ratio
},
bbox=bbox
)
else:
# Check if it's a diamond (rotated square)
center_x = x + w // 2
center_y = y + h // 2
return DetectedShape(
shape_type=ShapeType.DIAMOND,
confidence=0.85,
contour=contour,
parameters={
'center_x': center_x,
'center_y': center_y,
'width': w,
'height': h,
'vertices': vertices.tolist()
},
bbox=bbox
)
def _classify_circular(self, contour: np.ndarray,
bbox: Tuple[int, int, int, int]) -> DetectedShape:
"""
Classify a circular shape as circle or ellipse.
Args:
contour: Contour points
bbox: Bounding box
Returns:
DetectedShape object
"""
x, y, w, h = bbox
aspect_ratio = float(w) / h if h > 0 else 0
# Fit ellipse if contour has enough points
if len(contour) >= 5:
ellipse = cv2.fitEllipse(contour)
center, axes, angle = ellipse
# Check if it's a circle (axes approximately equal)
axis_ratio = min(axes) / max(axes) if max(axes) > 0 else 0
if axis_ratio > 0.9:
radius = (axes[0] + axes[1]) / 4 # Average radius
return DetectedShape(
shape_type=ShapeType.CIRCLE,
confidence=0.95,
contour=contour,
parameters={
'center_x': int(center[0]),
'center_y': int(center[1]),
'radius': int(radius)
},
bbox=bbox
)
else:
return DetectedShape(
shape_type=ShapeType.ELLIPSE,
confidence=0.9,
contour=contour,
parameters={
'center_x': int(center[0]),
'center_y': int(center[1]),
'major_axis': int(max(axes)),
'minor_axis': int(min(axes)),
'angle': angle
},
bbox=bbox
)
# Fallback to circle approximation
(cx, cy), radius = cv2.minEnclosingCircle(contour)
return DetectedShape(
shape_type=ShapeType.CIRCLE,
confidence=0.8,
contour=contour,
parameters={
'center_x': int(cx),
'center_y': int(cy),
'radius': int(radius)
},
bbox=bbox
)
def _create_triangle(self, contour: np.ndarray, approx: np.ndarray,
bbox: Tuple[int, int, int, int]) -> DetectedShape:
"""
Create triangle shape object.
Args:
contour: Original contour
approx: Approximated triangle vertices
bbox: Bounding box
Returns:
DetectedShape object
"""
vertices = approx.reshape(-1, 2)
return DetectedShape(
shape_type=ShapeType.TRIANGLE,
confidence=0.9,
contour=contour,
parameters={'vertices': vertices.tolist()},
bbox=bbox
)
def _classify_curve(self, contour: np.ndarray,
bbox: Tuple[int, int, int, int]) -> DetectedShape:
"""
Classify a contour as a curve.
Args:
contour: Contour points
bbox: Bounding box
Returns:
DetectedShape object
"""
# Sample points along the curve
num_samples = min(50, len(contour))
indices = np.linspace(0, len(contour) - 1, num_samples, dtype=int)
sampled_points = contour[indices].reshape(-1, 2)
return DetectedShape(
shape_type=ShapeType.CURVE,
confidence=0.75,
contour=contour,
parameters={'points': sampled_points.tolist()},
bbox=bbox
)
def _detect_lines(self, binary_image: np.ndarray) -> List[DetectedShape]:
"""
Detect straight lines using Hough transform.
Args:
binary_image: Binary input image
Returns:
List of DetectedShape objects for lines
"""
lines_shapes = []
# Apply Hough Line Transform
lines = cv2.HoughLinesP(
binary_image,
rho=1,
theta=np.pi / 180,
threshold=50,
minLineLength=30,
maxLineGap=10
)
if lines is not None:
for line in lines:
x1, y1, x2, y2 = line[0]
# Check if line has an arrowhead
is_arrow, arrow_params = self._detect_arrowhead(
binary_image, x1, y1, x2, y2
)
# Calculate bounding box
x_min = min(x1, x2)
y_min = min(y1, y2)
x_max = max(x1, x2)
y_max = max(y1, y2)
bbox = (x_min, y_min, x_max - x_min, y_max - y_min)
if is_arrow:
lines_shapes.append(DetectedShape(
shape_type=ShapeType.ARROW,
confidence=0.85,
contour=np.array([[x1, y1], [x2, y2]]),
parameters={
'start_x': x1, 'start_y': y1,
'end_x': x2, 'end_y': y2,
**arrow_params
},
bbox=bbox
))
else:
lines_shapes.append(DetectedShape(
shape_type=ShapeType.LINE,
confidence=0.9,
contour=np.array([[x1, y1], [x2, y2]]),
parameters={
'start_x': x1, 'start_y': y1,
'end_x': x2, 'end_y': y2
},
bbox=bbox
))
return lines_shapes
def _detect_arrowhead(self, image: np.ndarray, x1: int, y1: int,
x2: int, y2: int) -> Tuple[bool, Dict[str, Any]]:
"""
Detect if a line has an arrowhead at the end.
Args:
image: Binary image
x1, y1: Line start point
x2, y2: Line end point
Returns:
Tuple of (is_arrow, arrow_parameters)
"""
# Define search region at line end
search_radius = 20
# Extract region around line end
y_min = max(0, y2 - search_radius)
y_max = min(image.shape[0], y2 + search_radius)
x_min = max(0, x2 - search_radius)
x_max = min(image.shape[1], x2 + search_radius)
region = image[y_min:y_max, x_min:x_max]
# Find contours in region
contours, _ = cv2.findContours(region, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
# Look for triangular shapes near line end
for contour in contours:
area = cv2.contourArea(contour)
if 10 < area < 200: # Reasonable arrowhead size
epsilon = 0.1 * cv2.arcLength(contour, True)
approx = cv2.approxPolyDP(contour, epsilon, True)
if len(approx) == 3: # Triangle
return True, {'arrowhead_type': 'triangle'}
return False, {}
The shape detector provides comprehensive geometric analysis of hand-drawn diagrams. Contour approximation reduces complex shapes to their essential vertices, enabling robust classification. The circularity metric distinguishes between polygons and circular shapes. Hough line detection finds straight lines that might be missed by contour analysis. Arrowhead detection identifies directional connectors, which carry important semantic meaning in flowcharts and other diagram types.
DIAGRAM TYPE RECOGNITION
Different diagram types follow distinct visual conventions and semantic rules. UML class diagrams use rectangles divided into compartments. Flowcharts employ specific shapes to represent different operation types. Circuit diagrams use standardized symbols for electronic components. Recognizing the diagram type enables the system to apply domain-specific beautification rules and ensure semantic correctness.
The diagram classifier uses a vision-language model to analyze the overall structure and identify the diagram type. It examines shape distributions, spatial relationships, and detected symbols to classify the input as a specific diagram type or generic drawing.
from typing import List, Dict, Any, Optional
from enum import Enum
from dataclasses import dataclass
class DiagramType(Enum):
"""Enumeration of recognizable diagram types."""
FLOWCHART = "flowchart"
UML_CLASS = "uml_class"
UML_SEQUENCE = "uml_sequence"
UML_ACTIVITY = "uml_activity"
CIRCUIT = "circuit"
NETWORK = "network"
ER_DIAGRAM = "er_diagram"
MINDMAP = "mindmap"
ORGANIZATIONAL = "organizational"
GENERIC = "generic"
@dataclass
class DiagramAnalysis:
"""Results of diagram type recognition."""
diagram_type: DiagramType
confidence: float
characteristics: Dict[str, Any]
suggested_rules: Dict[str, Any]
class DiagramClassifier:
"""
Classifies hand-drawn diagrams into specific types.
Uses vision-language models and heuristic analysis.
"""
def __init__(self, vision_backend: VisionLanguageBackend):
"""
Initialize diagram classifier.
Args:
vision_backend: Vision-language model backend for analysis
"""
self.vision_backend = vision_backend
self.logger = logging.getLogger(__name__)
def classify_diagram(self, image: Image.Image,
shapes: List[DetectedShape],
text_regions: List[TextRegion]) -> DiagramAnalysis:
"""
Classify a diagram into a specific type.
Args:
image: Input image
shapes: List of detected shapes
text_regions: List of detected text regions
Returns:
DiagramAnalysis object with classification results
"""
# First, try heuristic classification based on shapes
heuristic_result = self._heuristic_classification(shapes, text_regions)
# If heuristic classification is confident, use it
if heuristic_result.confidence > 0.8:
self.logger.info(f"Heuristic classification: {heuristic_result.diagram_type.value} "
f"(confidence: {heuristic_result.confidence:.2f})")
return heuristic_result
# Otherwise, use vision-language model for deeper analysis
vlm_result = self._vlm_classification(image, shapes, text_regions)
# Combine results, preferring VLM if available
if vlm_result and vlm_result.confidence > heuristic_result.confidence:
self.logger.info(f"VLM classification: {vlm_result.diagram_type.value} "
f"(confidence: {vlm_result.confidence:.2f})")
return vlm_result
else:
return heuristic_result
def _heuristic_classification(self, shapes: List[DetectedShape],
text_regions: List[TextRegion]) -> DiagramAnalysis:
"""
Classify diagram using heuristic rules based on shape patterns.
Args:
shapes: List of detected shapes
text_regions: List of detected text regions
Returns:
DiagramAnalysis object
"""
# Count shape types
shape_counts = {}
for shape in shapes:
shape_type = shape.shape_type
shape_counts[shape_type] = shape_counts.get(shape_type, 0) + 1
total_shapes = len(shapes)
if total_shapes == 0:
return DiagramAnalysis(
diagram_type=DiagramType.GENERIC,
confidence=0.5,
characteristics={},
suggested_rules={}
)
# Calculate shape proportions
arrow_ratio = shape_counts.get(ShapeType.ARROW, 0) / total_shapes
rectangle_ratio = shape_counts.get(ShapeType.RECTANGLE, 0) / total_shapes
diamond_ratio = shape_counts.get(ShapeType.DIAMOND, 0) / total_shapes
circle_ratio = shape_counts.get(ShapeType.CIRCLE, 0) / total_shapes
# Flowchart detection: high arrow ratio, mix of rectangles and diamonds
if arrow_ratio > 0.3 and rectangle_ratio > 0.2 and diamond_ratio > 0.1:
return DiagramAnalysis(
diagram_type=DiagramType.FLOWCHART,
confidence=0.85,
characteristics={
'arrow_ratio': arrow_ratio,
'rectangle_ratio': rectangle_ratio,
'diamond_ratio': diamond_ratio
},
suggested_rules={
'align_shapes': True,
'standardize_sizes': True,
'arrow_routing': 'orthogonal'
}
)
# UML Class Diagram: rectangles with internal divisions
if rectangle_ratio > 0.6 and arrow_ratio < 0.4:
# Check for compartmentalized rectangles
compartmentalized = self._check_compartmentalized_rectangles(
shapes, text_regions
)
if compartmentalized:
return DiagramAnalysis(
diagram_type=DiagramType.UML_CLASS,
confidence=0.9,
characteristics={
'rectangle_ratio': rectangle_ratio,
'compartmentalized': True
},
suggested_rules={
'align_shapes': True,
'standardize_widths': True,
'compartment_lines': True
}
)
# Circuit diagram: high symbol density, specific shapes
if circle_ratio > 0.3 or self._has_circuit_symbols(shapes):
return DiagramAnalysis(
diagram_type=DiagramType.CIRCUIT,
confidence=0.8,
characteristics={
'circle_ratio': circle_ratio,
'has_circuit_symbols': True
},
suggested_rules={
'use_standard_symbols': True,
'wire_routing': 'orthogonal'
}
)
# Network diagram: circles/ellipses connected by lines
if (circle_ratio > 0.4 or shape_counts.get(ShapeType.ELLIPSE, 0) / total_shapes > 0.4) \
and arrow_ratio > 0.2:
return DiagramAnalysis(
diagram_type=DiagramType.NETWORK,
confidence=0.75,
characteristics={
'circle_ratio': circle_ratio,
'connection_ratio': arrow_ratio
},
suggested_rules={
'node_layout': 'force_directed',
'edge_routing': 'curved'
}
)
# Mind map: tree structure radiating from center
if self._is_tree_structure(shapes):
return DiagramAnalysis(
diagram_type=DiagramType.MINDMAP,
confidence=0.7,
characteristics={
'tree_structure': True
},
suggested_rules={
'layout': 'radial',
'edge_style': 'curved'
}
)
# Default to generic
return DiagramAnalysis(
diagram_type=DiagramType.GENERIC,
confidence=0.6,
characteristics=shape_counts,
suggested_rules={'preserve_layout': True}
)
def _check_compartmentalized_rectangles(self, shapes: List[DetectedShape],
text_regions: List[TextRegion]) -> bool:
"""
Check if rectangles contain internal horizontal divisions.
Args:
shapes: List of detected shapes
text_regions: List of detected text regions
Returns:
True if compartmentalized rectangles detected
"""
rectangles = [s for s in shapes if s.shape_type == ShapeType.RECTANGLE]
for rect in rectangles:
x, y, w, h = rect.bbox
# Look for horizontal lines inside rectangle
for shape in shapes:
if shape.shape_type == ShapeType.LINE:
line_x1 = shape.parameters['start_x']
line_x2 = shape.parameters['end_x']
line_y1 = shape.parameters['start_y']
line_y2 = shape.parameters['end_y']
# Check if line is horizontal and inside rectangle
if abs(line_y1 - line_y2) < 5: # Approximately horizontal
if x < line_x1 < x + w and x < line_x2 < x + w:
if y < line_y1 < y + h:
return True
return False
def _has_circuit_symbols(self, shapes: List[DetectedShape]) -> bool:
"""
Check for presence of circuit-specific symbols.
Args:
shapes: List of detected shapes
Returns:
True if circuit symbols detected
"""
# Look for patterns characteristic of circuit symbols
# Resistors: rectangles with specific aspect ratio
# Capacitors: parallel lines
# Batteries: combination of long and short parallel lines
for shape in shapes:
if shape.shape_type == ShapeType.RECTANGLE:
aspect_ratio = shape.parameters.get('aspect_ratio', 0)
# Resistor typically has aspect ratio around 3:1
if 2.5 < aspect_ratio < 4.0:
return True
return False
def _is_tree_structure(self, shapes: List[DetectedShape]) -> bool:
"""
Determine if shapes form a tree structure.
Args:
shapes: List of detected shapes
Returns:
True if tree structure detected
"""
# Build adjacency graph from arrows/lines
nodes = [s for s in shapes if s.shape_type in
[ShapeType.RECTANGLE, ShapeType.CIRCLE, ShapeType.ELLIPSE]]
edges = [s for s in shapes if s.shape_type in
[ShapeType.ARROW, ShapeType.LINE]]
if len(nodes) < 3 or len(edges) < 2:
return False
# Simple heuristic: check if number of edges is approximately nodes - 1
# (characteristic of trees)
return abs(len(edges) - (len(nodes) - 1)) <= 2
def _vlm_classification(self, image: Image.Image,
shapes: List[DetectedShape],
text_regions: List[TextRegion]) -> Optional[DiagramAnalysis]:
"""
Use vision-language model for diagram classification.
Args:
image: Input image
shapes: List of detected shapes
text_regions: List of detected text regions
Returns:
DiagramAnalysis object or None if classification fails
"""
try:
# Construct detailed prompt for diagram analysis
prompt = f"""Analyze this hand-drawn diagram and classify it into one of these types:
- flowchart: Process flow diagram with decision points
- uml_class: UML class diagram showing classes and relationships
- uml_sequence: UML sequence diagram showing interactions over time
- uml_activity: UML activity diagram showing workflow
- circuit: Electronic circuit schematic
- network: Network topology diagram
- er_diagram: Entity-relationship database diagram
- mindmap: Mind map or concept map
- organizational: Organizational chart or hierarchy
- generic: Generic drawing not matching above types
The diagram contains {len(shapes)} shapes and {len(text_regions)} text regions.
Respond with ONLY the diagram type (one of the values above) followed by a confidence score (0-1) and a brief explanation.
Format: TYPE|CONFIDENCE|EXPLANATION"""
# Get model response
response = self.vision_backend.analyze_image(image, prompt, max_tokens=200)
# Parse response
parts = response.strip().split('|')
if len(parts) >= 3:
type_str = parts[0].strip().lower()
confidence = float(parts[1].strip())
explanation = parts[2].strip()
# Map string to enum
try:
diagram_type = DiagramType(type_str)
except ValueError:
diagram_type = DiagramType.GENERIC
return DiagramAnalysis(
diagram_type=diagram_type,
confidence=confidence,
characteristics={'vlm_explanation': explanation},
suggested_rules=self._get_default_rules(diagram_type)
)
except Exception as e:
self.logger.error(f"VLM classification error: {e}")
return None
def _get_default_rules(self, diagram_type: DiagramType) -> Dict[str, Any]:
"""
Get default beautification rules for a diagram type.
Args:
diagram_type: Classified diagram type
Returns:
Dictionary of suggested rules
"""
rules_map = {
DiagramType.FLOWCHART: {
'align_shapes': True,
'standardize_sizes': True,
'arrow_routing': 'orthogonal'
},
DiagramType.UML_CLASS: {
'align_shapes': True,
'standardize_widths': True,
'compartment_lines': True
},
DiagramType.CIRCUIT: {
'use_standard_symbols': True,
'wire_routing': 'orthogonal'
},
DiagramType.NETWORK: {
'node_layout': 'force_directed',
'edge_routing': 'curved'
},
DiagramType.MINDMAP: {
'layout': 'radial',
'edge_style': 'curved'
}
}
return rules_map.get(diagram_type, {'preserve_layout': True})
The diagram classifier combines heuristic pattern matching with deep learning analysis. Heuristic rules provide fast classification for common patterns, while the vision-language model handles ambiguous cases and complex diagrams. The confidence scoring allows the system to fall back to generic handling when classification is uncertain. Suggested rules guide the beautification process to respect diagram-specific conventions.
SEMANTIC ANALYSIS AND GRAPH CONSTRUCTION
Beyond recognizing individual shapes and text, the system must understand the semantic relationships between elements. Arrows connect nodes, text labels describe components, and spatial proximity indicates grouping. The semantic analyzer constructs a graph representation capturing these relationships.
The graph construction process associates text with nearby shapes, identifies connections between elements, and builds a structured representation of diagram semantics. This graph serves as the foundation for digital reconstruction.
from typing import List, Dict, Any, Set, Tuple, Optional
from dataclasses import dataclass, field
import networkx as nx
import numpy as np
@dataclass
class DiagramNode:
"""Represents a node in the diagram graph."""
node_id: str
shape: DetectedShape
labels: List[str] = field(default_factory=list)
properties: Dict[str, Any] = field(default_factory=dict)
@dataclass
class DiagramEdge:
"""Represents an edge in the diagram graph."""
source_id: str
target_id: str
edge_shape: Optional[DetectedShape] = None
labels: List[str] = field(default_factory=list)
properties: Dict[str, Any] = field(default_factory=dict)
class SemanticAnalyzer:
"""
Analyzes diagram semantics and constructs graph representation.
Associates text with shapes and identifies relationships.
"""
def __init__(self, proximity_threshold: int = 50):
"""
Initialize semantic analyzer.
Args:
proximity_threshold: Maximum distance for text-shape association
"""
self.proximity_threshold = proximity_threshold
self.logger = logging.getLogger(__name__)
def analyze(self, shapes: List[DetectedShape],
text_regions: List[TextRegion],
diagram_type: DiagramType) -> nx.DiGraph:
"""
Analyze diagram semantics and build graph representation.
Args:
shapes: List of detected shapes
text_regions: List of detected text regions
diagram_type: Classified diagram type
Returns:
NetworkX directed graph representing diagram structure
"""
# Create graph
graph = nx.DiGraph()
# Separate shapes into nodes and edges
node_shapes, edge_shapes = self._separate_shapes(shapes)
# Create nodes
nodes = self._create_nodes(node_shapes, text_regions)
# Add nodes to graph
for node in nodes:
graph.add_node(
node.node_id,
shape=node.shape,
labels=node.labels,
properties=node.properties
)
# Create edges based on connections
edges = self._create_edges(edge_shapes, nodes, diagram_type)
# Add edges to graph
for edge in edges:
graph.add_edge(
edge.source_id,
edge.target_id,
shape=edge.edge_shape,
labels=edge.labels,
properties=edge.properties
)
self.logger.info(f"Constructed graph with {graph.number_of_nodes()} nodes "
f"and {graph.number_of_edges()} edges")
return graph
def _separate_shapes(self, shapes: List[DetectedShape]) -> Tuple[List[DetectedShape],
List[DetectedShape]]:
"""
Separate shapes into nodes (rectangles, circles, etc.) and edges (arrows, lines).
Args:
shapes: List of all detected shapes
Returns:
Tuple of (node_shapes, edge_shapes)
"""
node_types = {ShapeType.RECTANGLE, ShapeType.CIRCLE, ShapeType.ELLIPSE,
ShapeType.TRIANGLE, ShapeType.DIAMOND, ShapeType.POLYGON}
edge_types = {ShapeType.ARROW, ShapeType.LINE}
node_shapes = [s for s in shapes if s.shape_type in node_types]
edge_shapes = [s for s in shapes if s.shape_type in edge_types]
return node_shapes, edge_shapes
def _create_nodes(self, node_shapes: List[DetectedShape],
text_regions: List[TextRegion]) -> List[DiagramNode]:
"""
Create diagram nodes from shapes and associate text labels.
Args:
node_shapes: List of shapes representing nodes
text_regions: List of detected text regions
Returns:
List of DiagramNode objects
"""
nodes = []
used_text_indices = set()
for i, shape in enumerate(node_shapes):
node_id = f"node_{i}"
# Find text regions near this shape
associated_text = []
for j, text_region in enumerate(text_regions):
if j in used_text_indices:
continue
if self._is_text_inside_shape(text_region, shape):
associated_text.append(text_region.text)
used_text_indices.add(j)
elif self._is_text_near_shape(text_region, shape):
associated_text.append(text_region.text)
used_text_indices.add(j)
# Create node
node = DiagramNode(
node_id=node_id,
shape=shape,
labels=associated_text,
properties={
'shape_type': shape.shape_type.value,
'bbox': shape.bbox
}
)
nodes.append(node)
return nodes
def _is_text_inside_shape(self, text_region: TextRegion,
shape: DetectedShape) -> bool:
"""
Check if text region is inside a shape.
Args:
text_region: Text region to check
shape: Shape to check against
Returns:
True if text is inside shape
"""
text_x, text_y, text_w, text_h = text_region.bbox
text_center_x = text_x + text_w // 2
text_center_y = text_y + text_h // 2
shape_x, shape_y, shape_w, shape_h = shape.bbox
# Check if text center is inside shape bounding box
if (shape_x <= text_center_x <= shape_x + shape_w and
shape_y <= text_center_y <= shape_y + shape_h):
# For more accurate check, use contour
point = np.array([[text_center_x, text_center_y]], dtype=np.float32)
result = cv2.pointPolygonTest(shape.contour,
(float(text_center_x), float(text_center_y)),
False)
return result >= 0
return False
def _is_text_near_shape(self, text_region: TextRegion,
shape: DetectedShape) -> bool:
"""
Check if text region is near a shape.
Args:
text_region: Text region to check
shape: Shape to check against
Returns:
True if text is near shape
"""
text_x, text_y, text_w, text_h = text_region.bbox
text_center_x = text_x + text_w // 2
text_center_y = text_y + text_h // 2
shape_x, shape_y, shape_w, shape_h = shape.bbox
shape_center_x = shape_x + shape_w // 2
shape_center_y = shape_y + shape_h // 2
# Calculate distance between centers
distance = np.sqrt((text_center_x - shape_center_x) ** 2 +
(text_center_y - shape_center_y) ** 2)
return distance < self.proximity_threshold
def _create_edges(self, edge_shapes: List[DetectedShape],
nodes: List[DiagramNode],
diagram_type: DiagramType) -> List[DiagramEdge]:
"""
Create edges by connecting nodes based on arrows and lines.
Args:
edge_shapes: List of shapes representing edges (arrows, lines)
nodes: List of diagram nodes
diagram_type: Type of diagram
Returns:
List of DiagramEdge objects
"""
edges = []
for edge_shape in edge_shapes:
# Get edge endpoints
if edge_shape.shape_type in {ShapeType.ARROW, ShapeType.LINE}:
start_x = edge_shape.parameters['start_x']
start_y = edge_shape.parameters['start_y']
end_x = edge_shape.parameters['end_x']
end_y = edge_shape.parameters['end_y']
# Find nodes at endpoints
source_node = self._find_node_at_point(nodes, start_x, start_y)
target_node = self._find_node_at_point(nodes, end_x, end_y)
if source_node and target_node:
edge = DiagramEdge(
source_id=source_node.node_id,
target_id=target_node.node_id,
edge_shape=edge_shape,
properties={
'edge_type': edge_shape.shape_type.value
}
)
edges.append(edge)
return edges
def _find_node_at_point(self, nodes: List[DiagramNode],
x: int, y: int) -> Optional[DiagramNode]:
"""
Find node at or near a specific point.
Args:
nodes: List of nodes to search
x, y: Point coordinates
Returns:
DiagramNode if found, None otherwise
"""
search_radius = 30
for node in nodes:
shape_x, shape_y, shape_w, shape_h = node.shape.bbox
# Check if point is inside or near shape bounding box
if (shape_x - search_radius <= x <= shape_x + shape_w + search_radius and
shape_y - search_radius <= y <= shape_y + shape_h + search_radius):
# More precise check using contour
result = cv2.pointPolygonTest(
node.shape.contour,
(float(x), float(y)),
True # Measure distance
)
if result >= -search_radius:
return node
return None
The semantic analyzer transforms low-level shape detections into a high-level graph structure. Text association uses both containment and proximity heuristics to handle labels both inside and adjacent to shapes. Edge creation identifies connections by finding nodes at arrow endpoints. The resulting graph captures the essential semantic structure of the diagram, independent of visual presentation details.
DIGITAL RECONSTRUCTION AND BEAUTIFICATION
The final stage transforms the semantic graph into a clean digital representation. This involves layout optimization, shape regularization, text formatting, and rendering to vector graphics. The beautification process applies diagram-type-specific rules while preserving semantic equivalence to the original drawing.
The reconstruction engine uses the classified diagram type to select appropriate layout algorithms and styling rules. It generates SVG output that can be further edited or converted to other formats.
import svgwrite
from svgwrite import cm, mm
import networkx as nx
from typing import Dict, Any, Tuple, List
import math
class DiagramReconstructor:
"""
Reconstructs clean digital diagrams from semantic graph representation.
Applies beautification and layout optimization.
"""
def __init__(self, canvas_size: Tuple[int, int] = (800, 600)):
"""
Initialize diagram reconstructor.
Args:
canvas_size: Output canvas dimensions (width, height)
"""
self.canvas_size = canvas_size
self.logger = logging.getLogger(__name__)
def reconstruct(self, graph: nx.DiGraph,
diagram_analysis: DiagramAnalysis,
output_path: str) -> str:
"""
Reconstruct diagram as clean digital SVG.
Args:
graph: Semantic graph representation
diagram_analysis: Diagram type and characteristics
output_path: Path for output SVG file
Returns:
Path to generated SVG file
"""
# Create SVG drawing
dwg = svgwrite.Drawing(output_path, size=self.canvas_size)
# Add definitions for reusable elements
self._add_definitions(dwg, diagram_analysis.diagram_type)
# Optimize layout based on diagram type
layout = self._compute_layout(graph, diagram_analysis)
# Apply beautification rules
beautified_graph = self._beautify_graph(graph, diagram_analysis)
# Render nodes
self._render_nodes(dwg, beautified_graph, layout, diagram_analysis)
# Render edges
self._render_edges(dwg, beautified_graph, layout, diagram_analysis)
# Save SVG
dwg.save()
self.logger.info(f"Diagram reconstructed and saved to {output_path}")
return output_path
def _add_definitions(self, dwg: svgwrite.Drawing, diagram_type: DiagramType):
"""
Add SVG definitions for markers and reusable elements.
Args:
dwg: SVG drawing object
diagram_type: Type of diagram
"""
# Add arrowhead marker
marker = dwg.marker(
id='arrowhead',
insert=(10, 5),
size=(10, 10),
orient='auto'
)
marker.add(dwg.path(d='M 0 0 L 10 5 L 0 10 z', fill='black'))
dwg.defs.add(marker)
# Add diamond marker for UML aggregation
if diagram_type == DiagramType.UML_CLASS:
diamond = dwg.marker(
id='diamond',
insert=(10, 5),
size=(10, 10),
orient='auto'
)
diamond.add(dwg.path(d='M 0 5 L 5 0 L 10 5 L 5 10 z',
fill='white', stroke='black'))
dwg.defs.add(diamond)
def _compute_layout(self, graph: nx.DiGraph,
diagram_analysis: DiagramAnalysis) -> Dict[str, Tuple[float, float]]:
"""
Compute optimal layout for diagram nodes.
Args:
graph: Semantic graph
diagram_analysis: Diagram analysis results
Returns:
Dictionary mapping node IDs to (x, y) positions
"""
suggested_rules = diagram_analysis.suggested_rules
# Choose layout algorithm based on diagram type
if suggested_rules.get('layout') == 'radial':
return self._radial_layout(graph)
elif suggested_rules.get('node_layout') == 'force_directed':
return self._force_directed_layout(graph)
elif diagram_analysis.diagram_type == DiagramType.FLOWCHART:
return self._hierarchical_layout(graph)
elif diagram_analysis.diagram_type == DiagramType.UML_CLASS:
return self._grid_layout(graph)
else:
return self._preserve_layout(graph)
def _hierarchical_layout(self, graph: nx.DiGraph) -> Dict[str, Tuple[float, float]]:
"""
Compute hierarchical layout suitable for flowcharts.
Args:
graph: Semantic graph
Returns:
Node positions dictionary
"""
try:
# Use Sugiyama layout (hierarchical)
pos = nx.spring_layout(graph, k=2, iterations=50)
# Scale to canvas size with margins
margin = 50
width = self.canvas_size[0] - 2 * margin
height = self.canvas_size[1] - 2 * margin
scaled_pos = {}
for node_id, (x, y) in pos.items():
scaled_x = margin + (x + 1) * width / 2
scaled_y = margin + (y + 1) * height / 2
scaled_pos[node_id] = (scaled_x, scaled_y)
return scaled_pos
except Exception as e:
self.logger.error(f"Layout computation error: {e}")
return self._preserve_layout(graph)
def _grid_layout(self, graph: nx.DiGraph) -> Dict[str, Tuple[float, float]]:
"""
Compute grid layout suitable for UML class diagrams.
Args:
graph: Semantic graph
Returns:
Node positions dictionary
"""
nodes = list(graph.nodes())
n_nodes = len(nodes)
# Calculate grid dimensions
cols = math.ceil(math.sqrt(n_nodes))
rows = math.ceil(n_nodes / cols)
# Calculate spacing
margin = 50
h_spacing = (self.canvas_size[0] - 2 * margin) / cols
v_spacing = (self.canvas_size[1] - 2 * margin) / rows
# Position nodes
pos = {}
for i, node_id in enumerate(nodes):
col = i % cols
row = i // cols
x = margin + col * h_spacing + h_spacing / 2
y = margin + row * v_spacing + v_spacing / 2
pos[node_id] = (x, y)
return pos
def _radial_layout(self, graph: nx.DiGraph) -> Dict[str, Tuple[float, float]]:
"""
Compute radial layout suitable for mind maps.
Args:
graph: Semantic graph
Returns:
Node positions dictionary
"""
nodes = list(graph.nodes())
n_nodes = len(nodes)
if n_nodes == 0:
return {}
# Place first node at center
center_x = self.canvas_size[0] / 2
center_y = self.canvas_size[1] / 2
pos = {nodes[0]: (center_x, center_y)}
# Place remaining nodes in circle
radius = min(self.canvas_size) / 3
for i, node_id in enumerate(nodes[1:], 1):
angle = 2 * math.pi * i / (n_nodes - 1)
x = center_x + radius * math.cos(angle)
y = center_y + radius * math.sin(angle)
pos[node_id] = (x, y)
return pos
def _force_directed_layout(self, graph: nx.DiGraph) -> Dict[str, Tuple[float, float]]:
"""
Compute force-directed layout.
Args:
graph: Semantic graph
Returns:
Node positions dictionary
"""
pos = nx.spring_layout(graph, k=1.5, iterations=100)
# Scale to canvas
margin = 50
width = self.canvas_size[0] - 2 * margin
height = self.canvas_size[1] - 2 * margin
scaled_pos = {}
for node_id, (x, y) in pos.items():
scaled_x = margin + (x + 1) * width / 2
scaled_y = margin + (y + 1) * height / 2
scaled_pos[node_id] = (scaled_x, scaled_y)
return scaled_pos
def _preserve_layout(self, graph: nx.DiGraph) -> Dict[str, Tuple[float, float]]:
"""
Preserve original layout from hand-drawn diagram.
Args:
graph: Semantic graph
Returns:
Node positions dictionary
"""
pos = {}
for node_id in graph.nodes():
shape = graph.nodes[node_id]['shape']
x, y, w, h = shape.bbox
center_x = x + w / 2
center_y = y + h / 2
pos[node_id] = (center_x, center_y)
return pos
def _beautify_graph(self, graph: nx.DiGraph,
diagram_analysis: DiagramAnalysis) -> nx.DiGraph:
"""
Apply beautification rules to graph.
Args:
graph: Original semantic graph
diagram_analysis: Diagram analysis with suggested rules
Returns:
Beautified graph
"""
beautified = graph.copy()
rules = diagram_analysis.suggested_rules
# Standardize node sizes if requested
if rules.get('standardize_sizes'):
self._standardize_node_sizes(beautified)
# Align shapes if requested
if rules.get('align_shapes'):
self._align_shapes(beautified)
return beautified
def _standardize_node_sizes(self, graph: nx.DiGraph):
"""
Standardize sizes of similar node types.
Args:
graph: Graph to modify in place
"""
# Group nodes by shape type
shape_groups = {}
for node_id in graph.nodes():
shape_type = graph.nodes[node_id]['properties']['shape_type']
if shape_type not in shape_groups:
shape_groups[shape_type] = []
shape_groups[shape_type].append(node_id)
# Standardize each group
for shape_type, node_ids in shape_groups.items():
if len(node_ids) < 2:
continue
# Calculate average size
total_w = 0
total_h = 0
for node_id in node_ids:
_, _, w, h = graph.nodes[node_id]['shape'].bbox
total_w += w
total_h += h
avg_w = total_w / len(node_ids)
avg_h = total_h / len(node_ids)
# Apply average size
for node_id in node_ids:
shape = graph.nodes[node_id]['shape']
x, y, _, _ = shape.bbox
shape.bbox = (x, y, int(avg_w), int(avg_h))
def _align_shapes(self, graph: nx.DiGraph):
"""
Align shapes to grid.
Args:
graph: Graph to modify in place
"""
grid_size = 20
for node_id in graph.nodes():
shape = graph.nodes[node_id]['shape']
x, y, w, h = shape.bbox
# Snap to grid
aligned_x = round(x / grid_size) * grid_size
aligned_y = round(y / grid_size) * grid_size
shape.bbox = (aligned_x, aligned_y, w, h)
def _render_nodes(self, dwg: svgwrite.Drawing, graph: nx.DiGraph,
layout: Dict[str, Tuple[float, float]],
diagram_analysis: DiagramAnalysis):
"""
Render diagram nodes to SVG.
Args:
dwg: SVG drawing object
graph: Semantic graph
layout: Node positions
diagram_analysis: Diagram analysis
"""
for node_id in graph.nodes():
node_data = graph.nodes[node_id]
shape = node_data['shape']
labels = node_data['labels']
x, y = layout[node_id]
# Render based on shape type
if shape.shape_type == ShapeType.RECTANGLE:
self._render_rectangle(dwg, x, y, shape, labels, diagram_analysis)
elif shape.shape_type == ShapeType.CIRCLE:
self._render_circle(dwg, x, y, shape, labels)
elif shape.shape_type == ShapeType.DIAMOND:
self._render_diamond(dwg, x, y, shape, labels)
elif shape.shape_type == ShapeType.ELLIPSE:
self._render_ellipse(dwg, x, y, shape, labels)
def _render_rectangle(self, dwg: svgwrite.Drawing, x: float, y: float,
shape: DetectedShape, labels: List[str],
diagram_analysis: DiagramAnalysis):
"""Render a rectangle node."""
width = shape.parameters.get('width', 100)
height = shape.parameters.get('height', 60)
rect = dwg.rect(
insert=(x - width/2, y - height/2),
size=(width, height),
fill='white',
stroke='black',
stroke_width=2
)
dwg.add(rect)
# Add text labels
if labels:
text_y = y - height/2 + 20
for label in labels:
text = dwg.text(
label,
insert=(x, text_y),
text_anchor='middle',
font_size='14px',
font_family='Arial'
)
dwg.add(text)
text_y += 20
def _render_circle(self, dwg: svgwrite.Drawing, x: float, y: float,
shape: DetectedShape, labels: List[str]):
"""Render a circle node."""
radius = shape.parameters.get('radius', 30)
circle = dwg.circle(
center=(x, y),
r=radius,
fill='white',
stroke='black',
stroke_width=2
)
dwg.add(circle)
# Add text label
if labels:
text = dwg.text(
labels[0],
insert=(x, y + 5),
text_anchor='middle',
font_size='14px',
font_family='Arial'
)
dwg.add(text)
def _render_diamond(self, dwg: svgwrite.Drawing, x: float, y: float,
shape: DetectedShape, labels: List[str]):
"""Render a diamond node."""
width = shape.parameters.get('width', 80)
height = shape.parameters.get('height', 60)
points = [
(x, y - height/2),
(x + width/2, y),
(x, y + height/2),
(x - width/2, y)
]
polygon = dwg.polygon(
points=points,
fill='white',
stroke='black',
stroke_width=2
)
dwg.add(polygon)
# Add text label
if labels:
text = dwg.text(
labels[0],
insert=(x, y + 5),
text_anchor='middle',
font_size='14px',
font_family='Arial'
)
dwg.add(text)
def _render_ellipse(self, dwg: svgwrite.Drawing, x: float, y: float,
shape: DetectedShape, labels: List[str]):
"""Render an ellipse node."""
major_axis = shape.parameters.get('major_axis', 60)
minor_axis = shape.parameters.get('minor_axis', 40)
ellipse = dwg.ellipse(
center=(x, y),
r=(major_axis/2, minor_axis/2),
fill='white',
stroke='black',
stroke_width=2
)
dwg.add(ellipse)
# Add text label
if labels:
text = dwg.text(
labels[0],
insert=(x, y + 5),
text_anchor='middle',
font_size='14px',
font_family='Arial'
)
dwg.add(text)
def _render_edges(self, dwg: svgwrite.Drawing, graph: nx.DiGraph,
layout: Dict[str, Tuple[float, float]],
diagram_analysis: DiagramAnalysis):
"""
Render diagram edges to SVG.
Args:
dwg: SVG drawing object
graph: Semantic graph
layout: Node positions
diagram_analysis: Diagram analysis
"""
for source, target in graph.edges():
x1, y1 = layout[source]
x2, y2 = layout[target]
edge_data = graph.edges[source, target]
edge_shape = edge_data.get('shape')
# Determine if arrow or plain line
has_arrow = edge_shape and edge_shape.shape_type == ShapeType.ARROW
# Create path
if diagram_analysis.suggested_rules.get('arrow_routing') == 'orthogonal':
path_d = self._create_orthogonal_path(x1, y1, x2, y2)
else:
path_d = f'M {x1} {y1} L {x2} {y2}'
# Add path with optional arrowhead
path_attrs = {
'stroke': 'black',
'stroke_width': 2,
'fill': 'none'
}
if has_arrow:
path_attrs['marker_end'] = 'url(#arrowhead)'
path = dwg.path(d=path_d, **path_attrs)
dwg.add(path)
def _create_orthogonal_path(self, x1: float, y1: float,
x2: float, y2: float) -> str:
"""
Create orthogonal (right-angle) path between two points.
Args:
x1, y1: Start point
x2, y2: End point
Returns:
SVG path data string
"""
# Simple orthogonal routing with midpoint
mid_x = (x1 + x2) / 2
return f'M {x1} {y1} L {mid_x} {y1} L {mid_x} {y2} L {x2} {y2}'
The reconstruction engine produces publication-quality diagrams from rough hand-drawn input. Layout algorithms respect diagram-type conventions while optimizing visual clarity. Shape rendering uses clean geometric primitives with consistent styling. Text placement ensures readability without overlapping elements. The SVG output format enables further editing and scaling without quality loss.
ALTERNATIVE TOOLS AND APPROACHES
Vectorization Tools: Services like Kittl, Linearity Curve, Vectorizer.AI, Adobe Illustrator (Auto Trace), and Vectr scan and convert raster sketches to scalable, clean vector graphics, including shape and text recognition.
AI Sketch Enhancement: Platforms such as Fotor, Deep-image.ai, and Playform.io offer AI-powered beautification, enhancing rough sketches and even generating digital art or realistic renders in a chosen style.
Diagram and Structure Recognition: Systems like Sketch2Scheme focus on schematic diagrams and flowcharts, recognizing drawn elements and transforming them into clean, digital diagrams.
Customization and Style: Some tools provide style transfer and editing capabilities, letting you apply specific artistic or technical styles to your improved sketch.
Supported Elements:
- Shapes and Lines: AI can detect, clarify, and vectorize freehand shapes and lines for a clean, professional result.
- Text: For hand-written labels or annotations, certain AI tools can enhance, straighten, or even convert the handwriting to editable or beautified digital text.
- Artistic and Technical Styles: Choose between technical (blueprint, diagram, architectural) and artistic (cartoon, painting, minimalist) output according to your final use case.
This technology lets anyone—from engineers and teachers to artists and hobbyists—transform a scanned sketch into a refined digital asset quickly, without needing advanced design or image editing skills.
COMPLETE PRODUCTION SYSTEM
The following complete implementation integrates all components into a production-ready system. This code provides a full pipeline from image input to SVG output, supporting all discussed features including multi-GPU acceleration, multiple LLM backends, and comprehensive diagram type recognition.
#!/usr/bin/env python3
"""
Hand-Drawn Diagram Digitization System
A complete production-ready system for converting hand-drawn diagrams
to clean digital representations with semantic preservation.
Supports:
- Multiple GPU architectures (NVIDIA CUDA, AMD ROCm, Apple MPS, Intel XPU)
- Local and remote vision-language models
- Multiple diagram types (flowcharts, UML, circuits, etc.)
- Comprehensive shape and text detection
- Semantic analysis and graph construction
- Beautification and layout optimization
"""
import logging
import argparse
import sys
from pathlib import Path
from typing import Optional
from PIL import Image
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
class DiagramDigitizationPipeline:
"""
Complete pipeline for digitizing hand-drawn diagrams.
Integrates all components from preprocessing to SVG generation.
"""
def __init__(self, backend_type: str = 'local_llama',
backend_config: Optional[dict] = None):
"""
Initialize the digitization pipeline.
Args:
backend_type: Type of vision-language backend to use
backend_config: Configuration dictionary for the backend
"""
self.logger = logging.getLogger(__name__)
# Initialize device manager
self.logger.info("Initializing device manager...")
self.device_manager = DeviceManager()
self.device_manager.optimize_for_inference()
# Initialize vision-language backend
self.logger.info(f"Initializing {backend_type} backend...")
backend_config = backend_config or {}
self.vision_backend = VisionBackendFactory.create_backend(
backend_type,
device_manager=self.device_manager,
**backend_config
)
# Initialize pipeline components
self.logger.info("Initializing pipeline components...")
self.preprocessor = ImagePreprocessor(target_size=(1024, 1024))
self.text_extractor = TextExtractor(languages=['en'])
self.shape_detector = ShapeDetector(min_area=100)
self.diagram_classifier = DiagramClassifier(self.vision_backend)
self.semantic_analyzer = SemanticAnalyzer(proximity_threshold=50)
self.reconstructor = DiagramReconstructor(canvas_size=(800, 600))
self.logger.info("Pipeline initialization complete")
def process_image(self, input_path: str, output_path: str) -> str:
"""
Process a hand-drawn diagram image and generate digital SVG.
Args:
input_path: Path to input image file
output_path: Path for output SVG file
Returns:
Path to generated SVG file
"""
self.logger.info(f"Processing image: {input_path}")
# Load image
image = Image.open(input_path)
self.logger.info(f"Loaded image: {image.size}")
# Preprocess image
self.logger.info("Preprocessing image...")
preprocessed = self.preprocessor.preprocess(image)
# Extract text
self.logger.info("Extracting text...")
text_regions = self.text_extractor.extract_text(preprocessed)
# Detect shapes
self.logger.info("Detecting shapes...")
shapes = self.shape_detector.detect_shapes(preprocessed)
# Classify diagram type
self.logger.info("Classifying diagram type...")
diagram_analysis = self.diagram_classifier.classify_diagram(
preprocessed, shapes, text_regions
)
self.logger.info(f"Diagram type: {diagram_analysis.diagram_type.value} "
f"(confidence: {diagram_analysis.confidence:.2f})")
# Perform semantic analysis
self.logger.info("Performing semantic analysis...")
semantic_graph = self.semantic_analyzer.analyze(
shapes, text_regions, diagram_analysis.diagram_type
)
# Reconstruct digital diagram
self.logger.info("Reconstructing digital diagram...")
output_svg = self.reconstructor.reconstruct(
semantic_graph, diagram_analysis, output_path
)
self.logger.info(f"Processing complete. Output saved to: {output_svg}")
return output_svg
def get_system_info(self) -> dict:
"""
Get information about the system configuration.
Returns:
Dictionary with system information
"""
return {
'device': self.device_manager.get_device_info(),
'backend': self.vision_backend.get_backend_info(),
'components': {
'preprocessor': 'ImagePreprocessor',
'text_extractor': 'TextExtractor',
'shape_detector': 'ShapeDetector',
'diagram_classifier': 'DiagramClassifier',
'semantic_analyzer': 'SemanticAnalyzer',
'reconstructor': 'DiagramReconstructor'
}
}
def main():
"""Main entry point for the diagram digitization system."""
parser = argparse.ArgumentParser(
description='Digitize hand-drawn diagrams to clean SVG representations'
)
parser.add_argument(
'input',
type=str,
help='Path to input image file'
)
parser.add_argument(
'output',
type=str,
help='Path for output SVG file'
)
parser.add_argument(
'--backend',
type=str,
choices=['local_llama', 'openai'],
default='local_llama',
help='Vision-language model backend to use'
)
parser.add_argument(
'--model',
type=str,
default='llava-hf/llava-1.5-7b-hf',
help='Model name for local backend'
)
parser.add_argument(
'--api-key',
type=str,
help='API key for remote backends (e.g., OpenAI)'
)
parser.add_argument(
'--info',
action='store_true',
help='Display system information and exit'
)
args = parser.parse_args()
# Build backend configuration
backend_config = {}
if args.backend == 'local_llama':
backend_config['model_name'] = args.model
elif args.backend == 'openai':
if not args.api_key:
print("Error: --api-key required for OpenAI backend")
sys.exit(1)
backend_config['api_key'] = args.api_key
# Initialize pipeline
try:
pipeline = DiagramDigitizationPipeline(
backend_type=args.backend,
backend_config=backend_config
)
# Display system info if requested
if args.info:
import json
info = pipeline.get_system_info()
print(json.dumps(info, indent=2))
sys.exit(0)
# Process image
output_path = pipeline.process_image(args.input, args.output)
print(f"\nSuccess! Digital diagram saved to: {output_path}")
except Exception as e:
logging.error(f"Pipeline error: {e}", exc_info=True)
sys.exit(1)
if __name__ == '__main__':
main()
This production system provides a complete command-line interface for diagram digitization. Users can specify input images, output paths, and backend configurations through command-line arguments. The system handles errors gracefully and provides detailed logging throughout the processing pipeline. The modular architecture allows easy extension with additional diagram types, backends, or processing stages.
USAGE EXAMPLES AND DEPLOYMENT
To use the system, first ensure all dependencies are installed. The required packages include PyTorch with appropriate GPU support, OpenCV, Tesseract OCR, EasyOCR, NetworkX, and svgwrite. For local vision-language models, install the Transformers library and download the desired model.
For NVIDIA CUDA systems, install PyTorch with CUDA support. For AMD ROCm, use the ROCm-enabled PyTorch build. Apple Silicon users should install PyTorch with MPS support. Intel GPU users need the Intel Extension for PyTorch.
Basic usage with a local model processes an image as follows:
python diagram_digitizer.py input_sketch.jpg output_diagram.svg --backend local_llama --model llava-hf/llava-1.5-7b-hf
For OpenAI GPT-4 Vision, provide an API key:
python diagram_digitizer.py input_sketch.jpg output_diagram.svg --backend openai --api-key YOUR_API_KEY
The system automatically detects available GPU hardware and configures acceleration accordingly. It processes the input image through all pipeline stages and generates a clean SVG representation preserving the semantic content of the original drawing.
PERFORMANCE CONSIDERATIONS AND OPTIMIZATION
The system performance depends on several factors including input image resolution, diagram complexity, and available hardware. GPU acceleration significantly improves processing speed for vision models and image processing operations. Local models provide faster inference than remote APIs but require more memory and computational resources.
Image preprocessing benefits from OpenCV's optimized implementations. The adaptive thresholding and morphological operations execute efficiently on both CPU and GPU. Text extraction with EasyOCR leverages GPU acceleration when available, substantially reducing OCR time for complex diagrams.
Shape detection using contour analysis scales linearly with image resolution. The Hough transform for line detection has higher computational complexity but remains practical for typical diagram sizes. Caching intermediate results between pipeline stages avoids redundant computation.
The semantic analysis and graph construction stages have minimal computational cost compared to vision processing. Layout optimization algorithms vary in complexity, with force-directed layouts requiring more iterations than grid or hierarchical layouts. The SVG rendering stage is lightweight and completes quickly regardless of diagram complexity.
For production deployment, consider implementing batch processing for multiple diagrams, caching preprocessed images, and using model quantization to reduce memory requirements. Distributed processing across multiple GPUs can accelerate throughput for high-volume scenarios.
CONCLUSION
This comprehensive system demonstrates how modern computer vision and artificial intelligence techniques enable robust digitization of hand-drawn diagrams. The modular architecture supports diverse hardware configurations and backend choices while maintaining clean separation of concerns. The integration of multiple OCR engines, shape detection algorithms, and vision-language models provides robust handling of varied input quality and diagram types.
The semantic preservation approach ensures that digitized diagrams maintain the meaning and intent of original drawings while applying beautification and standardization. Support for multiple diagram types with type-specific layout and styling rules produces outputs that respect domain conventions. The extensible design allows easy addition of new diagram types, backends, or processing techniques.
The production-ready implementation provides a solid foundation for building diagram digitization applications. The comprehensive error handling, logging, and configuration options enable deployment in diverse environments. The open-source technology stack ensures accessibility and allows customization for specific requirements.
Future enhancements could include interactive editing of semantic graphs before rendering, support for additional diagram types such as Gantt charts or BPMN diagrams, integration with diagramming tools through API interfaces, and machine learning-based improvement of shape classification accuracy through user feedback.
No comments:
Post a Comment