Upload ui_element_locator.py with huggingface_hub

ui_element_locator.py

@@ -0,0 +1,256 @@
+import cv2
+import numpy as np
+import json
+import os
+from pathlib import Path
+from typing import Dict, Any, Tuple, Optional, List
+from PIL import Image
+
+def to_rgb(img: np.ndarray) -> Optional[np.ndarray]:
+    """Converts image to BGR format (3 channels). Handles None input."""
+    if img is None:
+        return None
+    if len(img.shape) == 2:
+        # Grayscale to BGR
+        return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+    if img.shape[2] == 4:
+        # BGRA to BGR (removes alpha channel)
+        return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
+    # Already BGR or RGB (assuming OpenCV reads as BGR)
+    return img
+
+def match_ui_elements_in_image(
+    original_image_path: str,
+    cropped_templates_dir: str = 'cropped_images',
+    threshold: float = 0.7,
+    output_json: Optional[str] = None
+) -> Dict[str, Any]:
+    """
+    Matches cropped UI element templates against an original image.
+    Returns coordinates of all matched elements.
+    
+    Args:
+        original_image_path: Path to the original image (e.g., Screenshot.png)
+        cropped_templates_dir: Directory containing cropped UI images
+        threshold: Confidence threshold for matches (0-1)
+        output_json: Optional path to save results as JSON
+    
+    Returns:
+        Dictionary with match results
+    """
+    
+    print(f"[UI Locator] Loading original image: {original_image_path}")
+    original_img = cv2.imread(original_image_path, cv2.IMREAD_UNCHANGED)
+    original_img_rgb = to_rgb(original_img)
+    
+    if original_img_rgb is None:
+        raise ValueError(f"Failed to load image: {original_image_path}")
+    
+    print(f"[UI Locator] Original image size: {original_img_rgb.shape}")
+    img_height, img_width = original_img_rgb.shape[:2]
+    
+    # Load all cropped templates
+    print(f"[UI Locator] Loading cropped templates from: {cropped_templates_dir}")
+    templates = {}
+    template_files = sorted(Path(cropped_templates_dir).glob('crop_*.png'))
+    
+    if not template_files:
+        raise ValueError(f"No cropped templates found in {cropped_templates_dir}")
+    
+    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
+        else:
+            print(f"[WARNING] Could not load template: {template_file.name}")
+    
+    print(f"[UI Locator] Loaded {len(templates)} templates")
+    
+    # Match each template
+    matches = []
+    skipped = 0
+    
+    for i, (template_name, template_img) in enumerate(templates.items()):
+        try:
+            # Skip templates that are too large
+            if template_img.shape[0] > img_height or template_img.shape[1] > img_width:
+                skipped += 1
+                continue
+            
+            # Skip very small templates (likely noise)
+            if template_img.shape[0] < 4 or template_img.shape[1] < 4:
+                skipped += 1
+                continue
+            
+            # Perform template matching
+            result = cv2.matchTemplate(original_img_rgb, template_img, cv2.TM_CCOEFF_NORMED)
+            _, max_val, _, max_loc = cv2.minMaxLoc(result)
+            
+            # Only record matches above threshold
+            if max_val >= threshold:
+                template_h, template_w = template_img.shape[:2]
+                x1, y1 = max_loc
+                x2 = x1 + template_w
+                y2 = y1 + template_h
+                
+                # Calculate center
+                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)
+                    },
+                    'bbox_ratio': {
+                        'x1': x1 / img_width,
+                        'y1': y1 / img_height,
+                        'x2': x2 / img_width,
+                        'y2': y2 / img_height
+                    }
+                })
+                
+                if (i + 1) % 20 == 0:
+                    print(f"[UI Locator] Processed {i + 1}/{len(templates)} templates...")
+        
+        except Exception as e:
+            print(f"[WARNING] Failed to match {template_name}: {str(e)}")
+            skipped += 1
+            continue
+    
+    # Sort matches by confidence
+    matches.sort(key=lambda x: x['confidence'], reverse=True)
+    
+    result = {
+        'source_image': original_image_path,
+        'image_size': {
+            'width': img_width,
+            'height': img_height
+        },
+        'templates_directory': cropped_templates_dir,
+        'templates_loaded': len(templates),
+        'templates_skipped': skipped,
+        'threshold': threshold,
+        'matches_found': len(matches),
+        'matches': matches
+    }
+    
+    print(f"\n[UI Locator] Matching complete!")
+    print(f"[UI Locator] Found {len(matches)} matches above threshold {threshold}")
+    print(f"[UI Locator] Skipped {skipped} templates (too large or too small)")
+    
+    # Save results as JSON
+    if output_json:
+        with open(output_json, 'w') as f:
+            json.dump(result, f, indent=2)
+        print(f"[UI Locator] Results saved to: {output_json}")
+    
+    return result
+
+def visualize_matches(
+    original_image_path: str,
+    matches_data: Dict[str, Any],
+    output_image_path: Optional[str] = None
+) -> np.ndarray:
+    """
+    Visualize matched UI elements on the original image.
+    
+    Args:
+        original_image_path: Path to original image
+        matches_data: Results from match_ui_elements_in_image
+        output_image_path: Optional path to save visualization
+    
+    Returns:
+        Annotated image with bounding boxes
+    """
+    
+    print(f"[Visualization] Loading image: {original_image_path}")
+    img = cv2.imread(original_image_path)
+    
+    if img is None:
+        raise ValueError(f"Failed to load visualization image: {original_image_path}")
+    
+    # Draw bounding boxes for each match
+    for match in matches_data['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 center point
+        
+        # Draw label
+        label = f"ID:{template_id} ({confidence:.2f})"
+        cv2.putText(img, label, (bbox['x1'], bbox['y1'] - 5), 
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 0, 0), 1)
+    
+    if output_image_path:
+        cv2.imwrite(output_image_path, img)
+        print(f"[Visualization] Saved to: {output_image_path}")
+    
+    return img
+
+if __name__ == "__main__":
+    import sys
+    from pathlib import Path
+    from config import get_screenshot_path, get_output_path, CROPPED_IMAGES_DIR
+    
+    # Paths using config
+    original_image = get_screenshot_path('Screenshot.png')
+    cropped_dir = str(CROPPED_IMAGES_DIR)
+    output_json = get_output_path('ui_elements_coordinates.json')
+    output_viz = get_output_path('ui_elements_visualization.png')
+    
+    print("=" * 70)
+    print("UI Element Locator - Template Matching Tool")
+    print("=" * 70)
+    
+    # Match UI elements
+    results = match_ui_elements_in_image(
+        original_image_path=original_image,
+        cropped_templates_dir=cropped_dir,
+        threshold=0.7,
+        output_json=output_json
+    )
+    
+    print(f"\n[Summary]")
+    print(f"Total UI elements found: {results['matches_found']}")
+    print(f"Image size: {results['image_size']['width']}x{results['image_size']['height']}")
+    
+    # Show top matches
+    print(f"\n[Top 10 Matches by Confidence]")
+    print("-" * 70)
+    for i, match in enumerate(results['matches'][:10], 1):
+        bbox = match['bbox']
+        center = match['center']
+        print(f"{i}. {match['template_id']} - Confidence: {match['confidence']:.4f}")
+        print(f"   Center: ({center['x']}, {center['y']}) | Bbox: ({bbox['x1']}, {bbox['y1']}) -> ({bbox['x2']}, {bbox['y2']})")
+    
+    # Visualize matches
+    try:
+        print(f"\n[Visualization] Creating annotated image...")
+        visualize_matches(original_image, results, output_viz)
+    except Exception as e:
+        print(f"[ERROR] Visualization failed: {str(e)}")
+    
+    print(f"\n[Output Files]")
+    print(f"JSON Results: {output_json}")
+    print(f"Visualization: {output_viz}")