47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
import cv2
|
||
import numpy as np
|
||
from rknnlite import RKNNLite
|
||
|
||
# 初始化 RKNN
|
||
rknn = RKNNLite()
|
||
rknn.load_rknn("yolov8n.rknn")
|
||
rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
|
||
|
||
def preprocess(image):
|
||
# 调整尺寸并转换为 BCHW 格式
|
||
img = cv2.resize(image, (640, 640))
|
||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # BGR 转 RGB（根据模型配置）
|
||
img = img.transpose(2, 0, 1)[np.newaxis, ...].astype(np.uint8) # HWC -> BCHW
|
||
return img
|
||
|
||
# 加载测试图像
|
||
img = cv2.imread("test.jpg")
|
||
input_data = preprocess(img)
|
||
|
||
# 推理
|
||
outputs = rknn.inference(inputs=[input_data])[0] # 输出形状 [1, 84, 8400]
|
||
|
||
# 后处理（示例：提取检测框）
|
||
def postprocess(output, conf_thresh=0.5):
|
||
# 输出格式: [batch, 84, 8400]
|
||
# 84 = 4（坐标） + 80（类别）
|
||
boxes = []
|
||
for i in range(output.shape[2]):
|
||
scores = output[0, 4:84, i]
|
||
class_id = np.argmax(scores)
|
||
confidence = scores[class_id]
|
||
if confidence > conf_thresh:
|
||
cx, cy, w, h = output[0, :4, i]
|
||
# 转换为 x1,y1,x2,y2 格式
|
||
x1 = int((cx - w/2) * 640)
|
||
y1 = int((cy - h/2) * 640)
|
||
x2 = int((cx + w/2) * 640)
|
||
y2 = int((cy + h/2) * 640)
|
||
boxes.append([x1, y1, x2, y2, confidence, class_id])
|
||
return boxes
|
||
|
||
# 执行后处理
|
||
detections = postprocess(outputs)
|
||
print("检测结果:", detections)
|
||
|
||
rknn.release() |