Files
gobot/board-vision/src/main.py
2024-09-24 12:14:16 +02:00

74 lines
2.0 KiB (Stored with Git LFS)
Python

import cv2
import argparse
import enum
import uuid
import numpy as np
from vision.countourFiltering import filter_countours
from vision.preproc import edge_dection, thresholding_edge_detection
from vision.preproc_fsm import PREPROCESSING_MODE
# Parse arguments
parser = argparse.ArgumentParser()
parser.add_argument("camera", type=int, help="Camera index")
args = parser.parse_args()
cam = cv2.VideoCapture(args.camera)
WINDOW_NAME = "Output Go Board Detection"
cv2.namedWindow(WINDOW_NAME)
cv2.moveWindow(WINDOW_NAME, 1,0)
font = cv2.FONT_HERSHEY_SIMPLEX
font_color = (127, 255, 0)
fps = 0
try:
while True:
ret, frame = cam.read()
t_start = cv2.getTickCount()
width, height = frame.shape[1], frame.shape[0]
frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
frame_edges = edge_dection(frame_gray, PREPROCESSING_MODE.THRES)
res = np.concatenate((
np.concatenate((
frame,
frame,
frame
), axis=1),
np.concatenate((
cv2.cvtColor(frame_gray, cv2.COLOR_GRAY2BGR),
cv2.cvtColor(frame_edges, cv2.COLOR_GRAY2BGR),
frame
), axis=1)
), axis=0)
cv2.imshow(WINDOW_NAME, res)
t_end = cv2.getTickCount()
fps = cv2.getTickFrequency() / (t_end - t_start)
k = cv2.waitKey(1)
if k%256 == 27: # ESC pressed
print("Escape hit, closing...")
break
elif k%256 == 32: # SPACE pressed
print("Space hit, saving image")
image_id = uuid.uuid4()
cv2.imwrite(f"{image_id}_gray.png", frame_gray)
cv2.imwrite(f"{image_id}_thres.png", edge_dection(frame_gray, PREPROCESSING_MODE.THRES))
cv2.imwrite(f"{image_id}_canny.png", edge_dection(frame_gray, PREPROCESSING_MODE.CANNY))
except KeyboardInterrupt:
print("Exiting")
cam.release()
cv2.destroyAllWindows()