npu端基本测试
This commit is contained in:
Generated
+5
@@ -0,0 +1,5 @@
|
|||||||
|
# 默认忽略的文件
|
||||||
|
/shelf/
|
||||||
|
/workspace.xml
|
||||||
|
# 基于编辑器的 HTTP 客户端请求
|
||||||
|
/httpRequests/
|
||||||
+26
@@ -0,0 +1,26 @@
|
|||||||
|
<component name="InspectionProjectProfileManager">
|
||||||
|
<profile version="1.0">
|
||||||
|
<option name="myName" value="Project Default" />
|
||||||
|
<inspection_tool class="PyPep8NamingInspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
|
||||||
|
<option name="ignoredErrors">
|
||||||
|
<list>
|
||||||
|
<option value="N801" />
|
||||||
|
</list>
|
||||||
|
</option>
|
||||||
|
</inspection_tool>
|
||||||
|
<inspection_tool class="PyStubPackagesAdvertiser" enabled="true" level="WARNING" enabled_by_default="true">
|
||||||
|
<option name="ignoredPackages">
|
||||||
|
<list>
|
||||||
|
<option value="scipy" />
|
||||||
|
</list>
|
||||||
|
</option>
|
||||||
|
</inspection_tool>
|
||||||
|
<inspection_tool class="PyUnresolvedReferencesInspection" enabled="true" level="WARNING" enabled_by_default="true">
|
||||||
|
<option name="ignoredIdentifiers">
|
||||||
|
<list>
|
||||||
|
<option value="types.NoneType.*" />
|
||||||
|
</list>
|
||||||
|
</option>
|
||||||
|
</inspection_tool>
|
||||||
|
</profile>
|
||||||
|
</component>
|
||||||
+6
@@ -0,0 +1,6 @@
|
|||||||
|
<component name="InspectionProjectProfileManager">
|
||||||
|
<settings>
|
||||||
|
<option name="USE_PROJECT_PROFILE" value="false" />
|
||||||
|
<version value="1.0" />
|
||||||
|
</settings>
|
||||||
|
</component>
|
||||||
Generated
+7
@@ -0,0 +1,7 @@
|
|||||||
|
<?xml version="1.0" encoding="UTF-8"?>
|
||||||
|
<project version="4">
|
||||||
|
<component name="Black">
|
||||||
|
<option name="sdkName" value="D:\program file\conda\envs\pytorch" />
|
||||||
|
</component>
|
||||||
|
<component name="ProjectRootManager" version="2" project-jdk-name="D:\program file\conda\envs\pytorch" project-jdk-type="Python SDK" />
|
||||||
|
</project>
|
||||||
Generated
+8
@@ -0,0 +1,8 @@
|
|||||||
|
<?xml version="1.0" encoding="UTF-8"?>
|
||||||
|
<project version="4">
|
||||||
|
<component name="ProjectModuleManager">
|
||||||
|
<modules>
|
||||||
|
<module fileurl="file://$PROJECT_DIR$/.idea/opi-yolo.iml" filepath="$PROJECT_DIR$/.idea/opi-yolo.iml" />
|
||||||
|
</modules>
|
||||||
|
</component>
|
||||||
|
</project>
|
||||||
Generated
+12
@@ -0,0 +1,12 @@
|
|||||||
|
<?xml version="1.0" encoding="UTF-8"?>
|
||||||
|
<module type="PYTHON_MODULE" version="4">
|
||||||
|
<component name="NewModuleRootManager">
|
||||||
|
<content url="file://$MODULE_DIR$" />
|
||||||
|
<orderEntry type="jdk" jdkName="D:\program file\conda\envs\pytorch" jdkType="Python SDK" />
|
||||||
|
<orderEntry type="sourceFolder" forTests="false" />
|
||||||
|
</component>
|
||||||
|
<component name="PyDocumentationSettings">
|
||||||
|
<option name="format" value="PLAIN" />
|
||||||
|
<option name="myDocStringFormat" value="Plain" />
|
||||||
|
</component>
|
||||||
|
</module>
|
||||||
Generated
+7
@@ -0,0 +1,7 @@
|
|||||||
|
<?xml version="1.0" encoding="UTF-8"?>
|
||||||
|
<project version="4">
|
||||||
|
<component name="VcsDirectoryMappings">
|
||||||
|
<mapping directory="" vcs="Git" />
|
||||||
|
<mapping directory="$PROJECT_DIR$" vcs="Git" />
|
||||||
|
</component>
|
||||||
|
</project>
|
||||||
@@ -1,78 +1,198 @@
|
|||||||
|
import os
|
||||||
import cv2
|
import cv2
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
import argparse
|
||||||
|
from ais_bench.infer.interface import InferSession
|
||||||
|
|
||||||
# 导入昇腾核心推理库(板子自带)
|
class YOLO26Infer:
|
||||||
from acllite_model import AclLiteModel
|
def __init__(self, model_path, device_id=0):
|
||||||
from acllite_resource import AclLiteResource
|
"""
|
||||||
|
初始化推理器
|
||||||
|
: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):
|
def letterbox(self, img, new_shape=(640, 640), color=(114, 114, 114)):
|
||||||
"""图片预处理:缩放、通道转换、归一化"""
|
shape = img.shape[:2]
|
||||||
orig_img = cv2.imread(img_path)
|
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
||||||
h, w, _ = orig_img.shape
|
# 计算缩放后的新尺寸(四舍五入)
|
||||||
|
new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
|
||||||
# 1. 缩放到模型要求的 640x640
|
# 计算需要填充的总宽度和高度
|
||||||
img = cv2.resize(orig_img, (model_w, model_h))
|
dw = new_shape[1] - new_unpad[0]
|
||||||
# 2. BGR 转 RGB (YOLO 训练时通常是 RGB)
|
dh = new_shape[0] - new_unpad[1]
|
||||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
# 分别计算左右/上下填充量(允许奇数不对称)
|
||||||
# 3. 归一化并调整维度为 [1, 3, 640, 640]
|
left = dw // 2
|
||||||
img = img.astype(np.float32) / 255.0
|
right = dw - left
|
||||||
img = img.transpose(2, 0, 1) # HWC -> CHW
|
top = dh // 2
|
||||||
img = np.expand_dims(img, axis=0) # CHW -> NCHW
|
bottom = dh - top
|
||||||
|
|
||||||
# 必须将其转换为连续的内存,方便 NPU 读取
|
|
||||||
return np.ascontiguousarray(img), orig_img
|
|
||||||
|
|
||||||
def main():
|
# 缩放图片
|
||||||
# 1. 初始化昇腾 NPU 资源
|
if shape[::-1] != new_unpad:
|
||||||
acl_resource = AclLiteResource()
|
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
|
||||||
acl_resource.init()
|
|
||||||
|
|
||||||
# 2. 加载你的 YOLO26 .om 模型
|
# 添加填充(可能左右或上下不对称,但总尺寸正好是 new_shape)
|
||||||
model_path = "./yolo26_8t.om"
|
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
|
||||||
model = AclLiteModel(model_path)
|
return img, r, (left, top) # 返回填充的左上角偏移量
|
||||||
|
|
||||||
# 3. 图像预处理
|
def preprocess(self, image_path):
|
||||||
input_data, orig_img = preprocess("./test.jpg")
|
"""
|
||||||
|
预处理图片,使其符合模型输入要求
|
||||||
|
"""
|
||||||
|
# 使用OpenCV读取图片 (BGR格式)
|
||||||
|
img_bgr = cv2.imread(image_path)
|
||||||
|
if img_bgr is None:
|
||||||
|
raise FileNotFoundError(f"图片未找到: {image_path}")
|
||||||
|
|
||||||
# 4. 送入 NPU 进行硬件加速推理
|
# 获取原始图片尺寸,用于后处理时的坐标转换
|
||||||
# execute 接收一个列表,里面是模型的输入 Tensor
|
self.original_height, self.original_width = img_bgr.shape[:2]
|
||||||
result_list = model.execute([input_data])
|
|
||||||
|
|
||||||
# 5. 后处理(针对 YOLO26 的端到端特性)
|
# 将图片从BGR转换为RGB
|
||||||
# YOLO26 输出的直接是最终预测框,result_list[0] 形状通常是 [1, 300, 6] 或类似
|
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
||||||
# 6个维度的含义一般为:[x_min, y_min, x_max, y_max, score, class_id]
|
|
||||||
predictions = result_list[0][0]
|
|
||||||
|
|
||||||
CONF_THRESHOLD = 0.25 # 设定置信度阈值
|
# 调整图片尺寸并保持宽高比,进行Letterbox填充
|
||||||
orig_h, orig_w = orig_img.shape[:2]
|
img_resized, self.ratio, (self.pad_left, self.pad_top) = self.letterbox(
|
||||||
|
img_rgb, (self.input_width, self.input_height)
|
||||||
|
)
|
||||||
|
|
||||||
box_count = 0
|
# 归一化:将像素值从 [0, 255] 转为 [0, 1]
|
||||||
for pred in predictions:
|
img_input = (img_resized.astype(np.float32) / 255.0).astype(np.float32)
|
||||||
# 取出置信度和类别
|
|
||||||
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)
|
|
||||||
|
|
||||||
print(f"推理完成,共检测到 {box_count} 个目标。")
|
# 调整维度:HWC -> NCHW,并转换为连续数组
|
||||||
|
img_input = np.ascontiguousarray(np.transpose(img_input, (2, 0, 1))[np.newaxis, :, :, :])
|
||||||
# 6. 保存检测结果图
|
|
||||||
cv2.imwrite("./result.jpg", orig_img)
|
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__":
|
if __name__ == "__main__":
|
||||||
main()
|
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}")
|
||||||
+175
@@ -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')
|
||||||
Reference in New Issue
Block a user