Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """ | |
| Hugging Face Spaces deployment for SPaCial AI OCR Service | |
| FastAPI HTTP Server for OCR Service using PaddleOCR | |
| """ | |
| from fastapi import FastAPI, File, UploadFile, HTTPException, Query, Body | |
| from fastapi.responses import JSONResponse | |
| from fastapi.middleware.cors import CORSMiddleware | |
| import uvicorn | |
| import cv2 | |
| import numpy as np | |
| import os | |
| import tempfile | |
| import math | |
| from pathlib import Path | |
| import logging | |
| # Set up logging first | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| # Set environment variables BEFORE importing PaddleOCR | |
| os.environ['PADDLE_PDX_MODEL_SOURCE'] = 'BOS' | |
| os.environ['PADDLE_HOME'] = '/app/.paddlex' | |
| os.environ['PADDLEX_HOME'] = '/app/.paddlex' | |
| os.environ['PADDLEOCR_HOME'] = '/app/.paddleocr' | |
| os.environ['TEMP'] = '/app/temp' | |
| os.environ['TMP'] = '/app/temp' | |
| os.environ['TMPDIR'] = '/app/temp' | |
| os.environ['HOME'] = '/app' | |
| os.environ['USER'] = 'app' | |
| # Let PaddleOCR auto-detect GPU/CPU (newer versions handle this automatically) | |
| # Create directories with proper permissions for all PaddleOCR models | |
| try: | |
| os.makedirs('/app/.paddlex', mode=0o777, exist_ok=True) | |
| os.makedirs('/app/.paddleocr', mode=0o777, exist_ok=True) | |
| os.makedirs('/app/.paddleocr/whl', mode=0o777, exist_ok=True) | |
| # Detection model | |
| os.makedirs('/app/.paddleocr/whl/det/en/en_PP-OCRv3_det_infer', mode=0o777, exist_ok=True) | |
| # Recognition model | |
| os.makedirs('/app/.paddleocr/whl/rec/en/en_PP-OCRv3_rec_infer', mode=0o777, exist_ok=True) | |
| # Classification model | |
| os.makedirs('/app/.paddleocr/whl/cls/en_ppocr_mobile_v2.0_cls_infer', mode=0o777, exist_ok=True) | |
| os.makedirs('/app/temp', mode=0o777, exist_ok=True) | |
| logger.info("Successfully created all PaddleOCR model directories with write permissions") | |
| except Exception as e: | |
| logger.error(f"Failed to create directories: {e}") | |
| # Now import PaddleOCR | |
| from paddleocr import PaddleOCR | |
| # Initialize FastAPI app | |
| app = FastAPI( | |
| title="SPaCial AI OCR Service", | |
| description="OCR service for dimension detection using PaddleOCR", | |
| version="1.0.0" | |
| ) | |
| # Add CORS middleware for HF Spaces | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], # Allow all origins for HF Spaces | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # Global OCR instance (initialize once at startup) | |
| ocr = None | |
| def is_dimension_text(text): | |
| """Check if text looks like a dimension""" | |
| import re | |
| # Look for numbers, possibly with units or decimal points | |
| if re.search(r'\d+\.?\d*\s*(mm|cm|m|in|inch)?', text.strip(), re.IGNORECASE): | |
| return True | |
| # Look for pure numbers | |
| if re.match(r'^\d+\.?\d*$', text.strip()): | |
| return True | |
| return False | |
| def parse_tolerance(text): | |
| """Parse tolerance information from text with improved patterns""" | |
| import re | |
| if not text: | |
| return None | |
| # Clean the text and handle common OCR errors | |
| clean_text = text.strip() | |
| # Fix common OCR errors | |
| clean_text = clean_text.replace('0', 'Ø') # Fix Ø being read as 0 | |
| clean_text = clean_text.replace('ØØ', 'Ø') # Fix double Ø | |
| clean_text = clean_text.replace('O', 'Ø') # Fix O being read as Ø | |
| clean_text = clean_text.replace('°', '') # Remove degree symbols that might interfere | |
| clean_text = re.sub(r'\s+', ' ', clean_text) # Normalize spaces | |
| # Pattern 1: ± tolerance (e.g., "28±0.03", "104±0.04", "Ø48±0.03", "28 ± 0.03") | |
| plus_minus_match = re.search(r'(?:Ø|DIA|DIAMETER)?\s*(\d+\.?\d*)\s*±\s*(\d+\.?\d*)', clean_text, re.IGNORECASE) | |
| if plus_minus_match: | |
| value = float(plus_minus_match.group(1)) | |
| tolerance = float(plus_minus_match.group(2)) | |
| is_diameter = bool(re.search(r'[ØDIA]', clean_text, re.IGNORECASE)) | |
| return { | |
| "value": value, | |
| "tolerance_plus": tolerance, | |
| "tolerance_minus": tolerance, | |
| "tolerance_type": "±", | |
| "is_diameter": is_diameter, | |
| "original_text": text | |
| } | |
| # Pattern 2: +tolerance/-tolerance (e.g., "25.5+0.1/-0.05", "Ø48+0.1/-0.05", "25.5 +0.1 / -0.05") | |
| asymmetric_match = re.search(r'(?:Ø|DIA|DIAMETER)?\s*(\d+\.?\d*)\s*\+\s*(\d+\.?\d*)\s*/\s*-\s*(\d+\.?\d*)', clean_text, re.IGNORECASE) | |
| if asymmetric_match: | |
| value = float(asymmetric_match.group(1)) | |
| plus_tolerance = float(asymmetric_match.group(2)) | |
| minus_tolerance = float(asymmetric_match.group(3)) | |
| is_diameter = bool(re.search(r'[ØDIA]', clean_text, re.IGNORECASE)) | |
| return { | |
| "value": value, | |
| "tolerance_plus": plus_tolerance, | |
| "tolerance_minus": minus_tolerance, | |
| "tolerance_type": "+/-", | |
| "is_diameter": is_diameter, | |
| "original_text": text | |
| } | |
| # Pattern 3: Value with negative tolerance only (e.g., "Ø48-0.03", "48 -0.03") | |
| negative_only_match = re.search(r'(?:Ø|DIA|DIAMETER)?\s*(\d+\.?\d*)\s*-\s*(\d+\.?\d*)', clean_text, re.IGNORECASE) | |
| if negative_only_match: | |
| value = float(negative_only_match.group(1)) | |
| minus_tolerance = float(negative_only_match.group(2)) | |
| is_diameter = bool(re.search(r'[ØDIA]', clean_text, re.IGNORECASE)) | |
| return { | |
| "value": value, | |
| "tolerance_plus": 0.0, | |
| "tolerance_minus": minus_tolerance, | |
| "tolerance_type": "-only", | |
| "is_diameter": is_diameter, | |
| "original_text": text | |
| } | |
| # Pattern 4: Value with positive tolerance only (e.g., "Ø48+0.03", "48 +0.03") | |
| positive_only_match = re.search(r'(?:Ø|DIA|DIAMETER)?\s*(\d+\.?\d*)\s*\+\s*(\d+\.?\d*)', clean_text, re.IGNORECASE) | |
| if positive_only_match: | |
| value = float(positive_only_match.group(1)) | |
| plus_tolerance = float(positive_only_match.group(2)) | |
| is_diameter = bool(re.search(r'[ØDIA]', clean_text, re.IGNORECASE)) | |
| return { | |
| "value": value, | |
| "tolerance_plus": plus_tolerance, | |
| "tolerance_minus": 0.0, | |
| "tolerance_type": "+only", | |
| "is_diameter": is_diameter, | |
| "original_text": text | |
| } | |
| # Pattern 5: ISO tolerance classes (e.g., "48H7", "25.5f6", "Ø48H7") | |
| iso_tolerance_match = re.search(r'(?:Ø|DIA|DIAMETER)?\s*(\d+\.?\d*)\s*([A-Za-z]\d+)', clean_text, re.IGNORECASE) | |
| if iso_tolerance_match: | |
| value = float(iso_tolerance_match.group(1)) | |
| tolerance_class = iso_tolerance_match.group(2).upper() | |
| is_diameter = bool(re.search(r'[ØDIA]', clean_text, re.IGNORECASE)) | |
| # Basic ISO tolerance lookup (simplified) | |
| iso_tolerances = { | |
| 'H7': {'plus': 0.025, 'minus': 0.0}, | |
| 'H8': {'plus': 0.039, 'minus': 0.0}, | |
| 'H9': {'plus': 0.062, 'minus': 0.0}, | |
| 'F6': {'plus': 0.0, 'minus': 0.013}, | |
| 'F7': {'plus': 0.0, 'minus': 0.025}, | |
| 'G6': {'plus': 0.0, 'minus': 0.009}, | |
| 'G7': {'plus': 0.0, 'minus': 0.020} | |
| } | |
| tolerance_info = iso_tolerances.get(tolerance_class, {'plus': 0.0, 'minus': 0.0}) | |
| return { | |
| "value": value, | |
| "tolerance_plus": tolerance_info['plus'], | |
| "tolerance_minus": tolerance_info['minus'], | |
| "tolerance_type": f"ISO {tolerance_class}", | |
| "is_diameter": is_diameter, | |
| "original_text": text | |
| } | |
| # Pattern 6: Single value with no tolerance (e.g., "48", "Ø48", "DIA 48") | |
| single_value_match = re.search(r'(?:Ø|DIA|DIAMETER)?\s*(\d+\.?\d*)', clean_text, re.IGNORECASE) | |
| if single_value_match: | |
| value = float(single_value_match.group(1)) | |
| is_diameter = bool(re.search(r'[ØDIA]', clean_text, re.IGNORECASE)) | |
| return { | |
| "value": value, | |
| "tolerance_plus": 0.0, | |
| "tolerance_minus": 0.0, | |
| "tolerance_type": "none", | |
| "is_diameter": is_diameter, | |
| "original_text": text | |
| } | |
| return None | |
| def merge_vertical_text(zones): | |
| """Merge nearby vertical text that might have been split by OCR""" | |
| if len(zones) < 2: | |
| return zones | |
| merged_zones = [] | |
| used_indices = set() | |
| for i, zone1 in enumerate(zones): | |
| if i in used_indices: | |
| continue | |
| bbox1 = zone1.get('bbox', []) | |
| if len(bbox1) < 4: | |
| merged_zones.append(zone1) | |
| continue | |
| x1_1, y1_1, x2_1, y2_1 = bbox1[0], bbox1[1], bbox1[2], bbox1[3] | |
| text1 = zone1.get('text', '') | |
| # Look for nearby zones to merge | |
| merged_text = text1 | |
| merged_bbox = [x1_1, y1_1, x2_1, y2_1] | |
| for j, zone2 in enumerate(zones[i+1:], i+1): | |
| if j in used_indices: | |
| continue | |
| bbox2 = zone2.get('bbox', []) | |
| if len(bbox2) < 4: | |
| continue | |
| x1_2, y1_2, x2_2, y2_2 = bbox2[0], bbox2[1], bbox2[2], bbox2[3] | |
| text2 = zone2.get('text', '') | |
| # Check if zones are vertically close and horizontally aligned | |
| vertical_distance = abs((y1_1 + y2_1) / 2 - (y1_2 + y2_2) / 2) | |
| horizontal_overlap = not (x2_1 < x1_2 or x2_2 < x1_1) | |
| # If zones are close vertically and overlap horizontally, merge them | |
| if vertical_distance < 30 and horizontal_overlap: | |
| # Merge text | |
| if y1_2 < y1_1: # zone2 is above zone1 | |
| merged_text = text2 + text1 | |
| merged_bbox = [min(x1_1, x1_2), min(y1_1, y1_2), | |
| max(x2_1, x2_2), max(y2_1, y2_2)] | |
| else: # zone1 is above zone2 | |
| merged_text = text1 + text2 | |
| merged_bbox = [min(x1_1, x1_2), min(y1_1, y1_2), | |
| max(x2_1, x2_2), max(y2_1, y2_2)] | |
| used_indices.add(j) | |
| logger.info(f"Merged vertical text: '{text1}' + '{text2}' = '{merged_text}'") | |
| # Create merged zone | |
| merged_zone = { | |
| "id": f"ocr_zone_{len(merged_zones)}", | |
| "text": merged_text, | |
| "confidence": zone1.get('confidence', 0.9), | |
| "bbox": merged_bbox, | |
| "x": float(merged_bbox[0]), | |
| "y": float(merged_bbox[1]), | |
| "width": float(merged_bbox[2] - merged_bbox[0]), | |
| "height": float(merged_bbox[3] - merged_bbox[1]), | |
| "orientation": zone1.get('orientation', 0), | |
| "rotation": 0 | |
| } | |
| merged_zones.append(merged_zone) | |
| used_indices.add(i) | |
| # Add any remaining zones that weren't merged | |
| for i, zone in enumerate(zones): | |
| if i not in used_indices: | |
| merged_zones.append(zone) | |
| logger.info(f"Merged {len(zones)} zones into {len(merged_zones)} zones") | |
| return merged_zones | |
| def create_overlay_image(image_path, zones, lines=None, dimension_lines=None): | |
| """Create overlay image with bounding boxes and lines (using direct coordinates, no rotation)""" | |
| try: | |
| # Load original image | |
| img = cv2.imread(str(image_path)) | |
| if img is None: | |
| return None | |
| # Create overlay with better visualization | |
| overlay = img.copy() | |
| # Create a separate overlay for semi-transparent rectangles | |
| overlay_color = np.zeros_like(img) | |
| # Draw lines first (so they appear behind text boxes) | |
| if lines: | |
| for line in lines: | |
| x1, y1, x2, y2 = int(line['x1']), int(line['y1']), int(line['x2']), int(line['y2']) | |
| cv2.line(overlay, (x1, y1), (x2, y2), (128, 128, 128), 1) # Gray lines | |
| # Draw dimension lines in different color | |
| if dimension_lines: | |
| for dim_line in dimension_lines: | |
| line = dim_line['closest_line'] | |
| x1, y1, x2, y2 = int(line['x1']), int(line['y1']), int(line['x2']), int(line['y2']) | |
| cv2.line(overlay, (x1, y1), (x2, y2), (0, 255, 255), 2) # Yellow for dimension lines | |
| # Draw clean bounding boxes for each zone | |
| for i, zone in enumerate(zones): | |
| bbox = zone["bbox"] | |
| # Color for this zone (cycling through colors) | |
| colors = [(0, 255, 0), (255, 0, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255)] | |
| color = colors[i % len(colors)] | |
| # Draw semi-transparent rectangle | |
| cv2.rectangle(overlay_color, | |
| (bbox["x1"], bbox["y1"]), | |
| (bbox["x2"], bbox["y2"]), | |
| color, -1) # Filled rectangle | |
| # Draw clean border | |
| cv2.rectangle(overlay, | |
| (bbox["x1"], bbox["y1"]), | |
| (bbox["x2"], bbox["y2"]), | |
| color, 2) # Border only | |
| # Draw clean text label with background | |
| label_text = f"{i+1}. {zone['text']} ({int(zone['confidence']*100)}%)" | |
| font_scale = 0.5 | |
| thickness = 1 | |
| (text_width, text_height), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness) | |
| # Position label above the box | |
| label_x = bbox["x1"] | |
| label_y = max(bbox["y1"] - 10, text_height + 5) | |
| # Draw text background | |
| cv2.rectangle(overlay, | |
| (label_x - 2, label_y - text_height - 2), | |
| (label_x + text_width + 2, label_y + 2), | |
| (255, 255, 255), -1) # White background | |
| # Draw text border for better readability | |
| cv2.putText(overlay, label_text, | |
| (label_x, label_y), | |
| cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thickness + 1) # White border | |
| cv2.putText(overlay, label_text, | |
| (label_x, label_y), | |
| cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thickness) # Black text | |
| # Blend the colored overlay with the original image | |
| alpha = 0.3 | |
| overlay = cv2.addWeighted(overlay, 1, overlay_color, alpha, 0) | |
| # Save overlay image | |
| overlay_path = str(image_path).replace('.', '_overlay.') | |
| cv2.imwrite(overlay_path, overlay) | |
| return overlay_path | |
| except Exception as e: | |
| logger.error(f"Error creating overlay: {e}") | |
| return None | |
| def process_image(image_path, mode="fast"): | |
| """Process a single image and return OCR results with specified mode""" | |
| try: | |
| # Get image dimensions for coordinate transformation | |
| img = cv2.imread(str(image_path)) | |
| height, width = img.shape[:2] | |
| # Use single OCR model (mode parameter kept for future use) | |
| logger.info(f"Processing image with OCR mode: {mode}") | |
| # Use ocr.predict() method (PaddleOCR 2.8+ API) | |
| result = ocr.predict(str(image_path)) | |
| zones = [] | |
| for res in result: | |
| if hasattr(res, 'json') and 'res' in res.json: | |
| res_data = res.json['res'] | |
| if 'rec_texts' in res_data and 'rec_scores' in res_data: | |
| texts = res_data['rec_texts'] | |
| scores = res_data['rec_scores'] | |
| # Get detection polygon coordinates (more accurate than rec_boxes) | |
| dt_polys = res_data.get('dt_polys', []) | |
| # Get text orientation information | |
| text_orientations = res_data.get('textline_orientation_angles', []) | |
| for i, (text, score) in enumerate(zip(texts, scores)): | |
| if text.strip(): | |
| # Use coordinates as-is (no rotation) since orientation detection is disabled | |
| if i < len(dt_polys): | |
| poly = dt_polys[i] | |
| # Convert polygon to bounding box (direct coordinates) | |
| x_coords = [p[0] for p in poly] | |
| y_coords = [p[1] for p in poly] | |
| x1, x2 = min(x_coords), max(x_coords) | |
| y1, y2 = min(y_coords), max(y_coords) | |
| # Get text orientation (if available) | |
| text_orientation = 0 | |
| if i < len(text_orientations): | |
| orientation_angle = text_orientations[i] | |
| # Convert PaddleOCR orientation to rotation angle | |
| # PaddleOCR returns values: 0=horizontal, 1=90°, 2=180°, 3=270° | |
| if orientation_angle == 1: | |
| text_orientation = 90 | |
| elif orientation_angle == 2: | |
| text_orientation = 180 | |
| elif orientation_angle == 3: | |
| text_orientation = 270 | |
| else: | |
| text_orientation = 0 | |
| # Calculate text dimensions | |
| text_width = x2 - x1 | |
| text_height = y2 - y1 | |
| # Parse tolerance information | |
| tolerance_info = parse_tolerance(text) | |
| # Check if this looks like a dimension | |
| is_dimension = is_dimension_text(text) | |
| # Create zone with proper bbox format | |
| zone = { | |
| "text": text, | |
| "confidence": score, | |
| "bbox": { | |
| "x1": int(x1), | |
| "y1": int(y1), | |
| "x2": int(x2), | |
| "y2": int(y2), | |
| "width": int(text_width), | |
| "height": int(text_height) | |
| }, | |
| "orientation": text_orientation, | |
| "rotation": text_orientation, # Keep for compatibility | |
| "width": int(text_width), | |
| "height": int(text_height), | |
| "is_dimension": is_dimension, | |
| "tolerance_info": tolerance_info | |
| } | |
| zones.append(zone) | |
| # Create overlay image with bounding boxes | |
| overlay_path = create_overlay_image(image_path, zones) | |
| return { | |
| "zones": zones, | |
| "metadata": { | |
| "total_zones": len(zones), | |
| "overlay_path": overlay_path, | |
| "detected_angle": 0 # No rotation - detect as-is | |
| } | |
| } | |
| except Exception as e: | |
| logger.error(f"Error processing image: {e}") | |
| return { | |
| "zones": [], | |
| "metadata": {"error": str(e)} | |
| } | |
| async def startup_event(): | |
| """Initialize PaddleOCR models on startup""" | |
| global ocr | |
| logger.info("Initializing PaddleOCR model...") | |
| # Single OCR instance (auto GPU detection for newer PaddleOCR) | |
| try: | |
| logger.info("Initializing PaddleOCR with auto GPU detection...") | |
| ocr = PaddleOCR( | |
| use_doc_orientation_classify=False, # Disable for better performance | |
| use_doc_unwarping=False, | |
| use_textline_orientation=True, # Enable text line orientation | |
| lang='en' | |
| ) | |
| logger.info("PaddleOCR initialized successfully with auto GPU detection!") | |
| except Exception as error: | |
| logger.error(f"PaddleOCR initialization failed: {error}") | |
| raise | |
| logger.info("PaddleOCR initialized successfully!") | |
| async def root(): | |
| """Health check endpoint""" | |
| return { | |
| "message": "SPaCial AI OCR Service", | |
| "status": "running", | |
| "ocr_initialized": ocr is not None, | |
| "service": "SPaCial AI OCR Service", | |
| "version": "1.0.0" | |
| } | |
| async def process_ocr(file: UploadFile = File(...), mode: str = Query("fast", description="OCR mode: 'fast' for position detection, 'accurate' for hard text search")): | |
| """ | |
| Process uploaded image with OCR | |
| Returns detected zones with bounding boxes and text | |
| """ | |
| if ocr is None: | |
| raise HTTPException(status_code=503, detail="OCR service not initialized") | |
| # Validate file type | |
| if not file.content_type or not file.content_type.startswith('image/'): | |
| raise HTTPException(status_code=400, detail="File must be an image") | |
| # Save uploaded file to temporary location | |
| with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file: | |
| content = await file.read() | |
| temp_file.write(content) | |
| temp_path = temp_file.name | |
| try: | |
| logger.info(f"Processing image: {temp_path}, mode: {mode}") | |
| # Process image with specified mode | |
| result = process_image(temp_path, mode) | |
| logger.info(f"OCR result: {len(result.get('zones', []))} zones found") | |
| # Clean up temporary file | |
| os.unlink(temp_path) | |
| return JSONResponse(content=result) | |
| except Exception as e: | |
| # Clean up temporary file on error | |
| if os.path.exists(temp_path): | |
| os.unlink(temp_path) | |
| logger.error(f"Error processing uploaded file: {e}") | |
| import traceback | |
| logger.error(f"Traceback: {traceback.format_exc()}") | |
| raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}") | |
| async def process_ocr_path(image_path: str = Query(...), mode: str = Query("fast", description="OCR mode: 'fast' for position detection, 'accurate' for hard text search")): | |
| """ | |
| Process image from file path | |
| Returns detected zones with bounding boxes and text | |
| """ | |
| if ocr is None: | |
| raise HTTPException(status_code=503, detail="OCR service not initialized") | |
| if not os.path.exists(image_path): | |
| raise HTTPException(status_code=404, detail="Image file not found") | |
| try: | |
| result = process_image(image_path, mode) | |
| return JSONResponse(content=result) | |
| except Exception as e: | |
| logger.error(f"Error processing image path: {e}") | |
| raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}") | |
| async def process_center_point(request: dict = Body(...)): | |
| """Process center point to find dimension and fit box to text""" | |
| try: | |
| # Extract data from request | |
| image_data = request.get('image', '') | |
| center_point = request.get('center_point', {}) | |
| center_x = center_point.get('x', 0) | |
| center_y = center_point.get('y', 0) | |
| if not image_data: | |
| raise HTTPException(status_code=400, detail="No image data provided") | |
| # Decode base64 image | |
| import base64 | |
| if image_data.startswith('data:image'): | |
| image_data = image_data.split(',')[1] | |
| image_bytes = base64.b64decode(image_data) | |
| # Save temporary image | |
| with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file: | |
| tmp_file.write(image_bytes) | |
| tmp_path = tmp_file.name | |
| try: | |
| # Process image with OCR | |
| result = process_image(tmp_path, mode="fast") | |
| if not result or not result.get('zones'): | |
| return {"zones": [], "message": "No text detected"} | |
| # Find the closest zone to the center point | |
| zones = result['zones'] | |
| closest_zone = None | |
| min_distance = float('inf') | |
| for zone in zones: | |
| bbox = zone.get('bbox', {}) | |
| if not bbox: | |
| continue | |
| # Calculate zone center | |
| zone_center_x = (bbox.get('x1', 0) + bbox.get('x2', 0)) / 2 | |
| zone_center_y = (bbox.get('y1', 0) + bbox.get('y2', 0)) / 2 | |
| # Calculate distance to clicked point | |
| distance = ((center_x - zone_center_x) ** 2 + (center_y - zone_center_y) ** 2) ** 0.5 | |
| if distance < min_distance: | |
| min_distance = distance | |
| closest_zone = zone | |
| # If we found a close zone, return it | |
| if closest_zone and min_distance < 100: # Within 100 pixels | |
| # Ensure we return the zone with proper bbox coordinates | |
| zone_bbox = closest_zone.get('bbox', {}) | |
| if isinstance(zone_bbox, dict): | |
| # Return zone with image coordinates | |
| return { | |
| "zone": { | |
| "text": closest_zone.get('text', ''), | |
| "confidence": closest_zone.get('confidence', 0), | |
| "bbox": zone_bbox, | |
| "x": zone_bbox.get('x1', 0), | |
| "y": zone_bbox.get('y1', 0), | |
| "width": zone_bbox.get('width', 0), | |
| "height": zone_bbox.get('height', 0), | |
| "tolerance_info": closest_zone.get('tolerance_info', {}) | |
| }, | |
| "message": f"Found dimension at center point (distance: {min_distance:.1f}px)", | |
| "center_point": {"x": center_x, "y": center_y}, | |
| "found_zone": closest_zone.get('text', ''), | |
| "fitted": True | |
| } | |
| else: | |
| return { | |
| "zone": closest_zone, | |
| "message": f"Found dimension at center point (distance: {min_distance:.1f}px)", | |
| "center_point": {"x": center_x, "y": center_y}, | |
| "found_zone": closest_zone.get('text', ''), | |
| "fitted": True | |
| } | |
| else: | |
| return { | |
| "zones": [], | |
| "message": f"No dimension found near center point (closest: {min_distance:.1f}px)", | |
| "center_point": {"x": center_x, "y": center_y} | |
| } | |
| finally: | |
| # Clean up temporary file | |
| if os.path.exists(tmp_path): | |
| os.unlink(tmp_path) | |
| except Exception as e: | |
| logger.error(f"Error processing center point: {e}") | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def process_with_lines(request: dict = Body(...)): | |
| """Process image with OCR and detect dimension lines""" | |
| try: | |
| # Extract data from request | |
| image_data = request.get('image', '') | |
| if not image_data: | |
| raise HTTPException(status_code=400, detail="No image data provided") | |
| # Decode base64 image | |
| import base64 | |
| if image_data.startswith('data:image'): | |
| image_data = image_data.split(',')[1] | |
| image_bytes = base64.b64decode(image_data) | |
| # Save temporary image | |
| with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file: | |
| tmp_file.write(image_bytes) | |
| tmp_path = tmp_file.name | |
| try: | |
| # Process image with OCR | |
| result = process_image(tmp_path, mode="fast") | |
| if not result or not result.get('zones'): | |
| return {"zones": [], "lines": [], "dimension_lines": [], "message": "No text detected"} | |
| # Detect lines in the image | |
| detected_lines = detect_dimension_lines(tmp_path) | |
| # Find lines near dimensions | |
| dimension_lines = find_lines_near_dimensions(result['zones'], detected_lines) | |
| # Create overlay image with lines | |
| overlay_image = create_overlay_image(tmp_path, result['zones'], detected_lines, dimension_lines) | |
| # Convert overlay to base64 | |
| overlay_base64 = None | |
| if overlay_image is not None: | |
| import base64 | |
| _, buffer = cv2.imencode('.jpg', overlay_image) | |
| overlay_base64 = base64.b64encode(buffer).decode('utf-8') | |
| return { | |
| "zones": result['zones'], | |
| "lines": detected_lines, | |
| "dimension_lines": dimension_lines, | |
| "overlay_image": overlay_base64, | |
| "message": f"Found {len(result['zones'])} text zones, {len(detected_lines)} lines, {len(dimension_lines)} dimension-line pairs" | |
| } | |
| finally: | |
| # Clean up temporary file | |
| if os.path.exists(tmp_path): | |
| os.unlink(tmp_path) | |
| except Exception as e: | |
| logger.error(f"Error processing image with lines: {e}") | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def get_text_properties(request: dict = Body(...)): | |
| """Get original text properties for precise box fitting""" | |
| try: | |
| # Extract data from request | |
| image_data = request.get('image', '') | |
| bbox = request.get('bbox', {}) | |
| if not image_data or not bbox: | |
| raise HTTPException(status_code=400, detail="Image data and bbox required") | |
| # Decode base64 image | |
| import base64 | |
| if image_data.startswith('data:image'): | |
| image_data = image_data.split(',')[1] | |
| image_bytes = base64.b64decode(image_data) | |
| # Save temporary image | |
| with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file: | |
| tmp_file.write(image_bytes) | |
| tmp_path = tmp_file.name | |
| try: | |
| # Process image with OCR | |
| result = process_image(tmp_path, mode="fast") | |
| if not result or not result.get('zones'): | |
| return {"error": "No text detected"} | |
| # Find the zone that matches the bbox | |
| target_zone = None | |
| for zone in result['zones']: | |
| zone_bbox = zone.get('bbox', {}) | |
| if (isinstance(zone_bbox, dict) and | |
| abs(zone_bbox.get('x1', 0) - bbox.get('x1', 0)) < 10 and | |
| abs(zone_bbox.get('y1', 0) - bbox.get('y1', 0)) < 10): | |
| target_zone = zone | |
| break | |
| if not target_zone: | |
| return {"error": "Matching zone not found"} | |
| # Extract text properties | |
| text_properties = { | |
| "text": target_zone.get('text', ''), | |
| "confidence": target_zone.get('confidence', 0), | |
| "bbox": target_zone.get('bbox', {}), | |
| "orientation": target_zone.get('orientation', 0), | |
| "rotation": target_zone.get('rotation', 0), | |
| "width": target_zone.get('width', 0), | |
| "height": target_zone.get('height', 0) | |
| } | |
| return text_properties | |
| finally: | |
| # Clean up temporary file | |
| if os.path.exists(tmp_path): | |
| os.unlink(tmp_path) | |
| except Exception as e: | |
| logger.error(f"Error getting text properties: {e}") | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def process_baseline(request: dict = Body(...)): | |
| """Process baseline to find text and fit precise box""" | |
| try: | |
| # Extract data from request | |
| image_data = request.get('image', '') | |
| baseline = request.get('baseline', {}) | |
| x1 = baseline.get('x1', 0) | |
| y1 = baseline.get('y1', 0) | |
| x2 = baseline.get('x2', 0) | |
| y2 = baseline.get('y2', 0) | |
| if not image_data: | |
| raise HTTPException(status_code=400, detail="No image data provided") | |
| # Decode base64 image | |
| import base64 | |
| if image_data.startswith('data:image'): | |
| image_data = image_data.split(',')[1] | |
| image_bytes = base64.b64decode(image_data) | |
| # Save temporary image | |
| with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file: | |
| tmp_file.write(image_bytes) | |
| tmp_path = tmp_file.name | |
| try: | |
| # Process full image with OCR | |
| result = process_image(tmp_path, mode="fast") | |
| if not result or not result.get('zones'): | |
| return {"zones": [], "message": "No text detected"} | |
| # Find text zones that intersect with the baseline | |
| zones = result['zones'] | |
| intersecting_zones = [] | |
| for zone in zones: | |
| bbox = zone.get('bbox', {}) | |
| if not bbox: | |
| continue | |
| # Check if zone intersects with baseline | |
| if intersects_with_line(bbox, x1, y1, x2, y2): | |
| # Calculate distance from zone center to baseline | |
| zone_center_x = (bbox.get('x1', 0) + bbox.get('x2', 0)) / 2 | |
| zone_center_y = (bbox.get('y1', 0) + bbox.get('y2', 0)) / 2 | |
| distance = distance_point_to_line(zone_center_x, zone_center_y, x1, y1, x2, y2) | |
| intersecting_zones.append({ | |
| 'zone': zone, | |
| 'distance': distance | |
| }) | |
| # Sort by distance to baseline (closest first) | |
| intersecting_zones.sort(key=lambda x: x['distance']) | |
| if intersecting_zones: | |
| # Return the closest zone | |
| closest = intersecting_zones[0] | |
| zone = closest['zone'] | |
| zone_bbox = zone.get('bbox', {}) | |
| # Ensure we return the zone with proper bbox coordinates | |
| if isinstance(zone_bbox, dict): | |
| return { | |
| "zone": { | |
| "text": zone.get('text', ''), | |
| "confidence": zone.get('confidence', 0), | |
| "bbox": zone_bbox, | |
| "x": zone_bbox.get('x1', 0), | |
| "y": zone_bbox.get('y1', 0), | |
| "width": zone_bbox.get('width', 0), | |
| "height": zone_bbox.get('height', 0), | |
| "tolerance_info": zone.get('tolerance_info', {}) | |
| }, | |
| "message": f"Found text along baseline (distance: {closest['distance']:.1f}px)", | |
| "baseline": {"x1": x1, "y1": y1, "x2": x2, "y2": y2}, | |
| "found_zone": zone.get('text', ''), | |
| "fitted": True | |
| } | |
| else: | |
| return { | |
| "zone": zone, | |
| "message": f"Found text along baseline (distance: {closest['distance']:.1f}px)", | |
| "baseline": {"x1": x1, "y1": y1, "x2": x2, "y2": y2}, | |
| "found_zone": zone.get('text', ''), | |
| "fitted": True | |
| } | |
| else: | |
| return { | |
| "zones": [], | |
| "message": "No text found along baseline", | |
| "baseline": {"x1": x1, "y1": y1, "x2": x2, "y2": y2} | |
| } | |
| finally: | |
| # Clean up temporary file | |
| if os.path.exists(tmp_path): | |
| os.unlink(tmp_path) | |
| except Exception as e: | |
| logger.error(f"Error processing baseline: {e}") | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| def intersects_with_line(bbox, x1, y1, x2, y2): | |
| """Check if bounding box intersects with line""" | |
| bx1, by1 = bbox.get('x1', 0), bbox.get('y1', 0) | |
| bx2, by2 = bbox.get('x2', 0), bbox.get('y2', 0) | |
| # Check if any corner of the box is close to the line | |
| corners = [(bx1, by1), (bx2, by1), (bx2, by2), (bx1, by2)] | |
| for cx, cy in corners: | |
| distance = distance_point_to_line(cx, cy, x1, y1, x2, y2) | |
| if distance < 20: # Within 20 pixels | |
| return True | |
| return False | |
| def distance_point_to_line(px, py, x1, y1, x2, y2): | |
| """Calculate distance from point to line""" | |
| # Line vector | |
| line_dx = x2 - x1 | |
| line_dy = y2 - y1 | |
| # Point vector | |
| point_dx = px - x1 | |
| point_dy = py - y1 | |
| # Calculate distance using cross product | |
| if line_dx == 0 and line_dy == 0: | |
| # Line is a point | |
| return ((px - x1) ** 2 + (py - y1) ** 2) ** 0.5 | |
| # Distance = |cross_product| / |line_vector| | |
| cross_product = abs(point_dx * line_dy - point_dy * line_dx) | |
| line_length = (line_dx ** 2 + line_dy ** 2) ** 0.5 | |
| return cross_product / line_length | |
| def detect_dimension_lines(image_path): | |
| """Detect lines in the image that could be dimension lines""" | |
| import cv2 | |
| import numpy as np | |
| # Read image | |
| img = cv2.imread(image_path) | |
| if img is None: | |
| return [] | |
| # Convert to grayscale | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| # Apply edge detection | |
| edges = cv2.Canny(gray, 50, 150, apertureSize=3) | |
| # Detect lines using HoughLinesP | |
| lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=50, minLineLength=30, maxLineGap=10) | |
| detected_lines = [] | |
| if lines is not None: | |
| for line in lines: | |
| x1, y1, x2, y2 = line[0] | |
| length = ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5 | |
| # Filter out very short lines | |
| if length > 20: | |
| # Calculate angle | |
| angle = math.atan2(y2 - y1, x2 - x1) * 180 / math.pi | |
| if angle < 0: | |
| angle += 180 | |
| detected_lines.append({ | |
| 'x1': float(x1), | |
| 'y1': float(y1), | |
| 'x2': float(x2), | |
| 'y2': float(y2), | |
| 'length': float(length), | |
| 'angle': float(angle) | |
| }) | |
| return detected_lines | |
| def find_lines_near_dimensions(zones, lines, max_distance=50): | |
| """Find lines that are close to dimension text""" | |
| dimension_lines = [] | |
| for zone in zones: | |
| bbox = zone.get('bbox', {}) | |
| if not bbox: | |
| continue | |
| # Get zone center | |
| zone_center_x = (bbox.get('x1', 0) + bbox.get('x2', 0)) / 2 | |
| zone_center_y = (bbox.get('y1', 0) + bbox.get('y2', 0)) / 2 | |
| # Find closest lines | |
| nearby_lines = [] | |
| for line in lines: | |
| distance = distance_point_to_line(zone_center_x, zone_center_y, | |
| line['x1'], line['y1'], line['x2'], line['y2']) | |
| if distance < max_distance: | |
| nearby_lines.append({ | |
| 'line': line, | |
| 'distance': distance | |
| }) | |
| # Sort by distance | |
| nearby_lines.sort(key=lambda x: x['distance']) | |
| if nearby_lines: | |
| dimension_lines.append({ | |
| 'zone': zone, | |
| 'closest_line': nearby_lines[0]['line'], | |
| 'distance': nearby_lines[0]['distance'], | |
| 'all_nearby_lines': [nl['line'] for nl in nearby_lines[:3]] # Top 3 closest | |
| }) | |
| return dimension_lines | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=8000) |