advancedskrald/opencv_video.py

from datetime import datetime

import cv2

import runner
from util import FILE, RANK, PIECE, COLOR, imwrite, POSITION

cap = cv2.VideoCapture(0)

color = COLOR.BLACK
rank = RANK.EIGHT
pieces = {
    PIECE.rook: [POSITION(FILE.A, rank), POSITION(FILE.F, rank)],
    PIECE.knight: [POSITION(FILE.E, rank), POSITION(FILE.H, rank)],
    PIECE.bishop: [POSITION(FILE.C, rank), POSITION(FILE.D, rank)],
    PIECE.queen: [POSITION(FILE.B, rank)],
    PIECE.king: [POSITION(FILE.G, rank)],
}

while True:
    # Capture frame-by-frame
    ret, frame = cap.read()

    # Display the resulting frame
    cv2.imshow("frame", frame)

    if cv2.waitKey(100) & 0xFF == ord("c"):
        print(f"capturing frame")
        imwrite(f"whole_boards/boards_for_empty/board_{datetime.utcnow().timestamp()}_.png", frame)
        warped = runner.warp_board(frame)

        runner.save_empty_fields(warped, skip_rank=rank)

        for piece, positions in pieces.items():
            for position in positions:
                square = runner.get_square(warped, *position)
                x, y = position
                imwrite(f"training_images/{piece}/{position.color}_square/training_{x}{str(y)}_{datetime.utcnow().timestamp()}.png", square)


# When everything done, release the capture
cap.release()
cv2.destroyAllWindows()