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