albionbot/txt2

#!/usr/bin/env python3
import cv2
import numpy as np


def find_and_display_template(main_image_path: str, template_image_path: str) -> None:
    """
    Находит шаблонное изображение внутри основного изображения с помощью OpenCV
    и отображает результат в окне.

    Args:
        main_image_path (str): Путь к основному изображению (например, "fig.png").
        template_image_path (str): Путь к шаблонному изображению (например, "bober.png").
    """
    # Загрузка основного изображения и шаблона
    img_rgb = cv2.imread(main_image_path)
    template = cv2.imread(template_image_path)

    if img_rgb is None:
        print(f"Ошибка: Не удалось загрузить основное изображение по пути: {main_image_path}")
        return
    if template is None:
        print(f"Ошибка: Не удалось загрузить шаблонное изображение по пути: {template_image_path}")
        return

    # Преобразование изображений в оттенки серого
    img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2GRAY)
    template_gray = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)

    # Получение размеров шаблона
    w, h = template_gray.shape[::-1]

    # Выполнение сопоставления шаблонов
    # cv2.TM_CCOEFF_NORMED - один из лучших методов для этого
    res = cv2.matchTemplate(img_gray, template_gray, cv2.TM_CCOEFF_NORMED)

    # Находим максимальное значение совпадения
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
    
    # Автоматически устанавливаем порог на основе максимального значения
    # Используем 90% от максимального значения или минимум 0.5
    threshold = max(0.5, max_val * 0.9)
    print(f"Максимальная уверенность совпадения: {max_val:.2f}")
    print(f"Используемый порог: {threshold:.2f}")
    
    loc = np.where(res >= threshold)

    # Подсчет найденных совпадений
    matches = list(zip(*loc[::-1]))
    
    # Фильтруем близкие совпадения, оставляя только лучшее в каждой группе
    filtered_matches = []
    for pt in matches:
        # Проверяем, нет ли уже близкого совпадения
        is_duplicate = False
        for existing_pt, existing_conf in filtered_matches:
            # Если совпадения находятся близко друг к другу (в пределах размера шаблона)
            if abs(pt[0] - existing_pt[0]) < w and abs(pt[1] - existing_pt[1]) < h:
                # Оставляем то, у которого выше уверенность
                current_conf = res[pt[1], pt[0]]
                if current_conf > existing_conf:
                    filtered_matches.remove((existing_pt, existing_conf))
                    filtered_matches.append((pt, current_conf))
                is_duplicate = True
                break
        
        if not is_duplicate:
            confidence = res[pt[1], pt[0]]
            filtered_matches.append((pt, confidence))
    
    match_count = len(filtered_matches)
    
    if match_count > 0:
        print(f"Найдено уникальных совпадений: {match_count}")
        # Рисование прямоугольников вокруг найденных совпадений
        for pt, confidence in filtered_matches:
            cv2.rectangle(img_rgb, pt, (pt[0] + w, pt[1] + h), (0, 255, 0), 2)  # Зеленый прямоугольник толщиной 2
            print(f"  Совпадение найдено в точке: ({pt[0]}, {pt[1]}) с уверенностью: {confidence:.2f}")
    else:
        print("Совпадений не найдено даже с автоматическим порогом.")
        # Все равно покажем лучшее совпадение
        if max_val >= 0.3:  # Если есть хоть какое-то совпадение
            print(f"Показываю лучшее совпадение с уверенностью {max_val:.2f} в точке {max_loc}")
            cv2.rectangle(img_rgb, max_loc, (max_loc[0] + w, max_loc[1] + h), (0, 0, 255), 2)  # Красный прямоугольник

    # Получаем размеры изображения
    img_height, img_width = img_rgb.shape[:2]
    print(f"Размер изображения: {img_width}x{img_height}")
    
    # Масштабируем изображение, если оно слишком большое для экрана
    max_display_width = 1920
    max_display_height = 1080
    scale = 1.0
    
    if img_width > max_display_width or img_height > max_display_height:
        scale_w = max_display_width / img_width
        scale_h = max_display_height / img_height
        scale = min(scale_w, scale_h)
        new_width = int(img_width * scale)
        new_height = int(img_height * scale)
        img_display = cv2.resize(img_rgb, (new_width, new_height), interpolation=cv2.INTER_AREA)
        print(f"Изображение масштабировано до {new_width}x{new_height} (масштаб: {scale:.2f})")
    else:
        img_display = img_rgb
    
    # Сохраняем результат в файл
    output_file = "result.png"
    cv2.imwrite(output_file, img_rgb)
    print(f"Результат сохранен в файл: {output_file}")
    
    # Отображение результата в отдельном окне
    print("\nОткрываю окно с результатом...")
    cv2.namedWindow('Результат поиска шаблона', cv2.WINDOW_NORMAL)
    cv2.imshow('Результат поиска шаблона', img_display)
    print("Окно открыто. Нажмите любую клавишу в окне или 'q' для закрытия...")
    
    # Ждать, пока пользователь не нажмет любую клавишу
    key = cv2.waitKey(0) & 0xFF
    if key == ord('q') or key == 27:  # 'q' или ESC
        print("Окно закрыто пользователем")
    cv2.destroyAllWindows()  # Закрыть все окна OpenCV


if __name__ == "__main__":
    import os
    
    # Ищем bob.png на @fig.png или fig.png
    template_image = "bob.png"
    
    # Пробуем сначала @fig.png, если не найдется - используем fig.png
    if os.path.exists("@fig.png"):
        main_image = "@fig.png"
    elif os.path.exists("fig.png"):
        main_image = "fig.png"
        print("Файл @fig.png не найден, используется fig.png")
    else:
        print("Ошибка: Файлы @fig.png и fig.png не найдены!")
        exit(1)
    
    # Проверяем существование шаблона
    if not os.path.exists(template_image):
        print(f"Ошибка: Файл {template_image} не найден!")
        exit(1)
    
    print(f"Поиск '{template_image}' на '{main_image}'...")
    find_and_display_template(main_image, template_image)