56 lines
1.4 KiB
Python
56 lines
1.4 KiB
Python
|
import itertools
|
||
|
|
from pathlib import Path
|
||
|
|
from threading import Thread
|
||
|
|
from time import sleep
|
||
|
|
import numpy as np
|
||
|
|
import cv2
|
||
|
|
import runner
|
||
|
|
from datetime import datetime
|
||
|
|
import utils
|
||
|
|
|
||
|
|
|
||
|
|
#cap = cv2.VideoCapture(0)
|
||
|
|
#cap = cv2.VideoCapture("rtsp://10.192.49.108:8080/h264_ulaw.sdp")
|
||
|
|
cap = cv2.VideoCapture(0)
|
||
|
|
|
||
|
|
|
||
|
|
piece = "knight"
|
||
|
|
color = "black"
|
||
|
|
|
||
|
|
|
||
|
|
rank = 8
|
||
|
|
|
||
|
|
pieces = {
|
||
|
|
'knight': [("E", rank), ("H", rank)],
|
||
|
|
'rook': [("A", rank), ("F", rank)],
|
||
|
|
'bishop': [("C", rank), ("D", rank)],
|
||
|
|
'king': [("G", rank)],
|
||
|
|
'queen': [("B", 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")
|
||
|
|
# cv2.imwrite(f"single_frame_{counter}.png", frame)
|
||
|
|
utils.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[0], position[1])
|
||
|
|
x, y = position
|
||
|
|
utils.imwrite(f"training_images/{piece}/{runner.compute_color(x, y)}_square/training_{x}{str(y)}_{datetime.utcnow().timestamp()}.png", square)
|
||
|
|
|
||
|
|
|
||
|
|
# When everything done, release the capture
|
||
|
|
cap.release()
|
||
|
|
cv2.destroyAllWindows()
|