Загрузка данных
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
self.hog = cv2.HOGDescriptor(
_winSize=self.win_size,
_blockSize=(16, 16),
_blockStride=(8, 8),
_cellSize=(8, 8),
_nbins=9
)
self.piece_symbols = {
'white_pawn': 'wP', 'white_knight': 'wN', 'white_bishop': 'wB',
'white_rook': 'wR', 'white_queen': 'wQ', 'white_king': 'wK',
'black_pawn': 'bP', 'black_knight': 'bN', 'black_bishop': 'bB',
'black_rook': 'bR', 'black_queen': 'bQ', 'black_king': 'bK'
}
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 - alpha)).astype(np.uint8)
else:
composite_gray = cv2.cvtColor(img_resized, cv2.COLOR_BGR2GRAY) if len(img_resized.shape) == 3 else img_resized
composite_gray = cv2.equalizeHist(composite_gray)
hog_feat = self.hog.compute(composite_gray)
self.templates[piece_name] = hog_feat
def is_empty_cell(self, cell_img, stddev_threshold=16.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.2):int(h*0.8), int(w*0.2):int(w*0.8)]
_, stddev = cv2.meanStdDev(inner_crop)
return stddev[0][0] < stddev_threshold
def segment_piece_mask(self, cell_img):
"""
Сегментация: вырезает силуэт фигуры, отделяя его от фона клетки.
Возвращает бинарную маску (255 там, где фигура, и 0 там, где фон).
"""
gray = cv2.cvtColor(cell_img, cv2.COLOR_BGR2GRAY) if len(cell_img.shape) == 3 else cell_img
h, w = gray.shape
# 1. Сглаживаем шум
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
# 2. Вычисляем градиент (Canny / Sobel) для поиска контуров объекта
edges = cv2.Canny(blurred, 30, 100)
# 3. Применяем морфологическую закрывающую операцию, чтобы объединить контуры фигуры
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
closed_edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
# 4. Находим контуры и берём самый крупный в центре (это и будет фигура)
contours, _ = cv2.findContours(closed_edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
mask = np.zeros_like(gray)
if contours:
# Находим контур, расположенный ближе всего к центру
center_point = (w // 2, h // 2)
best_cnt = None
min_dist = float('inf')
for cnt in contours:
area = cv2.contourArea(cnt)
if area > (h * w * 0.05): # фильтруем мелкий шум
M = cv2.moments(cnt)
if M["m00"] != 0:
cx = int(M["m10"] / M["m00"])
cy = int(M["m01"] / M["m00"])
dist = (cx - center_point[0])**2 + (cy - center_point[1])**2
if dist < min_dist:
min_dist = dist
best_cnt = cnt
if best_cnt is not None:
cv2.drawContours(mask, [best_cnt], -1, 255, thickness=cv2.FILLED)
# Если контур выделить не удалось — используем стандартный центральный прямоугольник
if np.sum(mask) == 0:
mask[int(h*0.25):int(h*0.75), int(w*0.25):int(w*0.75)] = 255
return mask
def get_piece_color(self, cell_img):
"""
Определяет цвет фигуры, вычисляя яркость пикселей СТРОГО внутри выделенной маски.
"""
if cell_img is None or cell_img.size == 0:
return 'white'
gray = cv2.cvtColor(cell_img, cv2.COLOR_BGR2GRAY) if len(cell_img.shape) == 3 else cell_img
# Получаем сегментированную маску фигуры
mask = self.segment_piece_mask(cell_img)
# Извлекаем пиксели, принадлежащие строго фигуре
piece_pixels = gray[mask == 255]
if len(piece_pixels) == 0:
return 'white'
# Анализируем яркость самой фигуры
mean_brightness = np.mean(piece_pixels)
median_brightness = np.median(piece_pixels)
# Чёрная фигура редко бывает ярче 110, белая — обычно выше 130
if median_brightness < 115 or mean_brightness < 115:
return 'black'
else:
return 'white'
def identify_piece(self, cell_img):
"""Определяет фигуру на отдельной клетке"""
if cell_img is None or cell_img.size == 0 or not self.templates:
return ""
if self.is_empty_cell(cell_img):
return ""
detected_color = self.get_piece_color(cell_img)
cell_resized = cv2.resize(cell_img, self.win_size)
gray = cv2.cvtColor(cell_resized, cv2.COLOR_BGR2GRAY) if len(cell_resized.shape) == 3 else cell_resized
gray = cv2.equalizeHist(gray)
cell_hog = self.hog.compute(gray)
best_piece = ""
min_dist = float('inf')
for piece_name, template_hog in self.templates.items():
if not piece_name.startswith(detected_color):
continue
dist = np.linalg.norm(cell_hog - template_hog)
if dist < min_dist:
min_dist = dist
best_piece = piece_name
if min_dist < 22.0 and best_piece:
return self.piece_symbols.get(best_piece, "")
return ""
def get_board_state(self, image_input):
"""Возвращает матрицу 8x8 со стоящими фигурами"""
if image_input is None:
print("[!] Ошибка: На вход передано пустое изображение (None).")
return []
if isinstance(image_input, str):
if not os.path.exists(image_input):
print(f"[!] Файл '{image_input}' не найден!")
return []
board_img = cv2.imread(image_input)
else:
board_img = image_input
if board_img is None:
print("[!] Ошибка: Не удалось прочитать изображение.")
return []
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 analyze_initial_board(self, image_input):
"""Псевдоним для совместимости со старыми вызовами"""
return self.get_board_state(image_input)