new_thea/Control/Jetson/tests/__vision_threaded_cuda.py
2021-09-21 12:11:46 +01:00

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')