draft of image tracker

untested
This commit is contained in:
Mathematical Michael 2023-09-19 03:47:00 -06:00
parent 4d67c97b44
commit 63f9be0a4d

53
track.py Normal file
View File

@ -0,0 +1,53 @@
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()