You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
51 lines
1.7 KiB
51 lines
1.7 KiB
import cv2
|
|
import numpy as np
|
|
|
|
from model import Detector
|
|
from utils import get_bbox_by_point
|
|
|
|
if __name__ == '__main__':
|
|
detector = Detector(classes=[0, 2, 5, 7])
|
|
|
|
cap = cv2.VideoCapture(0)
|
|
display_name = 'detector'
|
|
cv2.namedWindow(display_name, cv2.WINDOW_NORMAL)
|
|
cv2.setWindowProperty(display_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
|
|
# cv2.resizeWindow(display_name, 960, 720)
|
|
|
|
while True:
|
|
# Capture frame-by-frame
|
|
ret, frame = cap.read()
|
|
frame_disp = frame.copy()
|
|
|
|
results = detector.predict(frame)
|
|
for result in results:
|
|
cls = result[0]
|
|
bbox = result[1:]
|
|
color = (255, 0, 0) if cls == 0 else (0, 255, 255)
|
|
cv2.rectangle(frame_disp, (bbox[0], bbox[1]), (bbox[2] + bbox[0], bbox[3] + bbox[1]),
|
|
color, 5)
|
|
|
|
point = np.array([frame.shape[1] // 2, frame.shape[0] // 2])
|
|
cv2.rectangle(frame_disp, (point[0], point[1]),
|
|
(point[0] + 1, point[1] + 1),
|
|
(0, 0, 255), 2)
|
|
if len(results):
|
|
point_bbox = get_bbox_by_point(results[:, 1:], point)
|
|
if point_bbox is not None:
|
|
cv2.rectangle(frame_disp, (point_bbox[0], point_bbox[1]), (point_bbox[2] + point_bbox[0], point_bbox[3] + point_bbox[1]),
|
|
(0, 0, 255), 2)
|
|
|
|
# Display the resulting frame
|
|
cv2.imshow(display_name, frame_disp)
|
|
key = cv2.waitKey(1)
|
|
if key == ord('q'):
|
|
break
|
|
elif key == ord('r'):
|
|
pass
|
|
elif key == ord('s'):
|
|
pass
|
|
|
|
# When everything done, release the capture
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|