npu端基本测试
This commit is contained in:
+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