Загрузка данных
import cv2
import numpy as np
import os
class ChessAnalyzer:
def __init__(self, templates_dir='templates', win_size=(64, 64)):
self.templates_dir = templates_dir
self.win_size = win_size
# Символы FEN
self.piece_symbols = {
'white_pawn': 'P', 'white_knight': 'N', 'white_bishop': 'B',
'white_rook': 'R', 'white_queen': 'Q', 'white_king': 'K',
'black_pawn': 'p', 'black_knight': 'n', 'black_bishop': 'b',
'black_rook': 'r', 'black_queen': 'q', 'black_king': 'k'
}
self.templates = {}
self.load_templates()
def load_templates(self):
"""Загрузка PNG-шаблонов"""
if not os.path.exists(self.templates_dir):
print(f"[!] Ошибка: Папка '{self.templates_dir}' не найдена.")
return
for filename in os.listdir(self.templates_dir):
if filename.lower().endswith('.png'):
piece_name = os.path.splitext(filename)[0]
path = os.path.join(self.templates_dir, filename)
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
if img is None:
continue
img_resized = cv2.resize(img, self.win_size)
if len(img_resized.shape) == 3 and img_resized.shape[2] == 4:
alpha = img_resized[:, :, 3] / 255.0
gray_piece = cv2.cvtColor(img_resized[:, :, :3], cv2.COLOR_BGR2GRAY)
bg = np.ones_like(gray_piece, dtype=np.uint8) * 128
composite_gray = (gray_piece * alpha + bg * (1.0 - alpha)).astype(np.uint8)
else:
composite_gray = cv2.cvtColor(img_resized, cv2.COLOR_BGR2GRAY) if len(img_resized.shape) == 3 else img_resized
self.templates[piece_name] = composite_gray
def remove_highlights(self, cell_img):
"""Убирает подсвеченные области, точки ходов и шах"""
if cell_img is None or cell_img.size == 0 or len(cell_img.shape) < 3:
return cell_img
clean_img = cell_img.copy()
h, w = cell_img.shape[:2]
# 1. ТОЧКИ ХОДОВ В ЦЕНТРЕ КЛЕТКИ
center_mask = np.zeros((h, w), dtype=np.uint8)
cv2.circle(center_mask, (w // 2, h // 2), int(min(h, w) * 0.28), 255, -1)
gray = cv2.cvtColor(cell_img, cv2.COLOR_BGR2GRAY)
mean_val, _ = cv2.meanStdDev(gray)
dot_mask = cv2.inRange(gray, int(mean_val[0][0] - 50), int(mean_val[0][0] + 50))
dot_mask = cv2.bitwise_and(dot_mask, center_mask)
# 2. HSV ЦВЕТОВАЯ МАСКА
hsv = cv2.cvtColor(cell_img, cv2.COLOR_BGR2HSV)
mask_green_blue = cv2.inRange(hsv, np.array([22, 10, 30]), np.array([170, 255, 255]))
mask_red1 = cv2.inRange(hsv, np.array([0, 15, 40]), np.array([21, 255, 255]))
mask_red2 = cv2.inRange(hsv, np.array([171, 15, 40]), np.array([180, 255, 255]))
color_mask = mask_green_blue | mask_red1 | mask_red2
full_mask = cv2.bitwise_or(dot_mask, color_mask)
if not np.any(full_mask):
return cell_img
highlight_pixels = np.sum(full_mask > 0)
total_pixels = h * w
if highlight_pixels < (total_pixels * 0.40):
return cv2.inpaint(cell_img, full_mask, inpaintRadius=5, flags=cv2.INPAINT_TELEA)
s_channel = hsv[:, :, 1]
v_channel = hsv[:, :, 2]
piece_mask = (s_channel < 10) | (v_channel < 40)
bg_mask = cv2.bitwise_and(full_mask, cv2.bitwise_not(piece_mask.astype(np.uint8) * 255))
clean_img[bg_mask != 0] = [128, 128, 128]
return clean_img
def is_highlighted_cell(self, cell_img):
"""Проверяет, подсвечена ли вся клетка зелёным цветом (фон совершенного хода)"""
if cell_img is None or cell_img.size == 0 or len(cell_img.shape) < 3:
return False
hsv = cv2.cvtColor(cell_img, cv2.COLOR_BGR2HSV)
mask_green = cv2.inRange(hsv, np.array([25, 30, 30]), np.array([85, 255, 255]))
# Исключаем центр, чтобы точки возможных ходов при выделении не путали систему
h, w = cell_img.shape[:2]
center_mask = np.zeros((h, w), dtype=np.uint8)
cv2.circle(center_mask, (w // 2, h // 2), int(min(h, w) * 0.25), 255, -1)
border_mask = cv2.bitwise_not(center_mask)
green_in_border = cv2.bitwise_and(mask_green, border_mask)
border_pixels = np.sum(border_mask > 0)
green_border_pixels = np.sum(green_in_border > 0)
# Подсвечена ли рамка/фон клетки
return (green_border_pixels / border_pixels) > 0.15
def is_empty_cell(self, cell_img, stddev_threshold=12.0):
"""Проверка клетки на пустоту"""
gray = cv2.cvtColor(cell_img, cv2.COLOR_BGR2GRAY) if len(cell_img.shape) == 3 else cell_img
h, w = gray.shape
inner_crop = gray[int(h * 0.25):int(h * 0.75), int(w * 0.25):int(w * 0.75)]
stddev = cv2.meanStdDev(inner_crop)
return stddev[0][0] < stddev_threshold
def identify_piece(self, cell_img):
"""Распознавание фигуры"""
if cell_img is None or cell_img.size == 0 or not self.templates:
return ""
cleaned_cell = self.remove_highlights(cell_img)
if self.is_empty_cell(cleaned_cell):
return ""
cell_resized = cv2.resize(cleaned_cell, self.win_size)
gray_cell = cv2.cvtColor(cell_resized, cv2.COLOR_BGR2GRAY) if len(cell_resized.shape) == 3 else cell_resized
best_piece = ""
max_corr = -1.0
for piece_name, template_img in self.templates.items():
res = cv2.matchTemplate(gray_cell, template_img, cv2.TM_CCOEFF_NORMED)
max_val = res[0][0]
if max_val > max_corr:
max_corr = max_val
best_piece = piece_name
if max_corr > 0.35 and best_piece:
return self.piece_symbols.get(best_piece, "")
return ""
def get_board_state(self, image_input):
"""Получение матрицы 8х8 с символами FEN"""
if image_input is None:
return None
if isinstance(image_input, str):
if not os.path.exists(image_input):
print(f"[!] Файл '{image_input}' не найден!")
return None
board_img = cv2.imread(image_input)
else:
board_img = image_input
if board_img is None:
return None
h, w = board_img.shape[:2]
cell_h = h // 8
cell_w = w // 8
board_matrix = []
for row in range(8):
row_pieces = []
for col in range(8):
y1, y2 = row * cell_h, (row + 1) * cell_h
x1, x2 = col * cell_w, (col + 1) * cell_w
cell_crop = board_img[y1:y2, x1:x2]
symbol = self.identify_piece(cell_crop)
row_pieces.append(symbol)
board_matrix.append(row_pieces)
return board_matrix
def get_board_state_with_color(self, image_input):
"""Возвращает матрицу 8х8 и определенный цвет следующего хода (только при завершенном ходе)"""
board_matrix = self.get_board_state(image_input)
if not board_matrix:
return None, None
if isinstance(image_input, str):
board_img = cv2.imread(image_input)
else:
board_img = image_input
if board_img is None:
return board_matrix, None
h, w = board_img.shape[:2]
cell_h, cell_w = h // 8, w // 8
highlighted_cells = []
for row in range(8):
for col in range(8):
y1, y2 = row * cell_h, (row + 1) * cell_h
x1, x2 = col * cell_w, (col + 1) * cell_w
cell_crop = board_img[y1:y2, x1:x2]
if self.is_highlighted_cell(cell_crop):
highlighted_cells.append((row, col))
detected_next_turn = None
# ПРОВЕРКА: Подсветка выполненного хода всегда занимает минимум 2 клетки (откуда и куда).
# Если подсвечена только 1 клетка — это обычное выделение фигуры (кликом)
if len(highlighted_cells) >= 2:
for (r, c) in highlighted_cells:
symbol = board_matrix[r][c]
if symbol:
if symbol.isupper():
detected_next_turn = 'b' # Сходили белые -> следующий ход Чёрных
elif symbol.islower():
detected_next_turn = 'w' # Сходили чёрные -> следующий ход Белых
break
return board_matrix, detected_next_turn