def detect(q1, q2): model_config = parser.parse_args() print("Loading detector") model = "ssd_mobilenet_v2_coco" conf_th = 0.3 INPUT_HW = (300, 300) cls_dict = get_cls_dict("coco") vis = BBoxVisualization(cls_dict) trt_ssd = TrtSSD(model, INPUT_HW) print("Loading detector complete") if model_config.ui == 1: cv2.startWindowThread() cv2.namedWindow("window") while 1: try: frame, frame_time = q1.get() delay = time.time() - frame_time if delay > 0.4: print("Skipping frame") continue boxes, confs, clss = trt_ssd.detect(frame, conf_th) print([get_cls_dict("coco")[c] for c in clss]) if model_config.ui == 1: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) img = vis.draw_bboxes(frame, boxes, confs, clss) cv2.imshow('window', img[..., ::-1]) except Exception as e: traceback.print_exc()
def run(self): """Run until 'running' flag is set to False by main thread. NOTE: CUDA context is created here, i.e. inside the thread which calls CUDA kernels. In other words, creating CUDA context in __init__() doesn't work. """ global s_img, s_boxes, s_confs, s_clss print('TrtThread: loading the TRT SSD engine...') self.cuda_ctx = cuda.Device(0).make_context() # GPU 0 self.trt_ssd = TrtSSD(self.model, INPUT_HW) print('TrtThread: start running...') self.running = True while self.running: img = self.cam.read() if img is None: break boxes, confs, clss = self.trt_ssd.detect(img, self.conf_th) with self.condition: s_img, s_boxes, s_confs, s_clss = img, boxes, confs, clss self.condition.notify() del self.trt_ssd self.cuda_ctx.pop() del self.cuda_ctx print('TrtThread: stopped...')
def main(): args = parse_args() cam = Camera(args) cam.open() if not cam.is_opened: sys.exit('Failed to open camera!') cls_dict = get_cls_dict(args.model) trt_ssd = TrtSSD(args.model, INPUT_HW) cam.start() if args.use_console: loop_and_detect_console(cam, trt_ssd, conf_th=0.3, loop=args.loop, cls_dict=cls_dict) else: open_window(WINDOW_NAME, args.image_width, args.image_height, 'Camera TensorRT SSD Demo for Jetson Nano') vis = BBoxVisualization(cls_dict) loop_and_detect(cam, trt_ssd, conf_th=0.3, vis=vis) cam.stop() cam.release() cv2.destroyAllWindows()
class TrtThread(threading.Thread): """TrtThread This implements the child thread which continues to read images from cam (input) and to do TRT engine inferencing. The child thread stores the input image and detection results into global variables and uses a condition varaiable to inform main thread. In other words, the TrtThread acts as the producer while the main thread is the consumer. """ def __init__(self, condition, cam, model, conf_th): """__init__ # Arguments condition: the condition variable used to notify main thread about new frame and detection result cam: the camera object for reading input image frames model: a string, specifying the TRT SSD model conf_th: confidence threshold for detection """ threading.Thread.__init__(self) self.condition = condition self.cam = cam self.model = model self.conf_th = conf_th self.cuda_ctx = None # to be created when run self.trt_ssd = None # to be created when run self.running = False def run(self): """Run until 'running' flag is set to False by main thread. NOTE: CUDA context is created here, i.e. inside the thread which calls CUDA kernels. In other words, creating CUDA context in __init__() doesn't work. """ global s_img, s_boxes, s_confs, s_clss print('TrtThread: loading the TRT SSD engine...') self.cuda_ctx = cuda.Device(0).make_context() # GPU 0 self.trt_ssd = TrtSSD(self.model, INPUT_HW) print('TrtThread: start running...') self.running = True while self.running: img = self.cam.read() if img is None: break boxes, confs, clss = self.trt_ssd.detect(img, self.conf_th) with self.condition: s_img, s_boxes, s_confs, s_clss = img, boxes, confs, clss self.condition.notify() del self.trt_ssd self.cuda_ctx.pop() del self.cuda_ctx print('TrtThread: stopped...') def stop(self): self.running = False self.join()
def main(config): args = parse_args() data_sys = Data_sys() pub = Publish(host=config['Mqtt_pub']) cls_dict = get_cls_dict(args.model.split('_')[-1]) trt_ssd = TrtSSD(args.model, INPUT_HW) img_list = [] # video cap = cv2.VideoCapture('xiasha.avi') i = 0 while cap.isOpened(): frame, img = cap.read() if img is not None: img_list.append([img]) result = trt_ssd.detect(img_list[0], conf_th=0.3) # print(result) data_sys.dataSynchronization(result, img_list, args.model, ['boundary_intrude', None], config['Zone'], config['Channel'][0], config['device_id'], pub, config['Polygon']) img_list = [] i = i + 1 print(i) else: msg = json.dumps({ "nvr_id": config['Zone'], "device_id": config['device_id'], "channel_id": config['Channel'][0] }) pub.send_msg(topic="zs/ai_spxwfx/rtsp/" + config['Zone'] + "/" + config['Channel'][0], msg=msg, Zone=config['Zone'], device_id=config['device_id']) '''
def main(): # Parse arguments and get input args = parse_args() cam = Camera(args) if not cam.isOpened(): raise SystemExit('ERROR: failed to open camera!') # Create NN1 and NN2 models and load into memory cls_dict = get_cls_dict(args.model.split('_')[-1]) trt_ssd = TrtSSD(args.model, INPUT_HW) mtcnn = TrtMtcnn() # Create Preview Window open_window(WINDOW_NAME, 'Camera Preview', cam.img_width, cam.img_height) vis = BBoxVisualization(cls_dict) # Enter Detection Mode while True: # Get Image img = cam.read() out.write(img) nn1_results = [] # Run Neural Networks img, nn1_results, nn2_results, nn3_results = loop_and_detect( img, mtcnn, args.minsize, trt_ssd, conf_th=0.3, vis=vis) # Communicate to Arduino if (nn1_results != []): img = robot_drive(img, nn1_results) else: serial_port.write("N".encode()) print("N") # Display and save output cv2.imshow(WINDOW_NAME, img) outNN.write(img) # User/Keyboard Input key = cv2.waitKey(1) if key == ord('q'): out.release() outNN.release() break # Clean up and exit cam.release() cv2.destroyAllWindows() serial_port.close()
def main(): args = parse_args() cam = Camera(args) if not cam.get_is_opened(): raise SystemExit('ERROR: failed to open camera!') cls_dict = get_cls_dict(args.model.split('_')[-1]) trt_ssd = TrtSSD(args.model, INPUT_HW) open_window(WINDOW_NAME, 'Camera TensorRT SSD Demo', cam.img_width, cam.img_height) vis = BBoxVisualization(cls_dict) loop_and_detect(cam, trt_ssd, conf_th=0.3, vis=vis) cam.release() cv2.destroyAllWindows()
def main(): # Parse arguments and get input args = parse_args() # Create NN1 and NN2 models cls_dict = get_cls_dict(args.model.split('_')[-1]) trt_ssd = TrtSSD(args.model, INPUT_HW) mtcnn = TrtMtcnn() # Create Preview Window vis = BBoxVisualization(cls_dict) imageNum = 10 # Enter Detection Mode while True: # Get Image imageName = "/home/nvidia/Pictures/test13.jpg" #imageName = "/media/E76F-73E0/Faces/1 (" + str(imageNum) + ").jpg" #imageName = "/home/nvidia/Pictures/Benchmarks/Pedestrians/PennPed000" + str(imageNum) + ".png" imageNum = imageNum + 1 #print(imageName) img = cv2.imread(imageName) cv2.imshow(WINDOW_NAME, img) # Run Neural Networks img, nn1_results, nn2_results, nn3_results = loop_and_detect( img, mtcnn, args.minsize, trt_ssd, conf_th=0.3, vis=vis) # Display Results cv2.imshow(WINDOW_NAME, img) #cv2.waitKey(0) # User/Keyboard Input key = cv2.waitKey(1) if key == 27: # ESC key: quit program break # Clean up and exit cv2.destroyAllWindows() serial_port.close()
def main(args, cam): # args = parse_args() is_open = up() time.sleep(60) # cam = Camera(args) if is_open: cam.open() if not cam.is_opened: sys.exit('Failed to open camera!') cls_dict = get_cls_dict(args.model.split('_')[-1]) trt_ssd = TrtSSD(args.model, INPUT_HW) cam.start() vis = BBoxVisualization(cls_dict) loop_and_detect(cam, trt_ssd, conf_th=0.9, vis=vis) cam.stop() cam.release() cv2.destroyAllWindows()
def main(): args = parse_args() cam = Camera(args) is_open = up() #time.sleep(60) if is_open: cam.open() if not cam.is_opened: sys.exit('Failed to open camera!') cls_dict = get_cls_dict(args.model.split('_')[-1]) trt_ssd = TrtSSD(args.model, INPUT_HW) cam.start() open_window(WINDOW_NAME, args.image_width, args.image_height, 'Camera TensorRT SSD Demo for Jetson Nano') vis = BBoxVisualization(cls_dict) loop_and_detect(cam, trt_ssd, conf_th=0.9, vis=vis) cam.stop() cam.release() cv2.destroyAllWindows()
def main(): args = parse_args() check_args(args) results_file = 'ssd/results_%s_%s.json' % (args.model, args.mode) if args.mode == 'trt': ssd = TrtSSD(args.model, INPUT_HW) else: ssd = TfSSD(args.model, INPUT_HW) jpgs = [j for j in os.listdir(args.imgs_dir) if j.endswith('.jpg')] generate_results(ssd, args.imgs_dir, jpgs, results_file) # Run COCO mAP evaluation # Reference: https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb cocoGt = COCO(args.annotations) cocoDt = cocoGt.loadRes(results_file) imgIds = sorted(cocoGt.getImgIds()) cocoEval = COCOeval(cocoGt, cocoDt, 'bbox') cocoEval.params.imgIds = imgIds cocoEval.evaluate() cocoEval.accumulate() print(cocoEval.summarize())
from utils.ssd_classes import get_cls_dict from utils.ssd import TrtSSD from utils.visualization import BBoxVisualization from utils.display import open_window, set_display, show_fps import time import numpy as np model = "ssd_mobilenet_v2_coco" filename = "./dogs.jpg" conf_th = 0.3 INPUT_HW = (300, 300) cls_dict = get_cls_dict("coco") vis = BBoxVisualization(cls_dict) img = cv2.imread(filename) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) trt_ssd = TrtSSD(model, INPUT_HW) # Kick start the model. for _ in range(20): boxes, confs, clss = trt_ssd.detect(img, conf_th) print([get_cls_dict("coco")[c] for c in clss]) img = vis.draw_bboxes(img, boxes, confs, clss) cv2.imwrite("result.jpg", img[..., ::-1]) times = [] for _ in range(20): start_time = time.time() boxes, confs, clss = trt_ssd.detect(img, conf_th) delta = time.time() - start_time times.append(delta) mean_delta = np.array(times).mean() fps = 1 / mean_delta
'''parse args''' parser = argparse.ArgumentParser() parser = add_camera_args(parser) parser.add_argument('--model', type=str, default='ssd_mobilenet_v1_digger') parser.add_argument('--image_resize', default=300, type=int) parser.add_argument('--det_conf_thresh', default=0.8, type=float) parser.add_argument('--seq_dir', default="sequence/") parser.add_argument('--sort_max_age', default=5, type=int) parser.add_argument('--sort_min_hit', default=3, type=int) return parser.parse_args() if __name__ == "__main__": args = parse_args() trt_ssd = TrtSSD(args.model, (args.image_resize, args.image_resize)) mot_tracker = Sort(args.sort_max_age, args.sort_min_hit) video = '/home/mengjun/xianlu/data/digger.mp4' cap = cv2.VideoCapture(video) colours = np.random.rand(32, 3) * 255 fps = 0.0 tic = time.time() while cap.isOpened(): ret, frame = cap.read() boxes, confs, clss = trt_ssd.detect(frame, args.det_conf_thresh) #print(boxes, confs, clss) if len(boxes) != 0: result = [] for bb, cf, cl in zip(boxes, confs, clss):