cv2_template_matching

import cv2 as cv
import numpy as np
import os
from matplotlib import pyplot as plt

# Paths and constants
screenshot_path = 'images/breach_protocol_screenshot.png'
HEX_IMAGES_PATH = "hexcodes"
MATRIX = "/matrix/"
SEQUENCE = "/sequence/"

VALID_HEX_VALUES = ["55", "1C", "BD", "E9", "7A", "FF"]
threshold = 0.8  # adjust as needed

# Load screenshot in grayscale
screenshot = cv.imread(screenshot_path, cv.IMREAD_GRAYSCALE)
assert screenshot is not None, "Screenshot could not be loaded"

# Convert to BGR for visualization
output_img = cv.cvtColor(screenshot, cv.COLOR_GRAY2BGR)

# Generate colors for visualization
np.random.seed(42)
colors = {code: tuple(int(c) for c in np.random.randint(0, 255, 3)) for code in VALID_HEX_VALUES}

def build_grid(positions, row_threshold=15):
    """
    Build a 2D grid (list of lists) from detected positions.
    Positions format: [(x, y, code), ...]
    Sorted top-to-bottom, left-to-right.
    row_threshold defines max pixel difference to be in same row.
    """
    # Sort by y (top to bottom)
    positions = sorted(positions, key=lambda p: p[1])

    rows = []
    current_row = []
    last_y = None

    # Group into rows based on y threshold
    for pos in positions:
        x, y, code = pos
        if last_y is None or abs(y - last_y) <= row_threshold:
            current_row.append(pos)
        else:
            # sort current row by x (left to right)
            current_row.sort(key=lambda p: p[0])
            rows.append(current_row)
            current_row = [pos]
        last_y = y
    if current_row:
        current_row.sort(key=lambda p: p[0])
        rows.append(current_row)

    # Build the final grid of codes
    grid = []
    for row in rows:
        codes_row = [code for (_, _, code) in row]
        grid.append(codes_row)

    return grid

def detect_hex_codes(base_img, valid_codes, templates_folder, offset=10):
    """
    Detect hex codes in base_img using templates from templates_folder.
    Returns list of (x, y, code) for detected positions.
    """
    found_positions = []
    for hex_code in valid_codes:
        template_path = os.path.join(templates_folder, f"{hex_code}.png")
        template = cv.imread(template_path, cv.IMREAD_GRAYSCALE)
        assert template is not None, f"Template {hex_code} could not be loaded from {template_path}"
        w, h = template.shape[::-1]

        res = cv.matchTemplate(base_img, template, cv.TM_CCOEFF_NORMED)
        loc = np.where(res >= threshold)

        for pt in zip(*loc[::-1]):
            # filter out points too close to existing detections
            if all(np.linalg.norm(np.array(pt) - np.array(pos[:2])) > offset for pos in found_positions):
                found_positions.append((pt[0], pt[1], hex_code))

    return found_positions

# --- Detect MATRIX hex codes ---
found_positions_matrix = detect_hex_codes(screenshot, VALID_HEX_VALUES, HEX_IMAGES_PATH + MATRIX)
print(f"Detected {len(found_positions_matrix)} hex codes in MATRIX area.")

# Draw MATRIX detections
for x, y, code in found_positions_matrix:
    template_path = os.path.join(HEX_IMAGES_PATH + MATRIX, f"{code}.png")
    template = cv.imread(template_path, cv.IMREAD_GRAYSCALE)
    w, h = template.shape[::-1]
    cv.rectangle(output_img, (x, y), (x + w, y + h), colors[code], 2)
    cv.putText(output_img, "M:" + code, (x, y - 5), cv.FONT_HERSHEY_SIMPLEX, 0.5, colors[code], 1, cv.LINE_AA)

matrix_grid = build_grid(found_positions_matrix)
print("MATRIX GRID:")
for row in matrix_grid:
    print(row)

# --- Detect SEQUENCE hex codes ---
found_positions_sequence = detect_hex_codes(screenshot, VALID_HEX_VALUES, HEX_IMAGES_PATH + SEQUENCE)
print(f"Detected {len(found_positions_sequence)} hex codes in SEQUENCE area.")

# Draw SEQUENCE detections (different color - e.g. white)
for x, y, code in found_positions_sequence:
    template_path = os.path.join(HEX_IMAGES_PATH + SEQUENCE, f"{code}.png")
    template = cv.imread(template_path, cv.IMREAD_GRAYSCALE)
    w, h = template.shape[::-1]
    # Use a fixed color for SEQUENCE (e.g., white)
    cv.rectangle(output_img, (x, y), (x + w, y + h), (255, 255, 255), 2)
    cv.putText(output_img, "S:" + code, (x, y - 5), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv.LINE_AA)

sequence_grid = build_grid(found_positions_sequence)
print("SEQUENCE GRID:")
for row in sequence_grid:
    print(row)

# Show all detections
plt.figure(figsize=(15, 10))
plt.imshow(cv.cvtColor(output_img, cv.COLOR_BGR2RGB))
plt.title("Detected Hex Codes: MATRIX (colors) & SEQUENCE (white)")
plt.axis('off')
plt.show()