File size: 6,605 Bytes
57f1366 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | import cv2
import numpy as np
import mss
class VisionHandler:
def __init__(self):
self.templates = {}
self.is_calibrated = False
def capture_screen(self, region):
"""
Capture a specific region of the screen.
region: {'top': int, 'left': int, 'width': int, 'height': int}
"""
with mss.mss() as sct:
# mss requires int for all values
monitor = {
"top": int(region["top"]),
"left": int(region["left"]),
"width": int(region["width"]),
"height": int(region["height"])
}
screenshot = sct.grab(monitor)
img = np.array(screenshot)
return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
def split_board(self, board_image):
"""
Split board image into 64 squares.
Assumption: board_image is cropped exactly to the board edges.
"""
h, w, _ = board_image.shape
sq_h, sq_w = h // 8, w // 8
squares = []
# Row-major order (rank 8 down to rank 1)
# Standard FEN order: Rank 8 (index 0) to Rank 1 (index 7)
for r in range(8):
row_squares = []
for c in range(8):
# Add a small crop margin to avoid border noise
margin_h = int(sq_h * 0.1)
margin_w = int(sq_w * 0.1)
y1 = r * sq_h + margin_h
y2 = (r + 1) * sq_h - margin_h
x1 = c * sq_w + margin_w
x2 = (c + 1) * sq_w - margin_w
square = board_image[y1:y2, x1:x2]
row_squares.append(square)
squares.append(row_squares)
return squares
def calibrate(self, board_image):
"""
Calibrate by learning piece appearance from a starting position board image.
Standard Starting Position:
r n b q k b n r
p p p p p p p p
. . . . . . . .
. . . . . . . .
. . . . . . . .
. . . . . . . .
P P P P P P P P
R N B Q K B N R
"""
squares = self.split_board(board_image)
self.templates = {}
# Dictionary mapping piece chars to list of coordinates (row, col) in start pos
# We'll take the average or just the first instance as template
# 0-indexed rows: 0=BlackBack, 1=BlackPawn, ... 6=WhitePawn, 7=WhiteBack
piece_map = {
'r': [(0, 0), (0, 7)],
'n': [(0, 1), (0, 6)],
'b': [(0, 2), (0, 5)],
'q': [(0, 3)],
'k': [(0, 4)],
'p': [(1, i) for i in range(8)],
'R': [(7, 0), (7, 7)],
'N': [(7, 1), (7, 6)],
'B': [(7, 2), (7, 5)],
'Q': [(7, 3)],
'K': [(7, 4)],
'P': [(6, i) for i in range(8)],
'.': [] # Empty squares
}
# Add empty squares (rows 2, 3, 4, 5)
for r in range(2, 6):
for c in range(8):
piece_map['.'].append((r, c))
# Store one template per piece type (simplest approach)
# Better: Store all samples and use KNN, but simple template match might suffice for consistent graphics
for p_char, coords in piece_map.items():
if not coords:
continue
# Use the first coordinate as the primary template
r, c = coords[0]
template = squares[r][c]
# We convert to grayscale for simpler matching
gray_template = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)
# Store just one template for now
self.templates[p_char] = gray_template
self.is_calibrated = True
print("Calibration complete.")
def match_square(self, square_img):
if not self.templates:
return '.'
gray_sq = cv2.cvtColor(square_img, cv2.COLOR_BGR2GRAY)
best_score = float('inf')
best_piece = '.'
# Compare against all templates using MSE
# Note: Templates must be same size. If slightly different due to rounding, resize.
target_h, target_w = gray_sq.shape
for p_char, template in self.templates.items():
# Resize template to match square if needed (should be identical if grid is uniform)
if template.shape != gray_sq.shape:
template = cv2.resize(template, (target_w, target_h))
# Mean Squared Error
err = np.sum((gray_sq.astype("float") - template.astype("float")) ** 2)
err /= float(gray_sq.shape[0] * gray_sq.shape[1])
if err < best_score:
best_score = err
best_piece = p_char
return best_piece
def get_fen_from_image(self, board_image):
if not self.is_calibrated:
# Fallback or error
return None
squares = self.split_board(board_image)
fen_rows = []
for r in range(8):
empty_count = 0
row_str = ""
for c in range(8):
piece = self.match_square(squares[r][c])
if piece == '.':
empty_count += 1
else:
if empty_count > 0:
row_str += str(empty_count)
empty_count = 0
row_str += piece
if empty_count > 0:
row_str += str(empty_count)
fen_rows.append(row_str)
# Join with '/'
fen_board = "/".join(fen_rows)
# Add default game state info (w KQkq - 0 1)
# TODO: Detect turn based on external logic or diff?
# For now, we might alternate or assume it's always the user's turn to move?
# Actually, if we are just showing analysis, we want the engine to evaluate for the SIDE TO MOVE.
# But we don't know who is to move just from the static board.
# We can try to infer from previous state or just ask the user.
# For this version, let's return the board part, and handle the rest in logic.
return f"{fen_board} w KQkq - 0 1"
|