#!/usr/bin/env python3 """ OmniParser UI Element Detection - FastAPI Server for HF Spaces Full REST API + Web UI served on port 7860 Features: - /api/analyze - POST image for UI element detection - /api/health - Health check - / - HTML web interface - Automatic model initialization on startup - CORS enabled for cross-origin requests """ import os import sys import json import time import base64 import cv2 import numpy as np import io import csv from pathlib import Path from typing import Dict, Any, Optional, Tuple, List from contextlib import asynccontextmanager import threading from fastapi import FastAPI, File, UploadFile, HTTPException, Request from fastapi.responses import JSONResponse, HTMLResponse, FileResponse from fastapi.staticfiles import StaticFiles from fastapi.middleware.cors import CORSMiddleware import uvicorn from PIL import Image # Configure OmniParser os.environ["OMP_NUM_THREADS"] = "4" # Add OmniParser to path dynamically omoi_root = Path(__file__).parent sys.path.insert(0, str(omoi_root / 'OmniParser')) from util.omniparser import Omniparser from config import get_omniparser_config # ============ Utility Functions ============ def to_rgb(img: np.ndarray) -> Optional[np.ndarray]: """Convert image to BGR format.""" if img is None: return None if len(img.shape) == 2: return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) if img.shape[2] == 4: return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR) return img def extract_coordinates(matches: List[Dict]) -> List[Dict]: """Extract coordinates from matches for JSON export.""" coords = [] for i, match in enumerate(matches, 1): bbox = match['bbox'] center = match['center'] coords.append({ 'element_id': f"crop_{i:04d}", 'x': center['x'], 'y': center['y'], 'x1': bbox['x1'], 'y1': bbox['y1'], 'x2': bbox['x2'], 'y2': bbox['y2'], 'width': bbox['width'], 'height': bbox['height'], 'confidence': match['confidence'], 'template_file': match['template_file'] }) return coords def match_ui_elements( original_image_array: np.ndarray, cropped_images_dir: str, threshold: float = 0.7 ) -> Tuple[list, Dict]: """Match cropped UI templates against original image.""" original_img_rgb = to_rgb(original_image_array) if original_img_rgb is None: raise ValueError("Failed to convert original image") img_height, img_width = original_img_rgb.shape[:2] # Load templates templates = {} template_files = sorted(Path(cropped_images_dir).glob('crop_*.png')) for template_file in template_files: template_img = cv2.imread(str(template_file), cv2.IMREAD_UNCHANGED) if template_img is not None: template_img_rgb = to_rgb(template_img) templates[template_file.name] = template_img_rgb # Match templates matches = [] for template_name, template_img in templates.items(): try: if template_img.shape[0] > img_height or template_img.shape[1] > img_width: continue if template_img.shape[0] < 4 or template_img.shape[1] < 4: continue result = cv2.matchTemplate(original_img_rgb, template_img, cv2.TM_CCOEFF_NORMED) _, max_val, _, max_loc = cv2.minMaxLoc(result) if max_val >= threshold: template_h, template_w = template_img.shape[:2] x1, y1 = max_loc x2 = x1 + template_w y2 = y1 + template_h center_x = (x1 + x2) / 2 center_y = (y1 + y2) / 2 matches.append({ 'template_id': template_name.replace('.png', ''), 'template_file': template_name, 'confidence': float(max_val), 'bbox': { 'x1': int(x1), 'y1': int(y1), 'x2': int(x2), 'y2': int(y2), 'width': int(template_w), 'height': int(template_h) }, 'center': { 'x': int(center_x), 'y': int(center_y) } }) except Exception: continue matches.sort(key=lambda x: x['confidence'], reverse=True) metadata = { 'image_size': {'width': img_width, 'height': img_height}, 'templates_loaded': len(templates), 'threshold': threshold, 'matches_found': len(matches) } return matches, metadata def visualize_matches( original_image_array: np.ndarray, matches: list ) -> np.ndarray: """Create visualization with bounding boxes.""" img = original_image_array.copy() for match in matches: bbox = match['bbox'] center = match['center'] confidence = match['confidence'] template_id = match['template_id'] # Draw bounding box color = (0, 255, 0) # Green thickness = 2 cv2.rectangle(img, (bbox['x1'], bbox['y1']), (bbox['x2'], bbox['y2']), color, thickness) # Draw center point cv2.circle(img, (center['x'], center['y']), 3, (0, 0, 255), -1) # Red # Draw label label = f"{template_id} ({confidence:.2f})" cv2.putText(img, label, (bbox['x1'], bbox['y1'] - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 0, 0), 1) return img def matches_to_csv(matches: list) -> str: """Convert matches to CSV format.""" output = io.StringIO() writer = csv.writer(output) writer.writerow([ 'Element_ID', 'X', 'Y', 'X1', 'Y1', 'X2', 'Y2', 'Width', 'Height', 'Confidence' ]) for i, match in enumerate(matches, 1): bbox = match['bbox'] center = match['center'] writer.writerow([ f"crop_{i:04d}", center['x'], center['y'], bbox['x1'], bbox['y1'], bbox['x2'], bbox['y2'], bbox['width'], bbox['height'], f"{match['confidence']:.4f}" ]) return output.getvalue() # ============ FastAPI Setup ============ # Global OmniParser instance omniparser = None omniparser_lock = threading.Lock() @asynccontextmanager async def lifespan(app: FastAPI): """Initialize and cleanup on server startup/shutdown.""" global omniparser print("\n" + "="*60) print("šŸš€ Initializing OmniParser...") print("="*60) try: with omniparser_lock: config = get_omniparser_config() print(f"āœ“ Config loaded from: {config['omniparser_dir']}") print(f"āœ“ Loading YOLO model...") omniparser = Omniparser(config) print(f"āœ“ OmniParser initialized successfully!") print("="*60 + "\n") except Exception as e: print(f"āœ— ERROR during initialization: {str(e)}") import traceback traceback.print_exc() print("="*60 + "\n") yield # Application runs here # Cleanup print("\n[Server] Shutting down...") # Create FastAPI app app = FastAPI( title="OmniParser UI Detection API", description="Detects and locates UI elements in screenshots", version="1.0.0", lifespan=lifespan ) # Add CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ============ Web Interface ============ HTML_UI = """ OmniParser UI Detector

