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)
#cap.set(cv2.CAP_PROP_FRAME_WIDTH, 4096)
#cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 2160)

cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)

color = COLOR.WHITE
rank = RANK.FIVE
pieces = {
    PIECE.PAWN: [POSITION((FILE.A, rank)), POSITION((FILE.B, rank)), POSITION((FILE.C, rank)), POSITION((FILE.D, rank)),
                 POSITION((FILE.E, rank)), POSITION((FILE.F, rank)), POSITION((FILE.G, rank)), POSITION((FILE.H, 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)
        imwrite("whole_boards/boards_for_empty/lol_gg_xD.png", frame)

        break
        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)
                imwrite(f"training_images/{piece}/{position.color}_square/training_{position}_{datetime.utcnow().timestamp()}.png", square)


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