54 lines
2.2 KiB
Python
54 lines
2.2 KiB
Python
import cv2
|
|
import numpy as np
|
|
import chess
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.cm as cm
|
|
|
|
def capture_image(url):
|
|
"""Capture an image from the webcam stream."""
|
|
cap = cv2.VideoCapture(url)
|
|
ret, frame = cap.read()
|
|
cap.release()
|
|
return frame
|
|
|
|
def extract_fen_from_image(image):
|
|
"""Detect the chessboard from the image and extract the FEN string.
|
|
NOTE: This is a simplified version and might not work for all scenarios.
|
|
A more robust solution would involve training a model for piece detection."""
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
ret, corners = cv2.findChessboardCorners(gray, (8,8))
|
|
if not ret:
|
|
return None
|
|
# Here we would typically identify each square and its piece to extract the FEN.
|
|
# For simplicity, returning a default FEN for now.
|
|
return "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
|
|
|
|
def generate_heatmap_overlay(board_shape, aggregated_targets, cmap='viridis', alpha=0.5):
|
|
"""Generate heatmap overlay based on aggregated targets."""
|
|
heatmap = np.zeros(board_shape[:2], dtype=np.uint8)
|
|
for square, count in aggregated_targets.items():
|
|
x, y = chess.SQUARE_NAMES.index(square) % 8, chess.SQUARE_NAMES.index(square) // 8
|
|
heatmap[y, x] = count
|
|
heatmap_normalized = cv2.normalize(heatmap, None, 0, 255, cv2.NORM_MINMAX)
|
|
heatmap_colored = cm.get_cmap(cmap)(heatmap_normalized)[:, :, :3] * 255
|
|
heatmap_overlay = (heatmap_colored * alpha + board_shape * (1 - alpha)).astype(np.uint8)
|
|
return heatmap_overlay
|
|
|
|
def overlay_heatmap_on_image(image, heatmap_overlay):
|
|
"""Overlay heatmap onto the original image."""
|
|
return cv2.addWeighted(image, 1, heatmap_overlay, 1, 0)
|
|
|
|
if __name__ == "__main__":
|
|
url = "http://10.1.10.159:8000/stream.mjpg"
|
|
image = capture_image(url)
|
|
fen = extract_fen_from_image(image)
|
|
if fen:
|
|
board = chess.Board(fen)
|
|
moves = get_legal_moves(board)
|
|
aggregated_targets = aggregate_target_squares(moves)
|
|
heatmap_overlay = generate_heatmap_overlay(image.shape, aggregated_targets)
|
|
final_image = overlay_heatmap_on_image(image, heatmap_overlay)
|
|
cv2.imshow('Overlayed Chessboard', final_image)
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows()
|