smart_ocr.py

import easyocr
import pyautogui
import numpy as np
from PIL import ImageGrab
import cv2
from scipy.cluster.hierarchy import linkage, fcluster
from scipy.spatial.distance import pdist, squareform

class SmartDesktopController:
    def __init__(self):
        self.reader = easyocr.Reader(['tr', 'en'], gpu=False)
        self.last_results = []
    
    def get_screen_elements(self):
        """Get all texts and their locations on screen - advanced grouping"""
        screenshot = ImageGrab.grab()
        img_np = np.array(screenshot)
        
        # Detect texts with EasyOCR
        results = self.reader.readtext(img_np)
        
        elements = []
        for (bbox, text, confidence) in results:
            if confidence > 0.4:  # Lower threshold
                x_coords = [point[0] for point in bbox]
                y_coords = [point[1] for point in bbox]
                
                # Bbox info
                min_x, max_x = min(x_coords), max(x_coords)
                min_y, max_y = min(y_coords), max(y_coords)
                width = max_x - min_x
                height = max_y - min_y
                
                center_x = int(sum(x_coords) / len(x_coords))
                center_y = int(sum(y_coords) / len(y_coords))
                
                elements.append({
                    'text': text.lower().strip(),
                    'original_text': text,
                    'center': (center_x, center_y),
                    'bbox': bbox,
                    'min_x': min_x,
                    'max_x': max_x,
                    'min_y': min_y,
                    'max_y': max_y,
                    'width': width,
                    'height': height,
                    'confidence': confidence
                })
        
        if len(elements) > 0:
            elements = self._smart_grouping(elements)
        
        self.last_results = elements
        return elements
    
    def _smart_grouping(self, elements):
        """Smart grouping - hierarchical + geometric rules"""
        if len(elements) < 2:
            return elements
        
        # Custom distance function - consider both x and y
        coords = []
        for e in elements:
            # Increase Y weight (vertical alignment matters)
            coords.append([e['center'][0], e['center'][1] * 2])
        
        coords = np.array(coords)
        
        # Hierarchical clustering
        linkage_matrix = linkage(coords, method='ward')
        
        # Dynamic threshold - based on screen height
        screen_height = ImageGrab.grab().height
        distance_threshold = screen_height * 0.05  # 5% of the screen
        
        labels = fcluster(linkage_matrix, distance_threshold, criterion='distance')
        
        # Build groups
        groups = {}
        for idx, label in enumerate(labels):
            if label not in groups:
                groups[label] = []
            groups[label].append(elements[idx])
        
        # Analyze and merge each group
        grouped_elements = []
        
        for group_items in groups.values():
            if len(group_items) == 1:
                grouped_elements.append(group_items[0])
                continue
            
            # Geometric analysis
            group_items = self._analyze_group_geometry(group_items)
            
            if self._should_merge_group(group_items):
                # Merge
                merged = self._merge_group(group_items)
                grouped_elements.append(merged)
            else:
                # Keep separate
                grouped_elements.extend(group_items)
        
        return grouped_elements
    
    def _analyze_group_geometry(self, group_items):
        """Analyze group geometry"""
        # Sort by Y coordinate
        sorted_items = sorted(group_items, key=lambda x: x['center'][1])
        
        # Analyze relation to next element
        for i, item in enumerate(sorted_items):
            item['y_order'] = i
            
            # Compute distance to next element
            if i < len(sorted_items) - 1:
                next_item = sorted_items[i + 1]
                y_distance = next_item['min_y'] - item['max_y']
                x_overlap = self._calculate_x_overlap(item, next_item)
                
                item['next_y_distance'] = y_distance
                item['next_x_overlap'] = x_overlap
        
        return sorted_items
    
    def _calculate_x_overlap(self, item1, item2):
        """Calculate X-axis overlap ratio between two elements"""
        overlap_start = max(item1['min_x'], item2['min_x'])
        overlap_end = min(item1['max_x'], item2['max_x'])
        
        if overlap_end > overlap_start:
            overlap = overlap_end - overlap_start
            min_width = min(item1['width'], item2['width'])
            return overlap / min_width if min_width > 0 else 0
        return 0
    
    def _should_merge_group(self, group_items):
        """Should we merge the group?"""
        if len(group_items) < 2:
            return False
        
        # Conditions:
        # 1. Small vertical distance (<50px)
        # 2. High X-axis overlap (>0.5)
        # 3. Similar width
        
        for i in range(len(group_items) - 1):
            item = group_items[i]
            
            y_dist = item.get('next_y_distance', 999)
            x_overlap = item.get('next_x_overlap', 0)
            
            # If too much space between rows, don't merge
            if y_dist > 50:
                return False
            
            # If too little X overlap, don't merge
            if x_overlap < 0.3:
                return False
        
        return True
    
    def _merge_group(self, group_items):
        """Merge a group"""
        # Sorted by Y coordinate
        sorted_items = sorted(group_items, key=lambda x: x['center'][1])
        
        # Merge texts (line by line)
        lines = []
        current_line = [sorted_items[0]]
        
        for i in range(1, len(sorted_items)):
            prev = sorted_items[i-1]
            curr = sorted_items[i]
            
            y_distance = curr['min_y'] - prev['max_y']
            
            # Same line?
            if y_distance < prev['height'] * 0.5:
                current_line.append(curr)
            else:
                # New line
                lines.append(current_line)
                current_line = [curr]
        
        lines.append(current_line)
        
        # Sort each line by X coordinate
        combined_text_parts = []
        for line in lines:
            line_sorted = sorted(line, key=lambda x: x['center'][0])
            line_text = ' '.join([item['original_text'] for item in line_sorted])
            combined_text_parts.append(line_text)
        
        # Join lines
        combined_text = '\n'.join(combined_text_parts)
        combined_text_lower = combined_text.lower()
        
        # Average position
        avg_x = int(np.mean([item['center'][0] for item in sorted_items]))
        avg_y = int(np.mean([item['center'][1] for item in sorted_items]))
        
        # Highest confidence
        max_confidence = max([item['confidence'] for item in sorted_items])
        
        return {
            'text': combined_text_lower,
            'original_text': combined_text,
            'center': (avg_x, avg_y),
            'bbox': sorted_items[0]['bbox'],
            'confidence': max_confidence,
            'is_grouped': True,
            'group_size': len(sorted_items),
            'min_x': min([item['min_x'] for item in sorted_items]),
            'max_x': max([item['max_x'] for item in sorted_items]),
            'min_y': min([item['min_y'] for item in sorted_items]),
            'max_y': max([item['max_y'] for item in sorted_items]),
            'width': max([item['max_x'] for item in sorted_items]) - min([item['min_x'] for item in sorted_items]),
            'height': max([item['max_y'] for item in sorted_items]) - min([item['min_y'] for item in sorted_items])
        }
    
    def click_on_text(self, target_text):
        """Click on specified text - with fuzzy matching"""
        elements = self.get_screen_elements()
        target = target_text.lower().strip()
        
        # 1. Exact match
        for element in elements:
            if target == element['text']:
                x, y = element['center']
                pyautogui.click(x, y)
                print(f"✓ '{element['original_text']}' clicked ({x}, {y})")
                return True
        
        # 2. Substring match
        for element in elements:
            if target in element['text'] or element['text'] in target:
                x, y = element['center']
                pyautogui.click(x, y)
                print(f"✓ '{element['original_text']}' clicked ({x}, {y})")
                return True
        
        # 3. Fuzzy matching (similarity score)
        from difflib import SequenceMatcher
        
        best_match = None
        best_score = 0.6  # Minimum similarity
        
        for element in elements:
            score = SequenceMatcher(None, target, element['text']).ratio()
            if score > best_score:
                best_score = score
                best_match = element
        
        if best_match:
            x, y = best_match['center']
            pyautogui.click(x, y)
            print(f"✓ '{best_match['original_text']}' clicked (similarity: {best_score:.2f}) ({x}, {y})")
            return True
        
        # Not found
        print(f"✗ '{target_text}' not found")
        print(f"On-screen items:")
        for e in elements[:15]:
            grouped_mark = " [GROUP]" if e.get('is_grouped') else ""
            print(f"  - {e['original_text']}{grouped_mark}")
        return False
    
    def show_detected_elements(self):
        """Visualize detected elements"""
        try:
            screenshot = ImageGrab.grab()
            img = np.array(screenshot)
            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
            
            for element in self.last_results:
                # Convert coordinates to Python int
                min_x = int(element['min_x'])
                min_y = int(element['min_y'])
                max_x = int(element['max_x'])
                max_y = int(element['max_y'])
                
                # Different color for grouped elements
                if element.get('is_grouped'):
                    # Rectangle covering the whole group
                    cv2.rectangle(img, 
                                (min_x, min_y),
                                (max_x, max_y),
                                (0, 255, 255), 3)
                else:
                    bbox = element['bbox']
                    pts = np.array(bbox, dtype=np.int32)
                    cv2.polylines(img, [pts], True, (0, 255, 0), 2)
                
                x = int(element['center'][0])
                y = int(element['center'][1])
                cv2.circle(img, (x, y), 5, (255, 0, 0), -1)
                
                # Text label
                label = element['original_text'].replace('\n', ' ')[:30]
                if element.get('is_grouped'):
                    label += f" [{element['group_size']}]"
                
                # Improve label readability with background
                (text_w, text_h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
                cv2.rectangle(img, (x-5, y-text_h-15), (x+text_w+5, y-5), (0, 0, 0), -1)
                cv2.putText(img, label, (x, y-10), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
            
            cv2.imwrite('detected_elements.png', img)
            print("✓ Detected elements saved to 'detected_elements.png'")
        except Exception as e:
            print(f"❌ Visualization error: {e}")
            import traceback
            traceback.print_exc()

# Usage
if __name__ == '__main__':
    controller = SmartDesktopController()
    elements = controller.get_screen_elements()
    
    print(f"\n{'='*60}")
    print(f"Total of {len(elements)} elements found")
    print(f"{'='*60}\n")
    
    for i, element in enumerate(elements, 1):
        grouped = " [GROUPED]" if element.get('is_grouped') else ""
        print(f"{i}. 📌 {element['original_text']}{grouped}")
        print(f"   Position: {element['center']}")
        print(f"   Confidence: {element['confidence']:.2f}")
        if element.get('is_grouped'):
            print(f"   Size: {element['width']}x{element['height']}px")
        print()
    
    controller.show_detected_elements()