""" An example that uses TensorRT's Python api to make inferences. """ import ctypes import os import shutil import random import sys import threading import time import cv2 from datalink_serial import datalink import queue import socket import multiprocessing as mp TCP_IP = '0.0.0.0' TCP_PORT = 5005 W_img = 960 H_img = 540 FOV_x = 77 # 水平视场角，度 FOV_y = 44 # 垂直视场角，度 W_real = 0.8 # target实际宽度，米 H_real = 1.7 # target实际高度，米 safe_distance = 1.5 # 跟踪间距 width_up_threshold = 0.5 # 目标宽度上限阈值，比例 width_down_threshold = 0.25 # 目标宽度下限阈值 k = 1.0 # 系数 Kp_dx = 0.5 Kp_dy = 0.5 Kp_dalt = 0.5 Kp_dyaw = 0.3 def gstreamer_pipeline( sensor_id=0, capture_width=1920, capture_height=1080, display_width=W_img, display_height=H_img, framerate=30, flip_method=0, ): return ( "nvarguscamerasrc sensor-id=%d ! " "video/x-raw(memory:NVMM), width=(int)%d, height=(int)%d, framerate=(fraction)%d/1 ! " "nvvidconv flip-method=%d ! " "video/x-raw, width=(int)%d, height=(int)%d, format=(string)BGRx ! " "videoconvert ! " "video/x-raw, format=(string)BGR ! appsink" % ( sensor_id, capture_width, capture_height, framerate, flip_method, display_width, display_height, ) ) window_title = "CSI Camera" # To flip the image, modify the flip_method parameter (0 and 2 are the most common) print(gstreamer_pipeline(flip_method=2)) video_capture = cv2.VideoCapture(gstreamer_pipeline(flip_method=2), cv2.CAP_GSTREAMER) dl = datalink() # 实例化 datalink import numpy as np import pycuda.autoinit import pycuda.driver as cuda import tensorrt as trt CONF_THRESH = 0.5 IOU_THRESHOLD = 0.4 LEN_ALL_RESULT = 38001 LEN_ONE_RESULT = 38 def get_img_path_batches(batch_size, img_dir): ret = [] batch = [] for root, dirs, files in os.walk(img_dir): for name in files: if len(batch) == batch_size: ret.append(batch) batch = [] batch.append(os.path.join(root, name)) if len(batch) > 0: ret.append(batch) return ret def plot_one_box(x, img, color=None, label=None, line_thickness=None): """ description: Plots one bounding box on image img, this function comes from YoLov5 project. param: x: a box likes [x1,y1,x2,y2] img: a opencv image object color: color to draw rectangle, such as (0,255,0) label: str line_thickness: int return: no return """ tl = ( line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 ) # line/font thickness color = color or [random.randint(0, 255) for _ in range(3)] c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) if label: tf = max(tl - 1, 1) # font thickness t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled cv2.putText( img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA, ) class YoLov5TRT(object): """ description: A YOLOv5 class that warps TensorRT ops, preprocess and postprocess ops. """ def __init__(self, engine_file_path): # Create a Context on this device, self.ctx = cuda.Device(0).make_context() stream = cuda.Stream() TRT_LOGGER = trt.Logger(trt.Logger.INFO) runtime = trt.Runtime(TRT_LOGGER) # Deserialize the engine from file with open(engine_file_path, "rb") as f: engine = runtime.deserialize_cuda_engine(f.read()) context = engine.create_execution_context() host_inputs = [] cuda_inputs = [] host_outputs = [] cuda_outputs = [] bindings = [] for binding in engine: print('bingding:', binding, engine.get_binding_shape(binding)) size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size dtype = trt.nptype(engine.get_binding_dtype(binding)) # Allocate host and device buffers host_mem = cuda.pagelocked_empty(size, dtype) cuda_mem = cuda.mem_alloc(host_mem.nbytes) # Append the device buffer to device bindings. bindings.append(int(cuda_mem)) # Append to the appropriate list. if engine.binding_is_input(binding): self.input_w = engine.get_binding_shape(binding)[-1] self.input_h = engine.get_binding_shape(binding)[-2] host_inputs.append(host_mem) cuda_inputs.append(cuda_mem) else: host_outputs.append(host_mem) cuda_outputs.append(cuda_mem) # Store self.stream = stream self.context = context self.engine = engine self.host_inputs = host_inputs self.cuda_inputs = cuda_inputs self.host_outputs = host_outputs self.cuda_outputs = cuda_outputs self.bindings = bindings self.batch_size = engine.max_batch_size def infer(self, input_image_path): threading.Thread.__init__(self) # Make self the active context, pushing it on top of the context stack. self.ctx.push() self.input_image_path = input_image_path # Restore stream = self.stream context = self.context engine = self.engine host_inputs = self.host_inputs cuda_inputs = self.cuda_inputs host_outputs = self.host_outputs cuda_outputs = self.cuda_outputs bindings = self.bindings # Do image preprocess batch_origin_h = [] batch_origin_w = [] batch_input_image = np.empty(shape=[self.batch_size, 3, self.input_h, self.input_w]) input_image, image_raw, origin_h, origin_w = self.preprocess_image(input_image_path) batch_origin_h.append(origin_h) batch_origin_w.append(origin_w) np.copyto(batch_input_image, input_image) batch_input_image = np.ascontiguousarray(batch_input_image) # Copy input image to host buffer np.copyto(host_inputs[0], batch_input_image.ravel()) start = time.time() # Transfer input data to the GPU. cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream) # Run inference. context.execute_async(batch_size=self.batch_size, bindings=bindings, stream_handle=stream.handle) # Transfer predictions back from the GPU. cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream) # Synchronize the stream stream.synchronize() end = time.time() # Remove any context from the top of the context stack, deactivating it. self.ctx.pop() # Here we use the first row of output in that batch_size = 1 output = host_outputs[0] # Do postprocess result_boxes, result_scores, result_classid = self.post_process( output, origin_h, origin_w) # Draw rectangles and labels on the original image for j in range(len(result_boxes)): box = result_boxes[j] if result_scores[j]>0.5 and categories[int(result_classid[j])]=='person': # 计算相机的焦距 f，单位为像素 f_x = (W_img / 2) / np.tan(np.radians(FOV_x / 2)) f_y = (H_img / 2) / np.tan(np.radians(FOV_y / 2)) x1=box[0] y1=box[1] x2=box[2] y2=box[3] W_qr = x2 - x1 H_qr = y2 - y1 cx_qr = (x1 + x2) / 2 cy_qr = (y1 + y2) / 2 cx_img = W_img / 2 cy_img = H_img / 2 dx = cx_qr - cx_img dy = cy_qr - cy_img angle_x_rad = k * np.arctan(dx / f_x) dz_m = k * (W_real * f_x) / W_qr dy_m = k * (dy / f_y) * dz_m dx_m = k * (dx / f_x) * dz_m dx_1 = dz_m - safe_distance dy_1 = dx_m d_alt_1 = -dy_m d_yaw = angle_x_rad if W_qr > width_up_threshold * W_img or W_qr < width_down_threshold * W_img: dx_1 = 0.0 dl.set_pose(Kp_dx * dx_1, Kp_dy * dy_1, Kp_dalt * d_alt_1, Kp_dyaw * d_yaw) plot_one_box( box, image_raw, label="{}:{:.2f}".format( categories[int(result_classid[j])], result_scores[j] ), ) return image_raw, end - start def destroy(self): # Remove any context from the top of the context stack, deactivating it. self.ctx.pop() def get_raw_image(self, image_path_batch): """ description: Read an image from image path """ for img_path in image_path_batch: yield cv2.imread(img_path) def get_raw_image_zeros(self, image_path_batch=None): """ description: Ready data for warmup """ for _ in range(self.batch_size): yield np.zeros([self.input_h, self.input_w, 3], dtype=np.uint8) def preprocess_image(self, raw_bgr_image): """ description: Convert BGR image to RGB, resize and pad it to target size, normalize to [0,1], transform to NCHW format. param: input_image_path: str, image path return: image: the processed image image_raw: the original image h: original height w: original width """ image_raw = raw_bgr_image h, w, c = image_raw.shape image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB) # Calculate widht and height and paddings r_w = self.input_w / w r_h = self.input_h / h if r_h > r_w: tw = self.input_w th = int(r_w * h) tx1 = tx2 = 0 ty1 = int((self.input_h - th) / 2) ty2 = self.input_h - th - ty1 else: tw = int(r_h * w) th = self.input_h tx1 = int((self.input_w - tw) / 2) tx2 = self.input_w - tw - tx1 ty1 = ty2 = 0 # Resize the image with long side while maintaining ratio image = cv2.resize(image, (tw, th)) # Pad the short side with (128,128,128) image = cv2.copyMakeBorder( image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, None, (128, 128, 128) ) image = image.astype(np.float32) # Normalize to [0,1] image /= 255.0 # HWC to CHW format: image = np.transpose(image, [2, 0, 1]) # CHW to NCHW format image = np.expand_dims(image, axis=0) # Convert the image to row-major order, also known as "C order": image = np.ascontiguousarray(image) return image, image_raw, h, w def xywh2xyxy(self, origin_h, origin_w, x): """ description: Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right param: origin_h: height of original image origin_w: width of original image x: A boxes numpy, each row is a box [center_x, center_y, w, h] return: y: A boxes numpy, each row is a box [x1, y1, x2, y2] """ y = np.zeros_like(x) r_w = self.input_w / origin_w r_h = self.input_h / origin_h if r_h > r_w: y[:, 0] = x[:, 0] - x[:, 2] / 2 y[:, 2] = x[:, 0] + x[:, 2] / 2 y[:, 1] = x[:, 1] - x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2 y[:, 3] = x[:, 1] + x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2 y /= r_w else: y[:, 0] = x[:, 0] - x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2 y[:, 2] = x[:, 0] + x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2 y[:, 1] = x[:, 1] - x[:, 3] / 2 y[:, 3] = x[:, 1] + x[:, 3] / 2 y /= r_h return y def post_process(self, output, origin_h, origin_w): """ description: postprocess the prediction param: output: A numpy likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...] origin_h: height of original image origin_w: width of original image return: result_boxes: finally boxes, a boxes numpy, each row is a box [x1, y1, x2, y2] result_scores: finally scores, a numpy, each element is the score correspoing to box result_classid: finally classid, a numpy, each element is the classid correspoing to box """ # Get the num of boxes detected num = int(output[0]) # Reshape to a two dimentional ndarray pred = np.reshape(output[1:], (-1, LEN_ONE_RESULT))[:num, :] pred = pred[:, :6] # Do nms boxes = self.non_max_suppression(pred, origin_h, origin_w, conf_thres=CONF_THRESH, nms_thres=IOU_THRESHOLD) result_boxes = boxes[:, :4] if len(boxes) else np.array([]) result_scores = boxes[:, 4] if len(boxes) else np.array([]) result_classid = boxes[:, 5] if len(boxes) else np.array([]) return result_boxes, result_scores, result_classid def bbox_iou(self, box1, box2, x1y1x2y2=True): """ description: compute the IoU of two bounding boxes param: box1: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) box2: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) x1y1x2y2: select the coordinate format return: iou: computed iou """ if not x1y1x2y2: # Transform from center and width to exact coordinates b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2 b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2 b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2 b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2 else: # Get the coordinates of bounding boxes b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3] b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3] # Get the coordinates of the intersection rectangle inter_rect_x1 = np.maximum(b1_x1, b2_x1) inter_rect_y1 = np.maximum(b1_y1, b2_y1) inter_rect_x2 = np.minimum(b1_x2, b2_x2) inter_rect_y2 = np.minimum(b1_y2, b2_y2) # Intersection area inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, None) * \ np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, None) # Union Area b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1) b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1) iou = inter_area / (b1_area + b2_area - inter_area + 1e-16) return iou def non_max_suppression(self, prediction, origin_h, origin_w, conf_thres=0.5, nms_thres=0.4): """ description: Removes detections with lower object confidence score than 'conf_thres' and performs Non-Maximum Suppression to further filter detections. param: prediction: detections, (x1, y1, x2, y2, conf, cls_id) origin_h: original image height origin_w: original image width conf_thres: a confidence threshold to filter detections nms_thres: a iou threshold to filter detections return: boxes: output after nms with the shape (x1, y1, x2, y2, conf, cls_id) """ # Get the boxes that score > CONF_THRESH boxes = prediction[prediction[:, 4] >= conf_thres] # Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2] boxes[:, :4] = self.xywh2xyxy(origin_h, origin_w, boxes[:, :4]) # clip the coordinates boxes[:, 0] = np.clip(boxes[:, 0], 0, origin_w -1) boxes[:, 2] = np.clip(boxes[:, 2], 0, origin_w -1) boxes[:, 1] = np.clip(boxes[:, 1], 0, origin_h -1) boxes[:, 3] = np.clip(boxes[:, 3], 0, origin_h -1) # Object confidence confs = boxes[:, 4] # Sort by the confs boxes = boxes[np.argsort(-confs)] # Perform non-maximum suppression keep_boxes = [] while boxes.shape[0]: large_overlap = self.bbox_iou(np.expand_dims(boxes[0, :4], 0), boxes[:, :4]) > nms_thres label_match = boxes[0, -1] == boxes[:, -1] # Indices of boxes with lower confidence scores, large IOUs and matching labels invalid = large_overlap & label_match keep_boxes += [boxes[0]] boxes = boxes[~invalid] boxes = np.stack(keep_boxes, 0) if len(keep_boxes) else np.array([]) return boxes class inferThread(threading.Thread): def __init__(self, yolov5_wrapper): threading.Thread.__init__(self) self.yolov5_wrapper = yolov5_wrapper def infer(self, frame): batch_image_raw, use_time = self.yolov5_wrapper.infer(frame) return batch_image_raw, use_time class warmUpThread(threading.Thread): def __init__(self, yolov5_wrapper): threading.Thread.__init__(self) self.yolov5_wrapper = yolov5_wrapper def run(self): batch_image_raw, use_time = self.yolov5_wrapper.infer(self.yolov5_wrapper.get_raw_image_zeros()) print('warm_up->{}, time->{:.2f}ms'.format(batch_image_raw[0].shape, use_time * 1000)) class TranstreamThread(threading.Thread): def __init__(self, TCP_IP, TCP_PORT): super().__init__() self.TCP_IP = TCP_IP self.TCP_PORT = TCP_PORT self.sock = None self.conn = None self.addr = None self.running = True self.connected = False self.queue = queue.Queue() def run(self): self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) # 允许端口重用 self.sock.bind((self.TCP_IP, self.TCP_PORT)) self.sock.listen(1) print(f"服务器已启动，监听 {self.TCP_IP}:{self.TCP_PORT}") while self.running: print("等待客户端连接...") try: self.conn, self.addr = self.sock.accept() # 每次 accept 一个新的客户端 print(f"客户端已连接：{self.addr}") self.connected = True # 设置 TCP_NODELAY self.conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) # 处理当前客户端通信 while self.running and self.connected: try: frame = self.queue.get(timeout=1) # 等待图像帧 self._send_frame(frame) except queue.Empty: continue except socket.error as e: print(f"accept() 出错：{e}") self.connected = False continue def _send_frame(self, frame): try: size_data = len(frame).to_bytes(4, 'big') self.conn.sendall(size_data) self.conn.sendall(frame) except (socket.error, BrokenPipeError) as e: print(f"客户端断开连接：{e}") self.connected = False # 客户端断开后标记为未连接 self._close_connection() def send_frame(self, frame): if self.connected: self.queue.put(frame) return True return False def _close_connection(self): if self.conn: try: self.conn.close() except: pass self.conn = None def stop(self): self.running = False self.connected = False self._close_connection() if self.sock: self.sock.close() if __name__ == "__main__": # load custom plugin and engine PLUGIN_LIBRARY = "build/libmyplugins.so" engine_file_path = "build/person.engine" if len(sys.argv) > 1: engine_file_path = sys.argv[1] if len(sys.argv) > 2: PLUGIN_LIBRARY = sys.argv[2] ctypes.CDLL(PLUGIN_LIBRARY) # load coco labels categories = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"] if os.path.exists('output/'): shutil.rmtree('output/') os.makedirs('output/') # a YoLov5TRT instance yolov5_wrapper = YoLov5TRT(engine_file_path) try: # create a new thread to do warm_up thread1 = inferThread(yolov5_wrapper) thread1.start() thread1.join() thread2 = TranstreamThread(TCP_IP, TCP_PORT) # 这句是我们要加入的 thread2.start() # 这句是我们要加入的 print('data link is starting...\n') drone_thread=threading.Thread(target=dl.drone) heartbeat_thread=threading.Thread(target=dl.heartbeat) drone_thread.start() heartbeat_thread.start() if video_capture.isOpened(): try: window_handle = cv2.namedWindow(window_title, cv2.WINDOW_AUTOSIZE) while True: ret_val, frame = video_capture.read() img, use_time = thread1.infer(frame) # Check to see if the user closed the window # Under GTK+ (Jetson Default), WND_PROP_VISIBLE does not work correctly. Under Qt it does # GTK - Substitute WND_PROP_AUTOSIZE to detect if window has been closed by user if cv2.getWindowProperty(window_title, cv2.WND_PROP_AUTOSIZE) >= 0: cv2.imshow(window_title, img) encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 70] result, img_encoded = cv2.imencode('.jpg', frame, encode_param) data = img_encoded.tobytes() thread2.send_frame(data) # 通过队列异步发送 else: break keyCode = cv2.waitKey(10) & 0xFF # Stop the program on the ESC key or 'q' if keyCode == 27 or keyCode == ord('q'): break finally: video_capture.release() cv2.destroyAllWindows() else: print("Error: Unable to open camera") finally: # destroy the instance yolov5_wrapper.destroy() drone_thread.join() heartbeat_thread.join() dl.com.close()