import time import threading import cv2 import paho.mqtt.client as mqtt import numpy as np import pycuda.autoinit from utils.yolo_classes import get_cls_dict from utils.yolo_with_plugins import TrtYOLO from utils.display import open_window, set_display, show_fps from utils.visualization import BBoxVisualization class TrtThread(threading.Thread): def __init__(self, model, conf_th): threading.Thread.__init__(self) self.model = model self.conf_th = conf_th self.trt_yolo = None # to be created when run self.running = False def run(self): global image global cnbr print('TrtThread: loading the TRT YOLO engine...') self.trt_yolo = TrtYOLO(self.model, INPUT_HW) print('TrtThread: start running...waiting for images to analyze') self.running = True full_scrn = False cls_dict = get_cls_dict(80) vis = BBoxVisualization(cls_dict) open_window(WINDOW_NAME, 320, 240, 'Camera TensorRT YOLO Demo') while self.running: if image is not None: print('Found image from camera: ', cnbr) img = image image = None tic = time.time() boxes, confs, clss = self.trt_yolo.detect(img, self.conf_th) toc = time.time() print('Detection time: ', (toc - tic)) (H, W) = img.shape[:2] img = vis.draw_bboxes(img, boxes, confs, clss) cv2.imshow(WINDOW_NAME, img) # prepare & send image via mqtt client_thr = mqtt.Mosquitto() client_thr.connect("127.0.0.1", 1883, 60) client_thr.loop_start() img_str = cv2.imencode('.jpg', img)[1].tostring() client_thr.publish("result", img_str, 1, False) client_thr.loop_stop() client_thr.disconnect() key = cv2.waitKey(1) if key == 27: # ESC key: quit program break elif key == ord('F') or key == ord('f'): # Toggle fullscreen full_scrn = not full_scrn set_display(WINDOW_NAME, full_scrn) cv2.destroyAllWindows() del self.trt_yolo print('TrtThread: stopped...') def stop(self): self.running = False self.join() def initmqtt(mqttThreadEvent): def on_connect(client, userdata, flags, rc): print("Connected to broker:", rc) client.subscribe("image/#", 0) def on_subscribe(client, userdata, mid, granted_qos): print ('Subscribed to topics:', userdata, mid, granted_qos) def on_message(client, userdata, msg): global image global cnbr global th_abort # print (msg.payload) # print (len(msg.payload)) if(len(msg.payload)>1000): cnbr = msg.topic.split('/')[1] # convert string of image data to uint8 nparr = np.fromstring(msg.payload, np.uint8) # decode image image = cv2.imdecode(nparr, cv2.IMREAD_COLOR) if msg.payload.decode("utf-8") == 'Stop-YOLO-Analyzer': th_abort = True client = mqtt.Mosquitto() client.on_connect = on_connect client.on_message = on_message client.on_subscribe = on_subscribe resp = client.connect("127.0.0.1", 1883, 60) client.loop_start() while th_abort == False: time.sleep(1) client.loop_stop() client.disconnect() print ("mqttThread terminated") exit() WINDOW_NAME = 'TrtYOLOdemoAsync' INPUT_HW = (416, 416) image = None th_abort = False cnbr = 0 mqttThreadEvent = threading.Event() mqttThread = threading.Thread( target=initmqtt, args=(mqttThreadEvent,) ) mqttThread.start() trt_thread = TrtThread('yolov4-tiny-opt-416', conf_th=0.3) trt_thread.start() # start the child thread while th_abort == False: time.sleep(1) trt_thread.stop() # stop the child thread