From 4a5c81a9c3ffe76321d7ea559e8321a1d78e51f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=80=E5=96=B5?= Date: Fri, 12 Jun 2026 15:54:55 +0800 Subject: [PATCH] =?UTF-8?q?=E4=B9=8B=E5=89=8D=E7=9A=84=E8=84=9A=E6=9C=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- infer_camera.py | 371 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 371 insertions(+) create mode 100644 infer_camera.py diff --git a/infer_camera.py b/infer_camera.py new file mode 100644 index 0000000..f779f30 --- /dev/null +++ b/infer_camera.py @@ -0,0 +1,371 @@ +import os +import cv2 +import numpy as np +import time +import threading +import queue +import socket +import json +import sys +from ais_bench.infer.interface import InferSession + +# ================== 导入控制模块 ===================== +from drone_controller.base_control import DroneController +from drone_controller.pid_counter import PID +# =================================================== + +class VideoStreaming: + """ 📡 视频流采集与后台图传模块（支持 RTSP 自动重连） """ + + def __init__(self, video_cfg, net_cfg): + self.video_source = video_cfg.get("video_source", 0) + self.jpeg_quality = video_cfg.get("jpeg_quality", 75) + self.gs_ip = net_cfg.get("ground_station_ip", "127.0.0.1") + self.udp_port = net_cfg.get("udp_port", 9999) + + # RTSP 特殊参数 + self.is_rtsp = isinstance(self.video_source, str) and self.video_source.startswith("rtsp://") + self.cap = None + self.frame_buffer = queue.Queue(maxsize=1) # 只保留最新一帧 + self.is_running = True + self.capture_thread = None + + # 启动后台采集线程 + self._start_capture_thread() + + # 图传发送队列与线程（原样保留） + self.tx_queue = queue.Queue(maxsize=2) + self.sender_thread = threading.Thread(target=self._udp_stream_sender_worker, daemon=True) + self.sender_thread.start() + + def _start_capture_thread(self): + """ 独立线程负责从摄像头/RTSP 读取帧，避免主循环阻塞 """ + self.capture_thread = threading.Thread(target=self._capture_worker, daemon=True) + self.capture_thread.start() + + def _capture_worker(self): + """ 后台不断读取帧，存入 frame_buffer """ + while self.is_running: + # 如果摄像头未打开或已断开，尝试连接 + if self.cap is None or not self.cap.isOpened(): + self._open_camera() + if self.cap is None or not self.cap.isOpened(): + time.sleep(1) # 等待重连 + continue + + ret, frame = self.cap.read() + if not ret: + # 读取失败，关闭并标记需重连 + print("⚠️ RTSP 读取失败，准备重连...") + self.cap.release() + self.cap = None + continue + + # 清空旧帧，只保留最新一帧 + while not self.frame_buffer.empty(): + try: + self.frame_buffer.get_nowait() + except queue.Empty: + break + self.frame_buffer.put(frame) + + if self.cap is not None: + self.cap.release() + + def _open_camera(self): + """ 打开摄像头或 RTSP 流，并设置低延迟参数 """ + try: + print(f"📷 正在连接视频源: {self.video_source}") + self.cap = cv2.VideoCapture(self.video_source) + + if self.is_rtsp: + # 设置 RTSP 缓冲区大小（只保留 1 帧，降低延迟） + self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) + # 设置读取超时（OpenCV 没有直接超时参数，但可尝试设置后端属性） + # 某些平台支持以下设置： + self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('H', '2', '6', '4')) + # 强制使用 tcp 协议（某些 rtsp 默认 udp 会丢包严重） + # 如果 rtsp 地址支持，可以改成 "rtsp://...?tcp" + if "?" not in self.video_source: + self.video_source += "?tcp" + # 注意：有些相机需要重新初始化，这里简单打印提示 + print("💡 建议使用 TCP 传输: 在 RTSP URL 后添加 ?tcp") + else: + # USB 摄像头也可设置缓冲区 + self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) + + # 可选：强制指定分辨率（根据你的相机调整） + # self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640) + # self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480) + + if self.cap.isOpened(): + print("✅ 视频源连接成功") + else: + print("❌ 视频源打开失败") + except Exception as e: + print(f"❌ 打开视频源异常: {e}") + self.cap = None + + def read_frame(self): + """ 主循环调用：从缓冲区获取最新一帧（非阻塞） """ + try: + frame = self.frame_buffer.get_nowait() + return True, frame + except queue.Empty: + # 没有新帧时返回 False，主循环可根据需要 sleep + return False, None + + def push_to_stream(self, frame): + """ 原图传推流逻辑保持不变 """ + try: + send_frame = cv2.resize(frame, (640, 480)) + self.tx_queue.put_nowait(send_frame) + except queue.Full: + pass + + def _udp_stream_sender_worker(self): + """ 原图传发送线程（未改动） """ + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + server_address = (self.gs_ip, self.udp_port) + print(f"📡 图传后台线程已启动，目标地面站 -> {self.gs_ip}:{self.udp_port}") + + while self.is_running: + try: + frame_to_send = self.tx_queue.get(timeout=1.0) + if frame_to_send is None: + break + result, img_encode = cv2.imencode('.jpg', frame_to_send, + [cv2.IMWRITE_JPEG_QUALITY, self.jpeg_quality]) + if not result: + continue + data = img_encode.tobytes() + if len(data) > 65000: + continue + sock.sendto(data, server_address) + except queue.Empty: + continue + except Exception as e: + print(f"❌ 图传线程异常: {e}") + time.sleep(0.1) + sock.close() + + def release(self): + """ 释放资源 """ + self.is_running = False + if self.capture_thread is not None: + self.capture_thread.join(timeout=2.0) + if self.cap is not None: + self.cap.release() + self.tx_queue.put(None) + +class YOLO26UAVInfer: + """ 🚀 昇腾 NPU 推理与图像后处理模块 """ + + def __init__(self, model_cfg): + model_path = model_cfg.get("model_path", "./om/yolo26n-balloon.om") + device_id = model_cfg.get("device_id", 0) + self.conf_threshold = model_cfg.get("conf_threshold", 0.25) + + if not os.path.exists(model_path): + raise FileNotFoundError(f"❌ 找不到模型文件: {model_path}") + + self.session = InferSession(device_id, model_path) + self.input_width = 640 + self.input_height = 640 + + def letterbox(self, img, new_shape=(640, 640), color=(114, 114, 114)): + shape = img.shape[:2] # [height, width] + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r))) + dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] + dw /= 2 + dh /= 2 + if shape[::-1] != new_unpad: + img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) + top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) + left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) + img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) + return img, r, (dw, dh) + + def preprocess(self, frame): + """ 图像预处理，输出符合 NPU 输入的 Tensor """ + input_tensor, ratio, dwdh = self.letterbox(frame) + input_tensor = input_tensor.astype(np.float32) / 255.0 + input_tensor = np.transpose(input_tensor, (2, 0, 1)) + input_tensor = np.expand_dims(input_tensor, axis=0) + input_tensor = np.ascontiguousarray(input_tensor) + return input_tensor, ratio, dwdh + + def postprocess(self, outputs, ratio, dwdh, orig_shape): + out = outputs[0] + detections = [] + for i in range(out.shape[1]): + box = out[0, i, :4] + conf = out[0, i, 4] + cls_id = int(out[0, i, 5]) + if conf > self.conf_threshold: + x1, y1, x2, y2 = box + # 基于原图比例还原坐标 + x1 = (x1 - dwdh[0]) / ratio + y1 = (y1 - dwdh[1]) / ratio + x2 = (x2 - dwdh[0]) / ratio + y2 = (y2 - dwdh[1]) / ratio + x1 = max(0, min(x1, orig_shape[1])) + y1 = max(0, min(y1, orig_shape[0])) + x2 = max(0, min(x2, orig_shape[1])) + y2 = max(0, min(y2, orig_shape[0])) + detections.append([x1, y1, x2, y2, conf, cls_id]) + return detections + + def draw_boxes(self, img, detections): + """ 在图像上渲染精准目标框 """ + for det in detections: + x1, y1, x2, y2, conf, cls_id = det + cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) + label = f"Target: {conf:.2f}" + cv2.putText(img, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) + return img + + +class UAVControlLoop: + """ 🛸 无人机飞控闭环控制核心 """ + + def __init__(self, fc_cfg, pid_y_cfg, pid_z_cfg): + connection_string = fc_cfg.get("connection_string", "udpout:127.0.0.1:14550") + baud_rate = fc_cfg.get("baud_rate", 57600) + + # 初始化控制驱动端 + self.uav = DroneController(connection_string=connection_string, baud=baud_rate) + + # 从 JSON 参数动态初始化 PID 调节器 + self.pid_y = PID( + kp=pid_y_cfg["kp"], ki=pid_y_cfg["ki"], kd=pid_y_cfg["kd"], + max_out=pid_y_cfg["max_out"], min_out=pid_y_cfg["min_out"] + ) + self.pid_z = PID( + kp=pid_z_cfg["kp"], ki=pid_z_cfg["ki"], kd=pid_z_cfg["kd"], + max_out=pid_z_cfg["max_out"], min_out=pid_z_cfg["min_out"] + ) + + def start_uav(self): + """ 连接并解锁无人机 """ + self.uav.connect() + if not self.uav.is_armed and self.uav.current_mode not in ['GUIDED', 'OFFBOARD']: + self.uav.arm() + self.uav.print_status() + self.uav.set_mode('OFFBOARD') + + def process_control(self, detections, orig_shape, frame): + """ 根据检测结果执行 MAVLink 闭环控制逻辑 """ + img_center_x = orig_shape[1] / 2 + img_center_y = orig_shape[0] / 2 + current_mode = getattr(self.uav, 'current_mode', 'UNKNOWN') + + if len(detections) > 0: + # 寻找置信度最高的目标 + best_det = max(detections, key=lambda x: x[4]) + x1, y1, x2, y2, conf, cls_id = best_det + target_x = (x1 + x2) / 2 + target_y = (y1 + y2) / 2 + + # 计算归一化误差 + err_x = (target_x - img_center_x) / img_center_x + err_y = (target_y - img_center_y) / img_center_y + + # PID 计算转化为物理速度需求 + vy_cmd = self.pid_y.update(err_x) # 图像X正向 -> 飞机Vy右移 + vz_cmd = self.pid_z.update(err_y) # 图像Y正向 -> 飞机Vz下移 (NED系) + vx_cmd = 0.0 # 保持前后平移为0 + + # 视效叠加 + cv2.line(frame, (int(img_center_x), int(img_center_y)), (int(target_x), int(target_y)), (255, 0, 0), 2) + cv2.putText(frame, f"MAVLink Out -> Err_X: {err_x:+.3f} Err_Y: {err_y:+.3f}", (20, 70), + cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2) + + # 执行机体速度控制 + if current_mode in ['GUIDED', 'OFFBOARD']: + self.uav.send_body_velocity(vx=vx_cmd, vy=vy_cmd, vz=vz_cmd, yaw_rate=0.0) + else: + cv2.putText(frame, f"MANUAL MODE OVERRIDE ({current_mode})", (20, 100), cv2.FONT_HERSHEY_SIMPLEX, 0.6, + (0, 0, 255), 2) + else: + # 目标丢失保护逻辑 + cv2.putText(frame, "MAVLink Out -> Target LOST (HOVER)", (20, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.6, + (0, 0, 255), 2) + + # 清除 PID 积分，防止突跳 + if hasattr(self.pid_y, 'reset'): self.pid_y.reset() + if hasattr(self.pid_z, 'reset'): self.pid_z.reset() + + if current_mode in ['GUIDED', 'OFFBOARD']: + self.uav.send_body_velocity(vx=0.0, vy=0.0, vz=0.0, yaw_rate=0.0) + + +# ==================== 主程序入口 ====================== +def main(): + # 默认加载同目录下的 config.json + config_file = "config.json" + + # 允许在命令行指定 json 路径，例如: python3 infer_camera_modular.py config_outdoor.json + if len(sys.argv) > 1: + config_file = sys.argv[1] + + print(f"📖 正在从 {config_file} 加载系统配置...") + try: + with open(config_file, "r", encoding="utf-8") as f: + cfg = json.load(f) + except Exception as e: + print(f"❌ 读取配置文件失败: {e}") + sys.exit(1) + + # 实例化各子系统并注入配置 + streamer = VideoStreaming(cfg["video"], cfg["network"]) + detector = YOLO26UAVInfer(cfg["model"]) + controller = UAVControlLoop(cfg["flight_control"], cfg["pid_y"], cfg["pid_z"]) + + # 启动无人机连接 + controller.start_uav() + + + print("🚀 NPU 精准推理 [JSON配置动态加载版] 主循环启动...") + + try: + while True: + ret, frame = streamer.read_frame() + if not ret: + time.sleep(0.005) + continue + + orig_shape = frame.shape[:2] + + # 1. 图像预处理 + input_tensor, ratio, dwdh = detector.preprocess(frame) + + # 2. NPU 硬件推理 + t_start = time.time() + outputs = detector.session.infer([input_tensor]) + t_end = time.time() + + fps_pure = 1.0 / (t_end - t_start) + + # 3. 坐标反求与控制计算 + detections = detector.postprocess(outputs, ratio, dwdh, orig_shape) + controller.process_control(detections, orig_shape, frame) + + # 4. 界面渲染 + frame = detector.draw_boxes(frame, detections) + cv2.putText(frame, f"NPU Pure FPS: {fps_pure:.1f}", (20, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), + 2) + + # 5. 推送至后台异步图传队列 + streamer.push_to_stream(frame) + + except KeyboardInterrupt: + print("\n👋 接收到终止信号，安全退出中...") + finally: + streamer.release() + print("程序运行结束。") + + +if __name__ == "__main__": + main() \ No newline at end of file