diff --git a/check.py b/check.py index 21607a2..fad0073 100644 --- a/check.py +++ b/check.py @@ -1,7 +1,10 @@ +# =================================== +# 查看视频的帧率等属性 +# =================================== import cv2 import os -video_path = "./output_result.mp4" # 替换成你上传的视频文件名 +video_path = "saved_video/output_result.mp4" # 替换成你上传的视频文件名 if not os.path.exists(video_path): print(f"错误：找不到文件 {video_path}") @@ -34,9 +37,9 @@ print(f"视频总帧数: {int(total_frames)} 帧") print(f"视频总时长: {duration:.2f} 秒") print("==========================================") -if fps <= 25.5: - print("💡 结论：不出所料，你的视频本身就被录制/压制成了 25 帧！") - print("OpenCV 的 cap.read() 会卡在 25fps 的播放速度上等下一帧，导致 NPU 闲置。") - print("你想测试 NPU 的真实上限，需要用前面提到的『纯随机数矩阵压测脚本』。") +# if fps <= 25.5: +# print("💡 结论：不出所料，你的视频本身就被录制/压制成了 25 帧！") +# print("OpenCV 的 cap.read() 会卡在 25fps 的播放速度上等下一帧，导致 NPU 闲置。") +# print("你想测试 NPU 的真实上限，需要用前面提到的『纯随机数矩阵压测脚本』。") cap.release() \ No newline at end of file diff --git a/infer_camera_direct.py b/infer_camera_direct.py index 8d88062..f3e8e19 100644 --- a/infer_camera_direct.py +++ b/infer_camera_direct.py @@ -2,157 +2,182 @@ import os import cv2 import numpy as np import time +import threading +import queue +import socket from ais_bench.infer.interface import InferSession -# ==================== 硬编码配置区域 ==================== -MODEL_PATH = "./om/yolo26n-balloon.om" # 你的 YOLO26 模型路径 -CONF_THRESHOLD = 0.25 # 置信度过滤阈值 (YOLO26 推荐 0.25) -CAMERA_INDEX = 0 # 摄像头索引 (默认0，如果打不开可以尝试1或2) +# ==================== 配置区域 ==================== +MODEL_PATH = "./om/yolo26n-balloon.om" # YOLO26 模型路径 +INPUT_VIDEO = "./test_video.mp4" # 输入的测试视频（或 0 代表摄像头） +CONF_THRESHOLD = 0.25 # 置信度阈值 + +# -------------- 无线图传网络配置 -------------- +GROUND_STATION_IP = "192.168.31.239" # 地面站的局域网 IP +UDP_PORT = 9999 +JPEG_QUALITY = 75 # JPEG 压缩质量 (0-100) +# ================================================== + +# 限制队列大小，发不完就直接丢帧，保证飞控算法实时性 +tx_queue = queue.Queue(maxsize=2) -# ======================================================== +def udp_stream_sender_thread(): + """ 📡 后台图传专用线程 """ + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + server_address = (GROUND_STATION_IP, UDP_PORT) + print(f"📡 图传后台线程已启动...") -class YOLO26CameraInfer: + while True: + try: + frame_to_send = tx_queue.get() + if frame_to_send is None: + break + + # 压缩为 JPEG 字节流 + result, img_encode = cv2.imencode('.jpg', frame_to_send, [cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY]) + if not result: + continue + + data = img_encode.tobytes() + if len(data) > 65000: # 严格防丢包/报错限制 + continue + + sock.sendto(data, server_address) + except Exception as e: + print(f"❌ 图传线程异常: {e}") + time.sleep(0.1) + + sock.close() + + +class YOLO26UAVInfer: def __init__(self, model_path, device_id=0): - self.device_id = device_id - self.model_path = model_path - - # 自动检查模型文件是否存在，防止硬编码路径写错导致程序崩溃 if not os.path.exists(model_path): - raise FileNotFoundError(f"找不到硬编码的模型文件: {model_path}，请检查路径是否正确。") - - # 创建 ais_bench 推理会话 + 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] + 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 = new_shape[1] - new_unpad[0] - dh = new_shape[0] - new_unpad[1] - left = dw // 2 - right = dw - left - top = dh // 2 - bottom = dh - top - + 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, (left, top) + return img, r, (dw, dh) - 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): + def postprocess(self, outputs, ratio, dwdh, orig_shape): + out = outputs[0] detections = [] - if isinstance(infer_result, (list, tuple)): - output = infer_result[0] - else: - output = 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 = (x1_scale - self.pad_left) / self.ratio - y1 = (y1_scale - self.pad_top) / self.ratio - x2 = (x2_scale - self.pad_left) / self.ratio - y2 = (y2_scale - self.pad_top) / self.ratio - - x1 = max(0, min(x1, self.original_width)) - y1 = max(0, min(y1, self.original_height)) - x2 = max(0, min(x2, self.original_width)) - y2 = max(0, min(y2, self.original_height)) - - detections.append([x1, y1, x2, y2, confidence, cls_id]) - + 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 > 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(image, detections): +def draw_boxes(img, 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(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 - cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) - label = f"Class: {cls_id} {conf:.2f}" - (label_w, label_h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) - cv2.rectangle(image, (x1, y1 - label_h - 5), (x1 + label_w, y1), (0, 255, 0), -1) - cv2.putText(image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1) +def main(): + detector = YOLO26UAVInfer(MODEL_PATH) + cap = cv2.VideoCapture(INPUT_VIDEO) - return image + if not cap.isOpened(): + print("❌ 无法打开视频源") + return + + # 启动后台无线电图传线程 + sender_thread = threading.Thread(target=udp_stream_sender_thread, daemon=True) + sender_thread.start() + + print("🚀 NPU 精准推理主循环启动...") + while cap.isOpened(): + ret, frame = cap.read() + if not ret: + break + + # 【核心修正】完全保留读取到的原图大小，坚决不在这一步进行 cv2.resize！ + orig_shape = frame.shape[:2] + img_center_x = orig_shape[1] / 2 + img_center_y = orig_shape[0] / 2 + + # 1. 严格使用原图进行标准无畸变 letterbox 图像预处理 + input_tensor, ratio, dwdh = detector.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) + + t_start = time.time() + # 2. NPU 硬件推理 + outputs = detector.session.infer([input_tensor]) + # 3. 坐标反求（恢复到大图的分辨率，保证百分之百精准） + detections = detector.postprocess(outputs, ratio, dwdh, orig_shape) + t_end = time.time() + + fps_pure = 1.0 / (t_end - t_start) + + # 4. 寻找置信度最高的目标并计算 MAVLink 飞控误差 + 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 + + # 画中心对准线 + 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) + else: + cv2.putText(frame, "MAVLink Out -> Target LOST (HOVER)", (20, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.6, + (0, 0, 255), 2) + + # 5. 在原始分辨率大图上，渲染精准目标框 + frame = 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) + + # 🚀【核心优化点】只有在进入发送队列的前一秒，才将画好完美框的图 resize 为轻量化分辨率 + send_frame = cv2.resize(frame, (640, 480)) + + # 6. 非阻塞式尝试推入图传队列 + try: + tx_queue.put_nowait(send_frame) + except queue.Full: + pass # 队列满则丢弃该帧，绝不积压延迟 + + cap.release() + tx_queue.put(None) + print("程序运行结束。") if __name__ == "__main__": - print("正在初始化 YOLOv26 离线推理引擎...") - # 传入硬编码的模型路径 - yolo_infer = YOLO26CameraInfer(MODEL_PATH) - - # 打开硬编码指定的摄像头 - cap = cv2.VideoCapture(CAMERA_INDEX) - if not cap.isOpened(): - print(f"错误：无法打开摄像头索引 {CAMERA_INDEX}，请检查物理连接。") - exit() - - # 尝试设置合适的分辨率 - cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) - cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) - - print(f"初始化成功！当前模型：{MODEL_PATH} | 过滤阈值：{CONF_THRESHOLD}") - print("开始实时检测，在画面窗口按下 'q' 键可退出...") - - fps_start_time = 0 - - while True: - ret, frame = cap.read() - if not ret: - continue - - # 预处理与 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) - - # 绘制结果 - frame = draw_boxes(frame, detections) - - # FPS 实时计算与显示 - 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"FPS: {fps:.1f}", (20, 40), - cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2) - - cv2.imshow("YOLOv26 Real-time Detection", frame) - - if cv2.waitKey(1) & 0xFF == ord('q'): - break - - cap.release() - cv2.destroyAllWindows() - print("检测已结束。") \ No newline at end of file + main() \ No newline at end of file diff --git a/infer_video.py b/infer_video.py index 02d64dc..5ab173f 100644 --- a/infer_video.py +++ b/infer_video.py @@ -2,54 +2,17 @@ import os import cv2 import numpy as np import time -import threading -import queue -import socket from ais_bench.infer.interface import InferSession # ==================== 配置区域 ==================== -MODEL_PATH = "./om/yolo26n-balloon.om" # YOLO26 模型路径 -INPUT_VIDEO = "./test_video.mp4" # 输入的测试视频（首飞时可改成 0 代表机载 USB 相机） +MODEL_PATH = "./om/yolo26n-balloon.om" # 你的 YOLO26 模型路径 +INPUT_VIDEO = "./test_video.mp4" # 输入的 24 帧测试视频 +OUTPUT_VIDEO = "./output_result.mp4" # 检测完成后保存的新视频路径 CONF_THRESHOLD = 0.25 # 置信度阈值 -# -------------- 无线图传网络配置 -------------- -GROUND_STATION_IP = "192.168.31.239" # 地面站的局域网 IP -UDP_PORT = 9999 + # ================================================== -# 初始化一个图传专用的轻量级队列，最大长度为2。如果后台发得慢导致队列满了，主推理环会触发丢帧，绝对不卡死。 -tx_queue = queue.Queue(maxsize=2) - - -def udp_stream_sender_thread(): - """ 📡 后台图传专用线程 """ - sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - server_address = (GROUND_STATION_IP, UDP_PORT) - print(f"📡 图传后台线程已启动，正在持续向地面站 [{GROUND_STATION_IP}:{UDP_PORT}] 推流...") - - while True: - try: - # 阻塞等待主线程丢进来的实时画框图像 - frame = tx_queue.get() - - # 1. 预处理：将图片缩放到 640x480（极大降低空中无线带宽压力和延迟） - small_frame = cv2.resize(frame, (640, 480)) - - # 2. 图像进行轻量级 JPEG 压缩（质量设为 60，画面清晰且计算速度极快） - ret, jpeg_img = cv2.imencode('.jpg', small_frame, [int(cv2.IMWRITE_JPEG_QUALITY), 60]) - - if ret: - data = np.array(jpeg_img).tobytes() - # 3. 确保单个 UDP 数据包不超过缓冲区上限（~65KB） - if len(data) < 65000: - sock.sendto(data, server_address) - - tx_queue.task_done() - except Exception: - # 即使无线网络抖动或地面站未开启，后台线程也静默承受，绝不连累无人机主系统 - pass - - class YOLO26UAVInfer: def __init__(self, model_path, device_id=0): if not os.path.exists(model_path): @@ -58,111 +21,99 @@ class YOLO26UAVInfer: self.input_width = 640 self.input_height = 640 - def letterbox(self, img, new_shape=(640, 640), color=(114, 114, 114)) : + 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] - dw /= 2 - dh /= 2 + 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) - 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) + 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 -def preprocess(self, img): - img_canvas, ratio, (dw, dh) = self.letterbox(img, (self.input_width, self.input_height)) - img_in = img_canvas[:, :, ::-1].transpose(2, 0, 1) - img_in = np.ascontiguousarray(img_in).astype(np.float32) / 255.0 - img_in = np.expand_dims(img_in, axis=0) - return img_in, ratio, (dw, dh) - - -def infer(self, input_tensor): - outputs = self.session.infer([input_tensor]) - return outputs - - -def postprocess(self, outputs, ratio, pad, orig_shape, conf_thres=0.25): - det_out = outputs[0][0] - valid_detections = [] - for i in range(det_out.shape[0]): - row = det_out[i] - conf = row[4] - if conf < conf_thres: - continue - cls_id = int(row[5]) - cx, cy, w, h = row[0], row[1], row[2], row[3] - x1 = cx - w / 2 - y1 = cy - h / 2 - x2 = cx + w / 2 - y2 = cy + h / 2 - dw, dh = pad - x1 = (x1 - dw) / ratio - y1 = (y1 - dh) / ratio - x2 = (x2 - dw) / ratio - y2 = (y2 - dh) / ratio - x1 = max(0, min(orig_shape[1], x1)) - y1 = max(0, min(orig_shape[0], y1)) - x2 = max(0, min(orig_shape[1], x2)) - y2 = max(0, min(orig_shape[0], y2)) - valid_detections.append([x1, y1, x2, y2, conf, cls_id]) - return valid_detections - - -def draw_boxes(img, detections): - img_copy = img.copy() +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(img_copy, (x1, y1), (x2, y2), (0, 255, 0), 2) - label = f"Target: {conf:.2f}" - (label_w, label_h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) - cv2.rectangle(img_copy, (x1, y1 - label_h - 5), (x1 + label_w, y1), (0, 255, 0), -1) - cv2.putText(img_copy, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1) - return img_copy + # 画绿色的目标检测框 + 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__": - # 0. 提前拉起后台异步无线图传线程 - stream_thread = threading.Thread(target=udp_stream_sender_thread, daemon=True) - stream_thread.start() - - # 初始化推理会话 yolo_infer = YOLO26UAVInfer(MODEL_PATH) - # 视频输入源 + # 1. 打开输入视频 cap = cv2.VideoCapture(INPUT_VIDEO) if not cap.isOpened(): - print(f"❌ 无法打开输入视频: {INPUT_VIDEO}") + print(f"错误：无法打开输入视频 {INPUT_VIDEO}") exit() - print("🚀 NPU 推理主循环已启动！") + # 读取输入视频的各种属性，用于配置保存视频的写入器 + frame_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + frame_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + video_fps = cap.get(cv2.CAP_PROP_FPS) + img_center_x, img_center_y = frame_w / 2, frame_h / 2 + + # 2. 【核心点】初始化视频保存器 (VideoWriter) + # mp4v 是最通用、在 Linux 下不需要额外配置库的编码格式 + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + video_writer = cv2.VideoWriter(OUTPUT_VIDEO, fourcc, video_fps, (frame_w, frame_h)) + + print(f"开始处理视频并渲染可视化结果...") + print(f"原视频属性: {frame_w}x{frame_h} @ {video_fps:.2f} FPS") + print(f"处理后的视频将实时保存到: {OUTPUT_VIDEO}") + + fps_start_time = time.time() + frame_count = 0 while True: - t_start = time.time() ret, frame = cap.read() if not ret: - print("视频播放结束或未捕获到帧。") break - orig_shape = frame.shape[:2] - img_center_x = orig_shape[1] / 2 - img_center_y = orig_shape[0] / 2 + frame_count += 1 - # 1. 预处理 - input_tensor, ratio, pad = yolo_infer.preprocess(frame) + # 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. 执行昇腾 NPU 硬件推理 - outputs = yolo_infer.infer(input_tensor) - - # 3. 后处理 - detections = yolo_infer.postprocess(outputs, ratio, pad, orig_shape, conf_thres=CONF_THRESHOLD) - - # 4. 执行你的核心控制逻辑（寻找最显著的目标并计算偏差） + # 如果检测到目标，计算偏差并画可视化指示线 if len(detections) > 0: best_det = max(detections, key=lambda x: x[4]) x1, y1, x2, y2, conf, cls_id = best_det @@ -183,25 +134,22 @@ if __name__ == "__main__": cv2.putText(frame, "MAVLink Out -> Target LOST (HOVER)", (20, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2) - # 5. 渲染基础边界框 + # 渲染框和实时运行 FPS frame = draw_boxes(frame, detections) - # 6. 计算单帧总耗时和纯计算 FPS - t_end = time.time() - fps = 1.0 / (t_end - t_start) - cv2.putText(frame, f"NPU Process-FPS: {fps:.1f}", (20, 40), - cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) + end_time = time.time() + curr_fps = 1 / (end_time - fps_start_time) if (end_time - fps_start_time) > 0 else video_fps + fps_start_time = end_time + cv2.putText(frame, f"NPU Process-FPS: {curr_fps:.1f}", (20, 30), + cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) - # ==================== 异步丢帧推流 ==================== - try: - # block=False 意味着如果无线网络太慢队列满了，程序不会原地死等，而是抛出异常并直接走下一次循环 - tx_queue.put(frame, block=False) - except queue.Full: - pass # 主动丢帧 - # =================================================================== + # 3. 【核心点】把这一帧画好的画面写入新视频文件 + video_writer.write(frame) - # 打印调试信息到终端 - print(f"当前帧推理完成 | 纯计算主环帧率: {fps:.1f} FPS") + if frame_count % 24 == 0: + print(f"已处理并写入 {frame_count} 帧...") + # 4. 释放资源，这一步至关重要！不 release 视频文件会损坏打不开 cap.release() - print("香橙派端任务运行完毕。") \ No newline at end of file + video_writer.release() + print(f"🎉 处理完毕！完整的检测结果视频已成功保存至: {OUTPUT_VIDEO}") \ No newline at end of file diff --git a/receiver.py b/receiver.py new file mode 100644 index 0000000..3688fe3 --- /dev/null +++ b/receiver.py @@ -0,0 +1,34 @@ +import socket +import cv2 +import numpy as np + +PORT = 9999 +sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) +sock.bind(('0.0.0.0', PORT)) + +print("📡 [笔记本地面站] 已启动，正在等待香橙派连线...") + +# 初始化地面站的强悍 CPU 录像（对应发送端压缩后的分辨率） +fourcc = cv2.VideoWriter_fourcc(*'mp4v') +video_writer = cv2.VideoWriter("./saved_video/ground_station_record.mp4", fourcc, 24.0, (640, 480)) + +try: + while True: + packet, _ = sock.recvfrom(65535) + data = np.frombuffer(packet, dtype=np.uint8) + frame = cv2.imdecode(data, cv2.IMREAD_COLOR) + + if frame is not None: + # 1. 地面站录像 + video_writer.write(frame) + + # 2. 实时显示监控 + cv2.imshow("🚀 UAV AI Ground Station Monitor", frame) + + if cv2.waitKey(1) & 0xFF == ord('q'): + break +finally: + sock.close() + video_writer.release() + cv2.destroyAllWindows() + print("地面站已安全关闭，视频成功保存在本地。") \ No newline at end of file diff --git a/saved_video/ground_station_record.mp4 b/saved_video/ground_station_record.mp4 new file mode 100644 index 0000000..128a176 Binary files /dev/null and b/saved_video/ground_station_record.mp4 differ diff --git a/output_result.mp4 b/saved_video/output_result.mp4 similarity index 100% rename from output_result.mp4 rename to saved_video/output_result.mp4 diff --git a/uav_ai_companion.py b/uav_ai_companion.py deleted file mode 100644 index 4a11985..0000000 --- a/uav_ai_companion.py +++ /dev/null @@ -1,188 +0,0 @@ -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