diff --git a/infer_camera_direct.py b/infer_camera_direct.py index 26306f8..8d88062 100644 --- a/infer_camera_direct.py +++ b/infer_camera_direct.py @@ -5,7 +5,7 @@ import time from ais_bench.infer.interface import InferSession # ==================== 硬编码配置区域 ==================== -MODEL_PATH = "./yolo26_8t.om" # 你的 YOLO26 模型路径 +MODEL_PATH = "./om/yolo26n-balloon.om" # 你的 YOLO26 模型路径 CONF_THRESHOLD = 0.25 # 置信度过滤阈值 (YOLO26 推荐 0.25) CAMERA_INDEX = 0 # 摄像头索引 (默认0，如果打不开可以尝试1或2) diff --git a/infer_video.py b/infer_video.py new file mode 100644 index 0000000..5173558 --- /dev/null +++ b/infer_video.py @@ -0,0 +1,119 @@ +import os +import cv2 +import numpy as np +import time +from ais_bench.infer.interface import InferSession + +# ==================== 模拟测试硬编码配置 ==================== +MODEL_PATH = "./yolo26_8t.om" # 你的 YOLO26 模型路径 +VIDEO_PATH = "./drone_test.mp4" # 用于模拟相机的测试视频文件 +CONF_THRESHOLD = 0.25 # 置信度阈值 + + +# ============================================================ + +class YOLO26UAVInfer: + def __init__(self, model_path, device_id=0): + 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] + 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] + left, top = dw // 2, dh // 2 + right, bottom = dw - left, dh - top + if shape[::-1] != new_unpad: + img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) + img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) + return img, r, (left, top) + + def preprocess_frame(self, img_bgr): + self.original_height, self.original_width = img_bgr.shape[:2] + img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) + img_resized, self.ratio, (self.pad_left, self.pad_top) = self.letterbox( + img_rgb, (self.input_width, self.input_height) + ) + img_input = (img_resized.astype(np.float32) / 255.0).astype(np.float32) + img_input = np.ascontiguousarray(np.transpose(img_input, (2, 0, 1))[np.newaxis, :, :, :]) + return img_input + + def infer(self, img_input): + return self.session.infer([img_input]) + + def postprocess(self, infer_result, conf_threshold): + detections = [] + output = infer_result[0] if isinstance(infer_result, (list, tuple)) else infer_result + if output.ndim == 3: output = output[0] + for det in output: + confidence = det[4] + if confidence < conf_threshold: continue + x1_scale, y1_scale, x2_scale, y2_scale = det[0], det[1], det[2], det[3] + cls_id = int(det[5]) + x1 = max(0, min((x1_scale - self.pad_left) / self.ratio, self.original_width)) + y1 = max(0, min((y1_scale - self.pad_top) / self.ratio, self.original_height)) + x2 = max(0, min((x2_scale - self.pad_left) / self.ratio, self.original_width)) + y2 = max(0, min((y2_scale - self.pad_top) / self.ratio, self.original_height)) + detections.append([x1, y1, x2, y2, confidence, cls_id]) + return detections + + +if __name__ == "__main__": + print("====== 启动香橙派无人机检测算法模拟测试 ======") + yolo_infer = YOLO26UAVInfer(MODEL_PATH) + + # 用视频文件模拟机载相机输入 + cap = cv2.VideoCapture(VIDEO_PATH) + if not cap.isOpened(): + print(f"错误：无法打开模拟视频文件 {VIDEO_PATH}") + exit() + + frame_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + frame_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + img_center_x, img_center_y = frame_w / 2, frame_h / 2 + + print(f"成功加载视频，分辨率: {frame_w}x{frame_h} | 正在进行 NPU 循环推理...") + fps_start_time = time.time() + frame_count = 0 + + while True: + ret, frame = cap.read() + if not ret: + print("视频播放结束或无法读取。") + break + + frame_count += 1 + + # 1. 预处理与 NPU 推理 + input_tensor = yolo_infer.preprocess_frame(frame) + raw_result = yolo_infer.infer(input_tensor) + detections = yolo_infer.postprocess(raw_result, conf_threshold=CONF_THRESHOLD) + + # 2. 模拟数据下发接口（核心优化点） + if len(detections) > 0: + best_det = max(detections, key=lambda x: x[4]) # 追踪置信度最高的目标 + x1, y1, x2, y2, conf, cls_id = best_det + target_x, target_y = (x1 + x2) / 2, (y1 + y2) / 2 + + # 计算标准化像素偏差（-1.0 到 1.0） + err_x = (target_x - img_center_x) / img_center_x + err_y = (target_y - img_center_y) / img_center_y + + # 模拟发送给飞控 + print( + f"[Frame {frame_count:04d}] 检测到目标(类别:{cls_id}/置信度:{conf:.2f}) -> 模拟下发飞控控制量 err_x: {err_x:+.4f}, err_y: {err_y:+.4f}") + else: + print(f"[Frame {frame_count:04d}] 未检测到目标 -> 告知飞控保持悬停对准...") + + # 3. 计算端到端真实 FPS（包含预处理、推理和后处理） + if frame_count % 30 == 0: + end_time = time.time() + avg_fps = 30 / (end_time - fps_start_time) + print(f"----------------> 当前香橙派端到端平均处理速度: {avg_fps:.1f} FPS") + fps_start_time = time.time() + + cap.release() \ No newline at end of file diff --git a/test_video.mp4 b/test_video.mp4 new file mode 100644 index 0000000..a8da134 Binary files /dev/null and b/test_video.mp4 differ diff --git a/uav_ai_companion.py b/uav_ai_companion.py new file mode 100644 index 0000000..4a11985 --- /dev/null +++ b/uav_ai_companion.py @@ -0,0 +1,188 @@ +import os +import cv2 +import numpy as np +import time +from ais_bench.infer.interface import InferSession + +# 引入 MAVLink 核心库 +from pymavlink import mavutil + +# ==================== 进阶硬编码配置区域 ==================== +MODEL_PATH = "./om/yolo26n-balloon.om" # YOLO26 模型路径 +CONF_THRESHOLD = 0.25 # 置信度过滤阈值 +CAMERA_INDEX = 0 # 机载相机索引 + +# MAVLink 串口配置 (请根据香橙派实际使用的串口设备和飞控波特率修改) +MAVLINK_PORT = "/dev/ttyAMA1" # 香橙派引脚对应的串口设备，例如 ttyAMA1 或 ttyS1 +MAVLINK_BAUD = 921600 # 推荐 921600 保证高频数据不卡顿 + +# 地面站推流配置 +GCS_IP = "114.114.114.114" # 替换为您地面站电脑/平板的局域网静态 IP 地址 +GCS_UDP_PORT = 5600 # QGroundControl 默认的视频监听端口是 5600 + + +# ============================================================ + +class YOLO26UAVInfer: + def __init__(self, model_path, device_id=0): + 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] + 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 = new_shape[1] - new_unpad[0] + dh = new_shape[0] - new_unpad[1] + left, top = dw // 2, dh // 2 + right, bottom = dw - left, dh - top + + if shape[::-1] != new_unpad: + img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) + img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) + return img, r, (left, top) + + def preprocess_frame(self, img_bgr): + self.original_height, self.original_width = img_bgr.shape[:2] + img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) + img_resized, self.ratio, (self.pad_left, self.pad_top) = self.letterbox( + img_rgb, (self.input_width, self.input_height) + ) + img_input = (img_resized.astype(np.float32) / 255.0).astype(np.float32) + img_input = np.ascontiguousarray(np.transpose(img_input, (2, 0, 1))[np.newaxis, :, :, :]) + return img_input + + def postprocess(self, infer_result, conf_threshold): + detections = [] + output = infer_result[0] if isinstance(infer_result, (list, tuple)) else infer_result + if output.ndim == 3: output = output[0] + + for det in output: + confidence = det[4] + if confidence < conf_threshold: continue + + x1_scale, y1_scale, x2_scale, y2_scale = det[0], det[1], det[2], det[3] + cls_id = int(det[5]) + + x1 = max(0, min((x1_scale - self.pad_left) / self.ratio, self.original_width)) + y1 = max(0, min((y1_scale - self.pad_top) / self.ratio, self.original_height)) + x2 = max(0, min((x2_scale - self.pad_left) / self.ratio, self.original_width)) + y2 = max(0, min((y2_scale - self.pad_top) / self.ratio, self.original_height)) + + detections.append([x1, y1, x2, y2, confidence, cls_id]) + return detections + + +def draw_boxes(image, detections): + for det in detections: + x1, y1, x2, y2, conf, cls_id = det + x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) + cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText(image, f"Target:{cls_id} {conf:.2f}", (x1, y1 - 5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) + return image + + +if __name__ == "__main__": + print("====== 正在初始化机载 AI 计算机系统 ======") + + # 1. 初始化 MAVLink 串口连接到 PX4 + print(f"正在建立与 PX4 飞控的 MAVLink 串口连接 ({MAVLINK_PORT})...") + try: + master = mavutil.mavlink_connection(MAVLINK_PORT, baud=MAVLINK_BAUD) + # 等待飞控心跳，确保连接成功 + master.wait_heartbeat(timeout=5) + print("与 PX4 飞控握手成功（收到心跳包）!") + except Exception as e: + print(f"警告：无法连接到飞控串口 ({e})。数据发送将被跳过。") + master = None + + # 2. 初始化 YOLOv26 NPU 会话 + yolo_infer = YOLO26UAVInfer(MODEL_PATH) + + # 3. 初始化机载相机 + cap = cv2.VideoCapture(CAMERA_INDEX) + if not cap.isOpened(): + print(f"致命错误：无法开启机载相机 index {CAMERA_INDEX}") + exit() + + # 设定适合推流的相机分辨率 + frame_w, frame_h = 640, 480 + cap.set(cv2.CAP_PROP_FRAME_WIDTH, frame_w) + cap.set(cv2.CAP_PROP_FRAME_HEIGHT, frame_h) + + # 4. 初始化地面站 UDP GStreamer 推流管道 + # 这一行命令会通过 H.264 编码，将视频转化为连续帧以 UDP 格式推往地面站 IP 端口 + gst_pipeline = f"appsrc ! videoconvert ! x264enc tune=zerolatency bitrate=800 speed-preset=ultrafast ! rtph264pay config-interval=1 pt=96 ! udpsink host={GCS_IP} port={GCS_UDP_PORT}" + out_stream = cv2.VideoWriter(gst_pipeline, cv2.CAP_GSTREAMER, 0, 30.0, (frame_w, frame_h), True) + if not out_stream.isOpened(): + print("警告：GStreamer 推流管道初始化失败，请检查香橙派是否安装了 gstreamer 组件。将退回到纯后台模式。") + else: + print(f"地面站推流已开启：目标 IP -> {GCS_IP}:{GCS_UDP_PORT}") + + print("系统启动成功，正在进入任务循环...") + fps_start_time = 0 + + while True: + ret, frame = cap.read() + if not ret: continue + + # 计算图像物理中心点 + img_center_x = frame_w / 2 + img_center_y = frame_h / 2 + + # NPU 执行推理 + input_tensor = yolo_infer.preprocess_frame(frame) + raw_result = yolo_infer.infer(input_tensor) + detections = yolo_infer.postprocess(raw_result, conf_threshold=CONF_THRESHOLD) + + # 核心逻辑：如果检测到了目标，计算目标相对画面中心的物理/像素偏差 + if len(detections) > 0: + # 假设默认追踪置信度最高的那一个目标 + best_det = max(detections, key=lambda x: x[4]) + x1, y1, x2, y2, conf, cls_id = best_det + + # 计算目标的中心像素坐标 + target_center_x = (x1 + x2) / 2 + target_center_y = (y1 + y2) / 2 + + # 计算目标相对于画面正中心的视角绝对角偏差 (或粗暴计算像素比偏差) + # 弧度视角偏差计算公式例：(target_center - img_center) / 焦距 + # 这里简化为标准化像素误差 (范围 -1.0 到 1.0) 方便 PX4 接收 + err_x = (target_center_x - img_center_x) / img_center_x + err_y = (target_center_y - img_center_y) / img_center_y + + # 通过 MAVLink 串口发送给 PX4 飞控 + if master is not None: + # 使用 LANDING_TARGET 消息（常用于自主追踪降落或悬停对准） + master.mav.landing_target_send( + time_usec=int(time.time() * 1e6), # 时间戳 + target_num=0, # 目标编号 + frame=mavutil.mavlink.MAV_FRAME_BODY_NED, # 相对机体坐标系 + angle_x=err_x, # X轴视角弧度/误差 + angle_y=err_y, # Y轴视角弧度/误差 + distance=0.0, # 距离 (如果不准可以填0，飞控会结合气压计计算) + size_x=0.0, size_y=0.0 # 目标物理大小 + ) + + # 在图像上额外标出追踪器连线 + cv2.line(frame, (int(img_center_x), int(img_center_y)), + (int(target_center_x), int(target_center_y)), (255, 0, 0), 2) + + # 绘制检测框与渲染实时机载 FPS + frame = draw_boxes(frame, detections) + fps_end_time = time.time() + time_diff = fps_end_time - fps_start_time + fps = 1 / time_diff if time_diff > 0 else 0 + fps_start_time = fps_end_time + cv2.putText(frame, f"Air-FPS: {fps:.1f}", (20, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2) + + # 通过数传网络将带有框和 FPS 的画面推向地面站 + if out_stream.isOpened(): + out_stream.write(frame) + + cap.release() + if out_stream.isOpened(): out_stream.release() \ No newline at end of file