ui_element_locator.py

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}")