diff --git a/.idea/.gitignore b/.idea/.gitignore new file mode 100644 index 0000000..10b731c --- /dev/null +++ b/.idea/.gitignore @@ -0,0 +1,5 @@ +# 默认忽略的文件 +/shelf/ +/workspace.xml +# 基于编辑器的 HTTP 客户端请求 +/httpRequests/ diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml new file mode 100644 index 0000000..fd02ef8 --- /dev/null +++ b/.idea/inspectionProfiles/Project_Default.xml @@ -0,0 +1,26 @@ + + + + \ No newline at end of file diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml new file mode 100644 index 0000000..105ce2d --- /dev/null +++ b/.idea/inspectionProfiles/profiles_settings.xml @@ -0,0 +1,6 @@ + + + + \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml new file mode 100644 index 0000000..dc5f227 --- /dev/null +++ b/.idea/misc.xml @@ -0,0 +1,7 @@ + + + + + + \ No newline at end of file diff --git a/.idea/modules.xml b/.idea/modules.xml new file mode 100644 index 0000000..a3e7011 --- /dev/null +++ b/.idea/modules.xml @@ -0,0 +1,8 @@ + + + + + + + + \ No newline at end of file diff --git a/.idea/opi-yolo.iml b/.idea/opi-yolo.iml new file mode 100644 index 0000000..c506ebe --- /dev/null +++ b/.idea/opi-yolo.iml @@ -0,0 +1,12 @@ + + + + + + + + + + \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 0000000..8306744 --- /dev/null +++ b/.idea/vcs.xml @@ -0,0 +1,7 @@ + + + + + + + \ No newline at end of file diff --git a/infer.py b/infer.py index 7b51b10..f0cfd40 100644 --- a/infer.py +++ b/infer.py @@ -1,78 +1,198 @@ +import os import cv2 import numpy as np +import argparse +from ais_bench.infer.interface import InferSession -# 导入昇腾核心推理库(板子自带) -from acllite_model import AclLiteModel -from acllite_resource import AclLiteResource +class YOLO26Infer: + def __init__(self, model_path, device_id=0): + """ + 初始化推理器 + :param model_path: .om 模型的本地路径 + :param device_id: NPU 设备ID + """ + self.device_id = device_id + self.model_path = model_path + # 创建 ais_bench 推理会话 + self.session = InferSession(device_id, model_path) + # 获取模型输入尺寸,用于后续预处理 + self.input_width = 640 + self.input_height = 640 -def preprocess(img_path, model_w=640, model_h=640): - """图片预处理:缩放、通道转换、归一化""" - orig_img = cv2.imread(img_path) - h, w, _ = orig_img.shape - - # 1. 缩放到模型要求的 640x640 - img = cv2.resize(orig_img, (model_w, model_h)) - # 2. BGR 转 RGB (YOLO 训练时通常是 RGB) - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - # 3. 归一化并调整维度为 [1, 3, 640, 640] - img = img.astype(np.float32) / 255.0 - img = img.transpose(2, 0, 1) # HWC -> CHW - img = np.expand_dims(img, axis=0) # CHW -> NCHW - - # 必须将其转换为连续的内存,方便 NPU 读取 - return np.ascontiguousarray(img), orig_img + 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 = dw // 2 + right = dw - left + top = dh // 2 + bottom = dh - top -def main(): - # 1. 初始化昇腾 NPU 资源 - acl_resource = AclLiteResource() - acl_resource.init() + # 缩放图片 + if shape[::-1] != new_unpad: + img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) - # 2. 加载你的 YOLO26 .om 模型 - model_path = "./yolo26_8t.om" - model = AclLiteModel(model_path) + # 添加填充(可能左右或上下不对称,但总尺寸正好是 new_shape) + img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) + return img, r, (left, top) # 返回填充的左上角偏移量 - # 3. 图像预处理 - input_data, orig_img = preprocess("./test.jpg") + def preprocess(self, image_path): + """ + 预处理图片,使其符合模型输入要求 + """ + # 使用OpenCV读取图片 (BGR格式) + img_bgr = cv2.imread(image_path) + if img_bgr is None: + raise FileNotFoundError(f"图片未找到: {image_path}") - # 4. 送入 NPU 进行硬件加速推理 - # execute 接收一个列表,里面是模型的输入 Tensor - result_list = model.execute([input_data]) + # 获取原始图片尺寸,用于后处理时的坐标转换 + self.original_height, self.original_width = img_bgr.shape[:2] - # 5. 后处理(针对 YOLO26 的端到端特性) - # YOLO26 输出的直接是最终预测框,result_list[0] 形状通常是 [1, 300, 6] 或类似 - # 6个维度的含义一般为:[x_min, y_min, x_max, y_max, score, class_id] - predictions = result_list[0][0] + # 将图片从BGR转换为RGB + img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) - CONF_THRESHOLD = 0.25 # 设定置信度阈值 - orig_h, orig_w = orig_img.shape[:2] + # 调整图片尺寸并保持宽高比,进行Letterbox填充 + img_resized, self.ratio, (self.pad_left, self.pad_top) = self.letterbox( + img_rgb, (self.input_width, self.input_height) + ) - box_count = 0 - for pred in predictions: - # 取出置信度和类别 - score = pred[4] - class_id = int(pred[5]) - - # 过滤掉低置信度的框 - if score > CONF_THRESHOLD: - box_count += 1 - # 取出缩放后的坐标 [0~640] - x1, y1, x2, y2 = pred[0:4] - - # 将坐标还原到原图尺寸 - x1 = int(x1 * orig_w / 640) - y1 = int(y1 * orig_h / 640) - x2 = int(x2 * orig_w / 640) - y2 = int(y2 * orig_h / 640) - - # 在原图上画框和标签 - cv2.rectangle(orig_img, (x1, y1), (x2, y2), (0, 255, 0), 2) - cv2.putText(orig_img, f"Class {class_id}: {score:.2f}", (x1, y1 - 10), - cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) + # 归一化:将像素值从 [0, 255] 转为 [0, 1] + img_input = (img_resized.astype(np.float32) / 255.0).astype(np.float32) - print(f"推理完成,共检测到 {box_count} 个目标。") - - # 6. 保存检测结果图 - cv2.imwrite("./result.jpg", orig_img) + # 调整维度:HWC -> NCHW,并转换为连续数组 + img_input = np.ascontiguousarray(np.transpose(img_input, (2, 0, 1))[np.newaxis, :, :, :]) + + return img_input + + def infer(self, img_input): + """ + 执行推理 + """ + # 使用会话的 infer 方法进行推理 + # infer 方法接收一个列表,包含输入张量,返回一个列表,包含输出张量 + result = self.session.infer([img_input]) + print(f"Input tensor shape: {img_input.shape}") + print(f"Input tensor dtype: {img_input.dtype}") + print(f"Input tensor size: {img_input.nbytes}") + return result + + def postprocess(self, infer_result, conf_threshold=0.5): + """ + 对模型输出进行后处理 + :param infer_result: 模型原始输出,一个列表,通常第一个元素是输出张量 + :param conf_threshold: 置信度阈值 + :return: 包含检测结果的列表,每个元素为 [x1, y1, x2, y2, confidence, class_id] + """ + detections = [] + # 获取输出张量,假设输出列表的第一个元素是输出数组 + if isinstance(infer_result, (list, tuple)): + output = infer_result[0] + else: + output = infer_result + + # ========== 调试信息 ========== + # print(f"[DEBUG] output shape: {output.shape}") + # print(f"[DEBUG] output dtype: {output.dtype}") + # print(f"[DEBUG] output min: {output.min()}, max: {output.max()}") + # # 打印前5个检测框(如果形状是 [N,6] 或 [1,N,6]) + # if output.ndim == 3: + # flat_output = output[0] + # else: + # flat_output = output + # print(f"[DEBUG] first 5 rows:\n{flat_output[:5]}") + + # 假设输出维度为 [batch, num_detections, 6] + # 移除batch维度,并只取第一个batch的结果 + if output.ndim == 3: + output = output[0] + + # 遍历所有检测框 + for det in output: + confidence = det[4] + if confidence < conf_threshold: + continue + + # 获取原始输出坐标 (x1, y1, x2, y2) 和 类别ID + 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]) + + return detections + +def draw_boxes(image, detections, class_names=None): + """ + 在图片上绘制检测框和标签 + :param image: 原始图片数组 (BGR格式) + :param detections: 后处理得到的结果列表 + :param class_names: 类别名称列表,默认使用类别ID + """ + img_copy = image.copy() + 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"Class: {cls_id}" if class_names is None else class_names[cls_id] + label += f" {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 if __name__ == "__main__": - main() \ No newline at end of file + parser = argparse.ArgumentParser(description="YOLOv26 目标检测推理") + parser.add_argument("--model", type=str, required=True, help="YOLOv26 OM模型文件路径") + parser.add_argument("--image", type=str, required=True, help="输入图片路径") + parser.add_argument("--conf", type=float, default=0.5, help="置信度阈值") + args = parser.parse_args() + + # 初始化推理器 + yolo_infer = YOLO26Infer(args.model) + + # 图片预处理 + print("预处理图片...") + input_tensor = yolo_infer.preprocess(args.image) + + # 执行推理 + print("执行推理...") + raw_result = yolo_infer.infer(input_tensor) + + # 后处理 + print("后处理...") + detections = yolo_infer.postprocess(raw_result, conf_threshold=args.conf) + + # 打印结果 + print(f"检测到 {len(detections)} 个目标:") + for det in detections: + print(f"类别ID: {int(det[5])}, 置信度: {det[4]:.4f}, 边界框: ({det[0]:.2f}, {det[1]:.2f}, {det[2]:.2f}, {det[3]:.2f})") + + # 可选:绘制结果并保存 + original_img = cv2.imread(args.image) + result_img = draw_boxes(original_img, detections) + output_path = "output_" + os.path.basename(args.image) + cv2.imwrite(output_path, result_img) + print(f"结果图片已保存至: {output_path}") \ No newline at end of file diff --git a/mav-test.py b/mav-test.py new file mode 100644 index 0000000..1c570e4 --- /dev/null +++ b/mav-test.py @@ -0,0 +1,175 @@ +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.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 diff --git a/test.jpg b/test.jpg new file mode 100644 index 0000000..200912c Binary files /dev/null and b/test.jpg differ