Upload app_hf_spaces_server.py with huggingface_hub

app_hf_spaces_server.py

@@ -0,0 +1,773 @@
+#!/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 = """
+<!DOCTYPE html>
+<html>
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>OmniParser UI Detector</title>
+    <style>
+        * {
+            margin: 0;
+            padding: 0;
+            box-sizing: border-box;
+        }
+        
+        body {
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, 'Helvetica Neue', Arial, sans-serif;
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            min-height: 100vh;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            padding: 20px;
+        }
+        
+        .container {
+            background: white;
+            border-radius: 12px;
+            box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
+            max-width: 1000px;
+            width: 100%;
+            padding: 40px;
+        }
+        
+        h1 {
+            color: #333;
+            margin-bottom: 10px;
+            display: flex;
+            align-items: center;
+            gap: 10px;
+        }
+        
+        .subtitle {
+            color: #666;
+            margin-bottom: 30px;
+            font-size: 14px;
+        }
+        
+        .upload-area {
+            border: 2px dashed #667eea;
+            border-radius: 8px;
+            padding: 40px;
+            text-align: center;
+            cursor: pointer;
+            transition: all 0.3s;
+            margin-bottom: 20px;
+        }
+        
+        .upload-area:hover {
+            border-color: #764ba2;
+            background: #f8f9ff;
+        }
+        
+        .upload-area.dragover {
+            border-color: #764ba2;
+            background: #f0f2ff;
+        }
+        
+        input[type="file"] {
+            display: none;
+        }
+        
+        .upload-text {
+            font-size: 16px;
+            color: #667eea;
+            margin-bottom: 10px;
+        }
+        
+        .upload-hint {
+            font-size: 12px;
+            color: #999;
+        }
+        
+        button {
+            background: #667eea;
+            color: white;
+            border: none;
+            padding: 12px 30px;
+            border-radius: 6px;
+            cursor: pointer;
+            font-size: 14px;
+            font-weight: 600;
+            transition: background 0.3s;
+        }
+        
+        button:hover {
+            background: #764ba2;
+        }
+        
+        button:disabled {
+            background: #ccc;
+            cursor: not-allowed;
+        }
+        
+        .results {
+            display: none;
+            margin-top: 30px;
+        }
+        
+        .results.active {
+            display: block;
+        }
+        
+        .result-section {
+            margin-bottom: 25px;
+        }
+        
+        .result-section h3 {
+            color: #333;
+            margin-bottom: 10px;
+            font-size: 14px;
+            text-transform: uppercase;
+            letter-spacing: 1px;
+        }
+        
+        .preview-image {
+            max-width: 100%;
+            border-radius: 6px;
+            margin-bottom: 15px;
+        }
+        
+        .stats {
+            display: grid;
+            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
+            gap: 15px;
+            margin-bottom: 20px;
+        }
+        
+        .stat {
+            background: #f8f9ff;
+            padding: 15px;
+            border-radius: 6px;
+            border-left: 4px solid #667eea;
+        }
+        
+        .stat-value {
+            font-size: 24px;
+            font-weight: bold;
+            color: #667eea;
+        }
+        
+        .stat-label {
+            font-size: 12px;
+            color: #999;
+            margin-top: 5px;
+        }
+        
+        textarea {
+            width: 100%;
+            min-height: 200px;
+            padding: 12px;
+            border: 1px solid #ddd;
+            border-radius: 6px;
+            font-family: 'Monaco', 'Courier New', monospace;
+            font-size: 12px;
+            resize: vertical;
+        }
+        
+        .loading {
+            display: none;
+            text-align: center;
+            color: #667eea;
+        }
+        
+        .loading.active {
+            display: block;
+        }
+        
+        .spinner {
+            border: 3px solid #f3f3f3;
+            border-top: 3px solid #667eea;
+            border-radius: 50%;
+            width: 40px;
+            height: 40px;
+            animation: spin 1s linear infinite;
+            margin: 20px auto;
+        }
+        
+        @keyframes spin {
+            0% { transform: rotate(0deg); }
+            100% { transform: rotate(360deg); }
+        }
+        
+        .error {
+            background: #fee;
+            color: #c33;
+            padding: 12px;
+            border-radius: 6px;
+            margin-bottom: 20px;
+            display: none;
+        }
+        
+        .error.active {
+            display: block;
+        }
+        
+        .download-buttons {
+            display: flex;
+            gap: 10px;
+            margin-top: 15px;
+            flex-wrap: wrap;
+        }
+        
+        .download-btn {
+            background: #28a745;
+            font-size: 13px;
+            padding: 10px 20px;
+        }
+        
+        .download-btn:hover {
+            background: #218838;
+        }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <h1>🎯 OmniParser UI Detector</h1>
+        <p class="subtitle">Upload a UI screenshot to detect and locate all UI elements</p>
+        
+        <div class="upload-area" id="uploadArea">
+            <input type="file" id="fileInput" accept="image/*">
+            <div class="upload-text">Click to upload or drag and drop</div>
+            <div class="upload-hint">PNG, JPG (max 10MB)</div>
+        </div>
+        
+        <button id="analyzeBtn" disabled>🔍 Analyze Image</button>
+        
+        <div class="error" id="errorDiv"></div>
+        
+        <div class="loading" id="loadingDiv">
+            <div class="spinner"></div>
+            <p>Processing image... (this may take 30-60 seconds)</p>
+        </div>
+        
+        <div class="results" id="results">
+            <div class="result-section">
+                <h3>📊 Statistics</h3>
+                <div class="stats">
+                    <div class="stat">
+                        <div class="stat-value" id="elementCount">0</div>
+                        <div class="stat-label">Elements Detected</div>
+                    </div>
+                    <div class="stat">
+                        <div class="stat-value" id="processingTime">0s</div>
+                        <div class="stat-label">Processing Time</div>
+                    </div>
+                </div>
+            </div>
+            
+            <div class="result-section">
+                <h3>🖼️ Visualization</h3>
+                <img id="vizImage" class="preview-image" src="">
+            </div>
+            
+            <div class="result-section">
+                <h3>📋 Coordinates (JSON)</h3>
+                <textarea id="jsonOutput" readonly></textarea>
+                <div class="download-buttons">
+                    <button class="download-btn" onclick="downloadJSON()">⬇️ Download JSON</button>
+                </div>
+            </div>
+            
+            <div class="result-section">
+                <h3>📈 Coordinates (CSV)</h3>
+                <textarea id="csvOutput" readonly></textarea>
+                <div class="download-buttons">
+                    <button class="download-btn" onclick="downloadCSV()">⬇️ Download CSV</button>
+                </div>
+            </div>
+        </div>
+    </div>
+    
+    <script>
+        const uploadArea = document.getElementById('uploadArea');
+        const fileInput = document.getElementById('fileInput');
+        const analyzeBtn = document.getElementById('analyzeBtn');
+        const results = document.getElementById('results');
+        const loadingDiv = document.getElementById('loadingDiv');
+        const errorDiv = document.getElementById('errorDiv');
+        
+        uploadArea.addEventListener('click', () => fileInput.click());
+        
+        uploadArea.addEventListener('dragover', (e) => {
+            e.preventDefault();
+            uploadArea.classList.add('dragover');
+        });
+        
+        uploadArea.addEventListener('dragleave', () => {
+            uploadArea.classList.remove('dragover');
+        });
+        
+        uploadArea.addEventListener('drop', (e) => {
+            e.preventDefault();
+            uploadArea.classList.remove('dragover');
+            if (e.dataTransfer.files.length) {
+                fileInput.files = e.dataTransfer.files;
+                analyzeBtn.disabled = false;
+            }
+        });
+        
+        fileInput.addEventListener('change', () => {
+            if (fileInput.files.length) {
+                analyzeBtn.disabled = false;
+            }
+        });
+        
+        analyzeBtn.addEventListener('click', async () => {
+            if (!fileInput.files.length) return;
+            
+            const file = fileInput.files[0];
+            const formData = new FormData();
+            formData.append('file', file);
+            
+            errorDiv.classList.remove('active');
+            results.classList.remove('active');
+            loadingDiv.classList.add('active');
+            analyzeBtn.disabled = true;
+            
+            try {
+                const response = await fetch('/api/analyze', {
+                    method: 'POST',
+                    body: formData
+                });
+                
+                if (!response.ok) {
+                    throw new Error(`HTTP ${response.status}: ${await response.text()}`);
+                }
+                
+                const data = await response.json();
+                displayResults(data);
+            } catch (error) {
+                showError(`Error: ${error.message}`);
+            } finally {
+                loadingDiv.classList.remove('active');
+                analyzeBtn.disabled = false;
+            }
+        });
+        
+        function displayResults(data) {
+            document.getElementById('elementCount').textContent = data.analysis.total_elements_detected;
+            document.getElementById('processingTime').textContent = data.processing_time_seconds.toFixed(2) + 's';
+            document.getElementById('vizImage').src = 'data:image/png;base64,' + data.exports.visualization_png_base64;
+            document.getElementById('jsonOutput').value = JSON.stringify(data.analysis.elements, null, 2);
+            document.getElementById('csvOutput').value = data.exports.csv_data;
+            
+            results.classList.add('active');
+        }
+        
+        function showError(msg) {
+            errorDiv.textContent = msg;
+            errorDiv.classList.add('active');
+        }
+        
+        function downloadJSON() {
+            const json = document.getElementById('jsonOutput').value;
+            const blob = new Blob([json], { type: 'application/json' });
+            const url = URL.createObjectURL(blob);
+            const a = document.createElement('a');
+            a.href = url;
+            a.download = 'coordinates.json';
+            a.click();
+        }
+        
+        function downloadCSV() {
+            const csv = document.getElementById('csvOutput').value;
+            const blob = new Blob([csv], { type: 'text/csv' });
+            const url = URL.createObjectURL(blob);
+            const a = document.createElement('a');
+            a.href = url;
+            a.download = 'coordinates.csv';
+            a.click();
+        }
+    </script>
+</body>
+</html>
+"""
+
+# ============ 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)
+    )