98 lines
2.8 KiB
Python
98 lines
2.8 KiB
Python
from threading import Thread
|
|
|
|
import cv2
|
|
import time
|
|
import apriltag
|
|
import logging
|
|
|
|
class VideoStream:
|
|
def __init__(self, cam_id=0):
|
|
# self.stream = cv2.VideoCapture(cam_id)
|
|
# self.stream.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
|
|
# self.stream.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) test.mp4
|
|
self.stream = self.open_cam_usb(cam_id, 1280, 720)
|
|
_, img = self.stream.read()
|
|
self.buffer = img
|
|
self.frames = 0
|
|
self.running = False
|
|
|
|
def start(self):
|
|
t = Thread(target=self.update, args=())
|
|
t.daemon = True
|
|
self.running = True
|
|
t.start()
|
|
self.start_time = time.time()
|
|
return self
|
|
|
|
def update(self):
|
|
while self.running:
|
|
_, frame = self.stream.read()
|
|
self.buffer = frame
|
|
|
|
self.frames += 1
|
|
if self.frames == 30:
|
|
time_taken = time.time() - self.start_time
|
|
self.start_time = time.time()
|
|
self.frames = 0
|
|
fps = 30 / time_taken
|
|
print(f"##### INGEST FPS: {fps} #####")
|
|
# print('#')
|
|
self.stream.release()
|
|
|
|
def read(self):
|
|
return self.buffer
|
|
|
|
def stop(self):
|
|
self.running = False
|
|
|
|
def open_cam_usb(self, dev, width, height):
|
|
# We want to set width and height here, otherwise we could just do:
|
|
# return cv2.VideoCapture(dev)
|
|
# We might get better framerates with this implementation...
|
|
gst_str = ('v4l2src device=/dev/video{} ! '
|
|
'video/x-raw, width=(int){}, height=(int){} ! '
|
|
'videoconvert ! appsink').format(dev, width, height)
|
|
return cv2.VideoCapture(gst_str, cv2.CAP_GSTREAMER)
|
|
|
|
|
|
def cuda_gray_conveter(frame):
|
|
gpu_frame = cv2.cuda_GpuMat()
|
|
gpu_frame.upload(frame)
|
|
gray_cuda = cv2.cuda.cvtColor(gpu_frame, cv2.COLOR_BGR2GRAY)
|
|
gray = gray_cuda.download()
|
|
return gray
|
|
|
|
|
|
options = apriltag.DetectorOptions(families="tag36h11", nthreads=2)
|
|
detector = apriltag.Detector(options)
|
|
|
|
logging.basicConfig(level='INFO')
|
|
|
|
stream_0 = VideoStream(0).start()
|
|
# stream_2 = VideoStream(1).start()
|
|
|
|
frames = 0
|
|
start = time.time()
|
|
running = True
|
|
while running:
|
|
try:
|
|
frame = stream_0.read()
|
|
if frame.any():
|
|
gray_frame = cuda_gray_conveter(frame)
|
|
result = detector.detect(gray_frame)
|
|
if result:
|
|
logging.info(result[0].center)
|
|
frames += 1
|
|
if frames == 30:
|
|
time_taken = time.time() - start
|
|
start = time.time()
|
|
frames = 0
|
|
fps = 30 / time_taken
|
|
print(f"##### TOTAL PROCESSING FPS: {fps} #####")
|
|
except Exception as e:
|
|
print('Exiting cleanly...')
|
|
stream_0.stop()
|
|
running = False
|
|
|
|
print('Exited')
|