app_hf_spaces_server.py

#!/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)
    )