šŸŽÆ OmniParser UI Detector

Upload a UI screenshot to detect and locate all UI elements

Click to upload or drag and drop
PNG, JPG (max 10MB)

Processing image... (this may take 30-60 seconds)

šŸ“Š Statistics

0
Elements Detected
0s
Processing Time

šŸ–¼ļø Visualization

šŸ“‹ Coordinates (JSON)

šŸ“ˆ Coordinates (CSV)

""" # ============ API Endpoints ============ @app.get("/") async def web_ui(): """Serve web UI.""" return HTMLResponse(content=HTML_UI) @app.get("/api/health") async def health(): """Health check endpoint.""" status = "ok" if omniparser else "initializing" return { "status": status, "service": "OmniParser UI Detection API", "mode": "HF Spaces" } @app.post("/api/analyze") async def analyze_image(file: UploadFile = File(...)): """ Analyze an image for UI elements. Returns detailed JSON with coordinates, visualization, and CSV data. """ if not omniparser: raise HTTPException(status_code=503, detail="OmniParser not initialized. Please try again in a moment.") try: print(f"\n[API] Analyzing: {file.filename}") start_time = time.time() # 1. Read image print("[Step 1] Reading image...") content = await file.read() np_array = np.frombuffer(content, np.uint8) original_img = cv2.imdecode(np_array, cv2.IMREAD_UNCHANGED) if original_img is None: raise HTTPException(status_code=400, detail="Failed to decode image") print(f"[Step 1] āœ“ Image loaded: {original_img.shape}") # 2. Encode for OmniParser print("[Step 2] Encoding for OmniParser...") _, buffer = cv2.imencode('.png', original_img) image_base64 = base64.b64encode(buffer).decode() # 3. Run OmniParser print("[Step 3] Running OmniParser...") omni_start = time.time() _, parsed_content = omniparser.parse(image_base64) omni_time = time.time() - omni_start print(f"[Step 3] āœ“ Complete in {omni_time:.2f}s") # 4. Match templates print("[Step 4] Matching templates...") cropped_dir = '/tmp/omoi_cropped_images' if not Path(cropped_dir).exists(): print(f"āš ļø Creating cropped_images cache dir...") Path(cropped_dir).mkdir(parents=True, exist_ok=True) match_start = time.time() matches, metadata = match_ui_elements(original_img, cropped_dir, threshold=0.7) match_time = time.time() - match_start print(f"[Step 4] āœ“ Found {len(matches)} elements in {match_time:.2f}s") # 5. Create visualization print("[Step 5] Creating visualization...") viz_img = visualize_matches(original_img, matches) _, viz_buffer = cv2.imencode('.png', viz_img) viz_base64 = base64.b64encode(viz_buffer).decode() # 6. Extract coordinates print("[Step 6] Extracting coordinates...") coordinates = extract_coordinates(matches) # 7. Generate CSV print("[Step 7] Generating CSV...") csv_data = matches_to_csv(matches) # 8. Prepare response total_time = time.time() - start_time print(f"[API] āœ“ Complete in {total_time:.2f}s\n") return { "status": "success", "processing_time_seconds": total_time, "timing": { "omniparser_seconds": omni_time, "template_matching_seconds": match_time }, "image_info": { "filename": file.filename, "size": metadata['image_size'] }, "analysis": { "total_elements_detected": len(coordinates), "elements": coordinates }, "exports": { "csv_data": csv_data, "visualization_png_base64": viz_base64 } } except HTTPException: raise except Exception as e: print(f"[ERROR] {str(e)}") import traceback traceback.print_exc() raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}") # ============ Main ============ if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="OmniParser FastAPI Server for HF Spaces") parser.add_argument("--port", type=int, default=7860, help="Server port (default: 7860)") parser.add_argument("--host", default="0.0.0.0", help="Server host (default: 0.0.0.0)") args = parser.parse_args() print(f"\nšŸš€ Starting server on {args.host}:{args.port}") print(f" Web UI: http://localhost:{args.port}") print(f" API Docs: http://localhost:{args.port}/docs") print(f" Analyze endpoint: POST http://localhost:{args.port}/api/analyze\n") uvicorn.run( app, host=args.host, port=args.port, loop="asyncio" # Use asyncio instead of uvloop (more compatible) )