identify.py - master - gitdev/homm5stacks

identify.py
97 строк3 KB
gitdev
initial commit
23 май 2025, 13:24
23 май 2025, 13:249f8cdfd
from collections import Counter
import mss
import numpy as np
from PIL import Image
import cv2

import settings

def capture_section(x0, y0, x1, y1):
    with mss.mss() as sct:
        monitor = {"top": y0, "left": x0, "width": x1 - x0, "height": y1 - y0}
        screenshot = sct.grab(monitor)
        img = Image.frombytes("RGB", screenshot.size, screenshot.rgb)
        return img

def crop_center_square(subsection):
    # Get the dimensions of the subsection
    width, height = subsection.size
    # Calculate the coordinates of the central square
    side_length = min(width, height) // 2
    x0 = (width - side_length) // 2
    y0 = (height - side_length) // 2
    x1 = x0 + side_length
    y1 = y0 + side_length
    # Crop and return the central square
    return subsection.crop((x0, y0, x1, y1))

def get_dominant_color(subsection):
    # Resize to speed up processing (optional)
    small_img = subsection.resize((50, 50))  # Reduce size for faster computation
    # Convert to numpy array and reshape
    arr = np.array(small_img)
    pixels = arr.reshape(-1, arr.shape[-1])
    # Count unique colors
    dominant_color = Counter(map(tuple, pixels)).most_common(1)[0][0]
    return np.array(dominant_color)

def is_mono_colored(subsection):
    # Get the dominant color of the subsection
    dominant_color = get_dominant_color(subsection)
    # Convert subsection to numpy array
    arr = np.array(subsection)
    # Calculate the color difference for each pixel
    color_diff = np.linalg.norm(arr - dominant_color, axis=-1)
    # Check if most pixels are close to the dominant color
    return np.mean(color_diff) < 20  # Threshold for similarity

def detect_paint_style(subsection):
    # Convert to grayscale for edge detection
    gray = cv2.cvtColor(np.array(subsection), cv2.COLOR_RGB2GRAY)
    # Apply Canny edge detection
    edges = cv2.Canny(gray, 100, 200)
    # Count non-zero pixels in the edge image
    edge_density = np.sum(edges > 0) / edges.size
    return edge_density > 0.01  # Adjust threshold based on your images

def collect_empty_slots():
    x0, x1, y0, y1 = settings.BOX  # stack holder coordinates

    width_of_subsection = (x1 - x0) // 7

    selection=[]

    for i in range(7):
        sub_x0 = x0 + i * width_of_subsection
        sub_x1 = sub_x0 + width_of_subsection
        
        # Capture the subsection
        subsection = capture_section(sub_x0, y0, sub_x1, y1)
        # Crop the central square
        central_square = crop_center_square(subsection)

        if is_mono_colored(central_square):
            selection.append(i)
    
    return selection

if __name__=="__main__":
    # Main Processing
    x0, x1, y0, y1 = settings.BOX  # stack holder coordinates
    width_of_subsection = (x1 - x0) // 7

    for i in range(7):
        sub_x0 = x0 + i * width_of_subsection
        sub_x1 = sub_x0 + width_of_subsection
        
        # Capture the subsection
        subsection = capture_section(sub_x0, y0, sub_x1, y1)
        # Crop the central square
        central_square = crop_center_square(subsection)
        
        if is_mono_colored(central_square):
            print(f"Subsection {i} is empty (mono-colored).")
        elif detect_paint_style(central_square):
            print(f"Subsection {i} contains a painted-style portrait.")
        else:
            print(f"Subsection {i} does not match the expected styles.")