Spaces:
Runtime error
Runtime error
File size: 3,442 Bytes
b9e2109 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | import logging
import numpy as np
from typing import List, Dict, Any
from modules.detectors.base_detector import ObjectDetector
from modules.config import UI_MODEL_PATH, PYTORCH_ENABLE_MPS_FALLBACK
logger = logging.getLogger(__name__)
class RectangleDetector(ObjectDetector):
"""Detector for rectangular UI elements like buttons, input fields, etc."""
def __init__(self):
self.model = None
# self.ui_classes = ['button', 'card', 'field', 'heading', 'icon', 'image', 'link', 'paragraph', 'text']
self.ui_classes = ['card', 'field', 'icon', 'image', 'paragraph']
async def initialize(self):
"""Initialize the YOLO model for UI element detection."""
from ultralytics import YOLO
import torch
import os
# Check for available hardware acceleration
if torch.cuda.is_available():
self.device = "cuda:0"
logger.info("CUDA is available, using GPU acceleration for UI detection")
elif hasattr(torch, 'mps') and torch.mps.is_available():
self.device = "mps"
logger.info("MPS is available, using Apple Silicon acceleration for UI detection")
# Enable MPS fallback for operations not supported by MPS
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = PYTORCH_ENABLE_MPS_FALLBACK
else:
self.device = "cpu"
logger.warning("No hardware acceleration available, falling back to CPU for UI detection")
# Initialize UI element detection model with custom-trained weights
model_path = UI_MODEL_PATH
if os.path.exists(model_path):
self.model = YOLO(model_path)
logger.info(f"UI element detector initialized from {model_path} on {self.device}")
else:
logger.error(f"UI model file {model_path} not found")
raise FileNotFoundError(f"UI model file {model_path} not found")
def detect(self, image: np.ndarray) -> List[Dict]:
"""Detect UI elements in an image."""
if self.model is None:
raise RuntimeError("UI detection model not initialized")
# Ensure image has 3 channels (RGB)
if len(image.shape) == 3 and image.shape[2] == 4: # If RGBA, convert to RGB
image = image[:, :, :3]
# Run inference
results = self.model(image, conf=0.25, device=self.device)
# Process and format results
detected_elements = []
for result in results:
boxes = result.boxes
for box in boxes:
# Get bounding box coordinates
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
# Get class name
cls_id = int(box.cls[0].item())
cls_name = result.names[cls_id]
confidence = float(box.conf[0].item())
# Only include UI elements
# Note: This is a placeholder - with a proper UI element model,
# this filtering wouldn't be necessary
if cls_name.lower() in self.ui_classes:
detected_elements.append({
"label": cls_name.lower(),
"confidence": confidence,
"bbox": [x1, y1, x2, y2]
})
return detected_elements |