44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
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()
|