print("Distance detected by supersonic:", distance) if distance != None and distance < 20: robot.backward(0.35) time.sleep(0.5) robot.stop() continue # Capture frame-by-frame frame = cap.read( ) # ret = 1 if the video is captured; frame is the image # Our operations on the frame come here img = cv2.flip(frame, 1) # flip left-right img = cv2.flip(img, 0) # flip up-down boxes, confs, clss = trt_yolo.detect(img, 0.8) print("boxes:", boxes) print("confs:", confs) print("clss:", clss) # Find target target_clss_index = clss < 3 target_boxes = boxes[target_clss_index] if len(target_boxes) > 0: print("Found target ", len(target_boxes), " and set barrier_center = -1") barrier_center = -1 # Find barriers barriers_clss_index = clss > 2
class YOLOv4: def __init__(self, frame_info_pool=None, detection_result_pool=None, final_result_pool=None, model='yolov4-416', dataset='obstacle'): """ :param model: model name :param category_num: """ self.results = dict() self.model_name = model self.dataset = dataset self.frame_info_pool = frame_info_pool self.detection_result_pool = detection_result_pool self.final_result_pool = final_result_pool if dataset == "coco": category_num = 80 else: category_num = 15 self.cls_dict = get_cls_dict(category_num) yolo_dim = model.split('-')[-1] if 'x' in yolo_dim: dim_split = yolo_dim.split('x') if len(dim_split) != 2: raise SystemExit('ERROR: bad yolo_dim (%s)!' % yolo_dim) w, h = int(dim_split[0]), int(dim_split[1]) else: h = w = int(yolo_dim) if h % 32 != 0 or w % 32 != 0: raise SystemExit('ERROR: bad yolo_dim (%s)!' % yolo_dim) self.model = TrtYOLO(model, (h, w), category_num) PrintLog.i("Object detection model is loaded - {}\t{}".format( model, dataset)) def inference_by_image(self, image, confidence_threshold=0.1): """ :param image: input image :param confidence_threshold: confidence/score threshold for object detection. :return: bounding box(x1, y1, x2, y2), scores, classes """ boxes, scores, classes = self.model.detect(image, confidence_threshold) results = [] for box, score, cls in zip(boxes, scores, classes): results.append({ "label": [{ "description": self.cls_dict.get(int(cls), 'CLS{}'.format(int(cls))), "score": float(score), "class_idx": int(cls) }], "position": { "x": int(box[0]), "y": int(box[1]), "w": int(box[2] - box[0]), "h": int(box[3] - box[1]) } }) return results def run(self): PrintLog.i("Frame Number\tTimestamp") while True: if len(self.frame_info_pool) > 0: frame_info = self.frame_info_pool.pop(0) result = self.inference_by_image(frame_info["frame"]) detection_result = dict() detection_result["cam_address"] = frame_info["cam_address"] detection_result["timestamp"] = frame_info["timestamp"] detection_result["frame_number"] = frame_info["frame_number"] detection_result["results"] = [] detection_result["results"].append( {"detection_result": result}) self.detection_result_pool.append(detection_result)
def main(): # initialize camera class cam = Camera() if not cam.is_opened(): raise SystemExit('EROR: failed to open camera!') cls_dict = alprClassNames() # Yolo dimensions (416x416) yolo_dim = 416 h = w = int(yolo_dim) # Initialized model and tools cwd = os.getcwd() model_yolo = str(cwd) + '/weights/yolov4-tiny-416.trt' model_crnn = str(cwd) + '/weights/crnn.pth' trt_yolo = TrtYOLO(model_yolo, (h,w), category_num=1) # category number is number of classes crnn = alpr.AutoLPR(decoder='bestPath', normalise=True) crnn.load(crnn_path=model_crnn) open_window(WINDOW_NAME, TITLE, cam.img_width, cam.img_height) vis = BBoxVisualization(cls_dict) # Loop and detect full_scrn = False fps = 0.0 tic = time.time() while True: if cv2.getWindowProperty(WINDOW_NAME, 0) < 0: break img = cam.read() if img is None: break # Detect car plate boxes confs, clss = trt_yolo.detect(img, conf_th=0.5) # Crop and preprocess car plate cropped = vis.crop_plate(img, boxes, confs, clss) # Recognize car plate lp_plate = '' fileLocate = str(cwd) + '/detections/detection1.jpg' if os.path.exists(fileLocate): lp_plate = lpr.predict(fileLocate) # Draw boxes and fps img = vis.draw_bboxes(img, boxes, confs, clss, lp=lp_plate) img = show_fps(img, fps) # Show image cv2.imshow(WINDOW_NAME, img) # Calculate fps toc = time.time() curr_fps = 1.0 / (toc - tic) fps = curr_fps if fps == 0.0 else (fps*0.95 + curr_fps*0.05) tic = toc # Exit key key = cv2.waitKey(1) if key == 27: # ESC key: quit program break # Release capture and destroy all windows cam.release() cv2.destroyAllWindows()
class yolov4(object): def __init__(self): """ Constructor """ self.bridge = CvBridge() self.init_params() self.init_yolo() self.cuda_ctx = cuda.Device(0).make_context() self.trt_yolo = TrtYOLO((self.model_path + self.model), (self.h, self.w), self.category_num) def __del__(self): """ Destructor """ self.cuda_ctx.pop() del self.trt_yolo del self.cuda_ctx def clean_up(self): """ Backup destructor: Release cuda memory """ if self.trt_yolo is not None: self.cuda_ctx.pop() del self.trt_yolo del self.cuda_ctx def init_params(self): """ Initializes ros parameters """ rospack = rospkg.RosPack() package_path = rospack.get_path("yolov4_trt_ros") self.video_topic = rospy.get_param("/video_topic", "/video_source/raw") self.model = rospy.get_param("/model", "yolov4") self.model_path = rospy.get_param("/model_path", package_path + "/yolo/") self.input_shape = rospy.get_param("/input_shape", "416") self.category_num = rospy.get_param("/category_number", 80) self.conf_th = rospy.get_param("/confidence_threshold", 0.5) self.show_img = rospy.get_param("/show_image", True) self.image_sub = rospy.Subscriber(self.video_topic, Image, self.img_callback, queue_size=1, buff_size=1920 * 1080 * 3) self.detection_pub = rospy.Publisher("detections", Detector2DArray, queue_size=1) self.overlay_pub = rospy.Publisher("/result/overlay", Image, queue_size=1) def init_yolo(self): """ Initialises yolo parameters required for trt engine """ if self.model.find('-') == -1: self.model = self.model + "-" + self.input_shape yolo_dim = self.model.split('-')[-1] if 'x' in yolo_dim: dim_split = yolo_dim.split('x') if len(dim_split) != 2: raise SystemExit('ERROR: bad yolo_dim (%s)!' % yolo_dim) self.w, self.h = int(dim_split[0]), int(dim_split[1]) else: self.h = self.w = int(yolo_dim) if self.h % 32 != 0 or self.w % 32 != 0: raise SystemExit('ERROR: bad yolo_dim (%s)!' % yolo_dim) cls_dict = get_cls_dict(self.category_num) self.vis = BBoxVisualization(cls_dict) def img_callback(self, ros_img): """Continuously capture images from camera and do object detection """ tic = time.time() # converts from ros_img to cv_img for processing try: cv_img = self.bridge.imgmsg_to_cv2(ros_img, desired_encoding="bgr8") rospy.logdebug("ROS Image converted for processing") except CvBridgeError as e: rospy.loginfo("Failed to convert image %s", str(e)) if cv_img is not None: boxes, confs, clss = self.trt_yolo.detect(cv_img, self.conf_th) cv_img = self.vis.draw_bboxes(cv_img, boxes, confs, clss) toc = time.time() fps = 1.0 / (toc - tic) self.publisher(boxes, confs, clss) if self.show_img: cv_img = show_fps(cv_img, fps) cv2.imshow("YOLOv4 DETECTION RESULTS", cv_img) cv2.waitKey(1) # converts back to ros_img type for publishing try: overlay_img = self.bridge.cv2_to_imgmsg(cv_img, encoding="passthrough") rospy.logdebug("CV Image converted for publishing") self.overlay_pub.publish(overlay_img) except CvBridgeError as e: rospy.loginfo("Failed to convert image %s", str(e)) def publisher(self, boxes, confs, clss): """ Publishes to detector_msgs Parameters: boxes (List(List(int))) : Bounding boxes of all objects confs (List(double)) : Probability scores of all objects clss (List(int)) : Class ID of all classes """ detection2d = Detector2DArray() detection = Detector2D() for i in range(len(boxes)): # boxes : xmin, ymin, xmax, ymax for _ in boxes: detection.header.stamp = rospy.Time.now() detection.header.frame_id = "camera" # change accordingly detection.results.id = clss[i] detection.results.score = confs[i] detection.bbox.center.x = (boxes[i][2] - boxes[i][0]) / 2 detection.bbox.center.y = (boxes[i][3] - boxes[i][1]) / 2 detection.bbox.center.theta = 0.0 # change if required detection.bbox.size_x = abs(boxes[i][0] - boxes[i][2]) detection.bbox.size_y = abs(boxes[i][1] - boxes[i][3]) detection2d.detections.append(detection) self.detection_pub.publish(detection2d)
tic = time.time() def closedCapCallback(): global cap print("no connection in the stream, reconnecting") cap = cv2.VideoCapture('rtsp://localhost:8554/stream_input') sleep(0.5) while True: if not cap.isOpened(): print('error while subprocess not running') closedCapCallback() else: _,frame = cap.read() if not frame is None: boxes, confs, clss = trt_yolo.detect(frame, 0.3) frame = vis.draw_bboxes(frame, boxes, confs, clss) cv2.putText(frame, "FPS: " + str(round(fps, 2)), (10, 50), font, 3, (0, 0, 0), 3) toc = time.time() curr_fps = 1.0 / (toc - tic) fps = curr_fps if fps == 0.0 else (fps * 0.95 + curr_fps * 0.05) tic = toc ret_toRTSP, frame_toRTSP = cv2.imencode('.png', frame) process.stdin.write(frame_toRTSP.tobytes()) else: print('error while subprocess is running') closedCapCallback()