diff --git a/infer_camera_direct.py b/infer_camera_direct.py index 6cf9fd0..98463ca 100644 --- a/infer_camera_direct.py +++ b/infer_camera_direct.py @@ -9,7 +9,7 @@ from ais_bench.infer.interface import InferSession # ==================== 配置区域 ====================== MODEL_PATH = "./om/yolo26n-balloon.om" # YOLO26 模型路径 -INPUT_VIDEO = "./test_video.mp4" # 输入的测试视频(或 0 代表摄像头) +INPUT_VIDEO = 0 # 摄像头编号 CONF_THRESHOLD = 0.25 # 置信度阈值 # -------------- 无线图传网络配置 -------------- diff --git a/infer_video.py b/infer_video.py deleted file mode 100644 index 5ab173f..0000000 --- a/infer_video.py +++ /dev/null @@ -1,155 +0,0 @@ -import os -import cv2 -import numpy as np -import time -from ais_bench.infer.interface import InferSession - -# ==================== 配置区域 ==================== -MODEL_PATH = "./om/yolo26n-balloon.om" # 你的 YOLO26 模型路径 -INPUT_VIDEO = "./test_video.mp4" # 输入的 24 帧测试视频 -OUTPUT_VIDEO = "./output_result.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 - - -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__": - yolo_infer = YOLO26UAVInfer(MODEL_PATH) - - # 1. 打开输入视频 - cap = cv2.VideoCapture(INPUT_VIDEO) - if not cap.isOpened(): - print(f"错误:无法打开输入视频 {INPUT_VIDEO}") - exit() - - # 读取输入视频的各种属性,用于配置保存视频的写入器 - 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: - ret, frame = cap.read() - if not ret: - break - - frame_count += 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) - - # 如果检测到目标,计算偏差并画可视化指示线 - 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) - - # 渲染框和实时运行 FPS - frame = draw_boxes(frame, detections) - - 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) - - # 3. 【核心点】把这一帧画好的画面写入新视频文件 - video_writer.write(frame) - - if frame_count % 24 == 0: - print(f"已处理并写入 {frame_count} 帧...") - - # 4. 释放资源,这一步至关重要!不 release 视频文件会损坏打不开 - cap.release() - video_writer.release() - print(f"🎉 处理完毕!完整的检测结果视频已成功保存至: {OUTPUT_VIDEO}") \ No newline at end of file diff --git a/infer_video_direct.py b/infer_video_direct.py new file mode 100644 index 0000000..e420f04 --- /dev/null +++ b/infer_video_direct.py @@ -0,0 +1,191 @@ +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" # 输入的测试视频 +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"📡 图传后台线程已启动...") + + 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): + 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 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 > 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(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 + + +def main(): + detector = YOLO26UAVInfer(MODEL_PATH) + cap = cv2.VideoCapture(INPUT_VIDEO) + from control import UAVController + + uav = UAVController() + + 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 + + # 完全保留读取到的原图大小 + 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) + # ================ MAVLink 动作调用 ================= + uav.update_control(has_target=True, err_x=err_x, err_y=err_y) + # ================================================= + else: + cv2.putText(frame, "MAVLink Out -> Target LOST (HOVER)", (20, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.6, + (0, 0, 255), 2) + # ================ MAVLink 丢失保护 ================= + uav.update_control(has_target=False) + # ================================================= + + # 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__": + main() \ No newline at end of file diff --git a/mav-test.py b/mav-test.py deleted file mode 100644 index 0795ed7..0000000 --- a/mav-test.py +++ /dev/null @@ -1,175 +0,0 @@ -import cv2 -import math -import time -import serial -import numpy as np -from ais_bench.infer.interface import InferSession -from pymavlink import mavutil - - -class DroneDetector: - def __init__(self, model_path, serial_port='/dev/ttyAMA0', baudrate=57600): - self.session = InferSession(0, model_path) - self.input_size = 640 - # 串口/MAVLink - self.serial_port = serial_port - self.baudrate = baudrate - self.mav = None - self.init_serial() - - def init_serial(self): - """初始化串口(如果用于自定义协议)或 MAVLink 连接""" - try: - self.mav = mavutil.mavlink_connection(self.serial_port, baud=self.baudrate) - print(f"MAVLink 已连接到 {self.serial_port} 波特率 {self.baudrate}") - except Exception as e: - print(f"串口打开失败: {e}") - self.mav = None - - def letterbox(self, img, color=(114, 114, 114)): - shape = img.shape[:2] - r = min(self.input_size / shape[0], self.input_size / shape[1]) - new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r))) - dw = self.input_size - new_unpad[0] - dh = self.input_size - new_unpad[1] - left = dw // 2 - right = dw - left - top = dh // 2 - bottom = dh - top - 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(self, img): - img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - img_resized, self.ratio, (self.pad_left, self.pad_top) = self.letterbox(img_rgb) - img_input = img_resized.astype(np.float32) / 255.0 - img_input = np.ascontiguousarray(np.transpose(img_input, (2, 0, 1))[np.newaxis, ...]) - return img_input - - def postprocess(self, output, conf_thresh=0.5): - detections = [] - if output.ndim == 3: - output = output[0] - for det in output: - raw_conf = det[4] - conf = 1.0 / (1.0 + math.exp(-raw_conf)) - if conf < conf_thresh: - continue - x1, y1, x2, y2 = det[0], det[1], det[2], det[3] - # 映射回原图坐标 - x1 = (x1 - self.pad_left) / self.ratio - y1 = (y1 - self.pad_top) / self.ratio - x2 = (x2 - self.pad_left) / self.ratio - y2 = (y2 - self.pad_top) / self.ratio - cls_id = int(det[5]) - detections.append({ - 'bbox': (x1, y1, x2, y2), - 'confidence': conf, - 'class': cls_id, - 'center': ((x1 + x2) / 2, (y1 + y2) / 2) - }) - return detections - - def send_mavlink_obstacle(self, angle_deg, distance_m): - """发送 DISTANCE_SENSOR 消息,表示目标方向与距离""" - if not self.mav: - return - self.mav.mav.distance_sensor_send( - 0, # time_boot_ms - 0, # min_distance cm - 10000, # max_distance cm - int(distance_m * 100), # current_distance cm - 0, # type (0=laser) - 0, # id - int(angle_deg / 360.0 * 65535), # orientation (0~360° -> 0~65535) - 0 # covariance - ) - - def send_custom_serial(self, detections): - """通过串口发送自定义格式的数据(例如 JSON)""" - if not self.serial or not self.serial.is_open: - return - import json - data = { - 'timestamp': time.time(), - 'detections': detections - } - self.serial.write((json.dumps(data) + '\n').encode()) - - def detect_and_send(self, frame, send_method='mavlink', conf_thresh=0.5): - """对单帧图像进行检测,并发送结果""" - input_tensor = self.preprocess(frame) - outputs = self.session.infer([input_tensor]) - detections = self.postprocess(outputs[0], conf_thresh) - if not detections: - return detections - # 假设只发送置信度最高的那个目标 - best = detections[0] - if send_method == 'mavlink': - # 计算角度和距离(需要根据相机参数和实际尺寸估算) - # 这里简单示例:根据目标在图像中的偏移量估算角度 - cx = best['center'][0] - img_w = frame.shape[1] - angle = (cx - img_w / 2) / (img_w / 2) * 45 # 水平视场角90°,偏移量映射为±45° - # 距离估算:根据目标在图像中的高度(假设目标实际高度1m) - # 需要相机焦距等,简化:使用固定值10米 - distance = 10.0 - self.send_mavlink_obstacle(angle, distance) - elif send_method == 'custom': - self.send_custom_serial(detections) - return detections - - def run_on_camera(self, camera_id=0, conf_thresh=0.5, send_method='mavlink'): - cap = cv2.VideoCapture(camera_id) - if not cap.isOpened(): - print("摄像头打开失败") - return - print("开始检测,按 q 退出") - while True: - ret, frame = cap.read() - if not ret: - break - detections = self.detect_and_send(frame, send_method, conf_thresh) - # 绘制检测框(调试) - for det in detections: - x1, y1, x2, y2 = map(int, det['bbox']) - cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) - cv2.putText(frame, f"{det['class']}:{det['confidence']:.2f}", (x1, y1 - 5), - cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) - cv2.imshow("Detection", frame) - if cv2.waitKey(1) & 0xFF == ord('q'): - break - cap.release() - cv2.destroyAllWindows() - - def run_on_image(self, image_path, conf_thresh=0.5, send_method='mavlink'): - frame = cv2.imread(image_path) - if frame is None: - print("图片读取失败") - return - detections = self.detect_and_send(frame, send_method, conf_thresh) - print(f"检测到 {len(detections)} 个目标") - for d in detections: - print(d) - # 显示结果 - for det in detections: - x1, y1, x2, y2 = map(int, det['bbox']) - cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) - cv2.putText(frame, f"{det['class']}:{det['confidence']:.2f}", (x1, y1 - 5), - cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) - cv2.imshow("Detection", frame) - cv2.waitKey(0) - cv2.destroyAllWindows() - - -if __name__ == "__main__": - MODEL_PATH = "./om/yolo26n-drone.om" - SERIAL_PORT = "/dev/ttyAMA0" # 根据实际修改 - BAUD = 57600 - - detector = DroneDetector(MODEL_PATH, SERIAL_PORT, BAUD) - - # 选择模式:摄像头或单张图片 - # detector.run_on_camera(camera_id=0, conf_thresh=0.5, send_method='mavlink') - detector.run_on_image("test.jpg", conf_thresh=0.5, send_method='mavlink') \ No newline at end of file