92 lines
2.9 KiB
Python
92 lines
2.9 KiB
Python
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
import platform
|
|||
|
|
from rknnlite.api import RKNNLite
|
|||
|
|
|
|||
|
|
# ------------------- 全局变量 -------------------
|
|||
|
|
_global_rknn_instance = None
|
|||
|
|
|
|||
|
|
labels = {0: '夹具未夹紧', 1: '夹具夹紧'}
|
|||
|
|
|
|||
|
|
DEVICE_COMPATIBLE_NODE = '/proc/device-tree/compatible'
|
|||
|
|
|
|||
|
|
def get_host():
|
|||
|
|
system = platform.system()
|
|||
|
|
machine = platform.machine()
|
|||
|
|
os_machine = system + '-' + machine
|
|||
|
|
if os_machine == 'Linux-aarch64':
|
|||
|
|
try:
|
|||
|
|
with open(DEVICE_COMPATIBLE_NODE) as f:
|
|||
|
|
device_compatible_str = f.read()
|
|||
|
|
if 'rk3562' in device_compatible_str:
|
|||
|
|
host = 'RK3562'
|
|||
|
|
elif 'rk3576' in device_compatible_str:
|
|||
|
|
host = 'RK3576'
|
|||
|
|
elif 'rk3588' in device_compatible_str:
|
|||
|
|
host = 'RK3588'
|
|||
|
|
else:
|
|||
|
|
host = 'RK3566_RK3568'
|
|||
|
|
except IOError:
|
|||
|
|
print('Read device node {} failed.'.format(DEVICE_COMPATIBLE_NODE))
|
|||
|
|
exit(-1)
|
|||
|
|
else:
|
|||
|
|
host = os_machine
|
|||
|
|
return host
|
|||
|
|
|
|||
|
|
# ------------------- 图像预处理 -------------------
|
|||
|
|
def preprocess(raw_image, target_size=(640, 640)):
|
|||
|
|
img = cv2.resize(raw_image, target_size)
|
|||
|
|
img = np.expand_dims(img, 0) # 添加 batch 维度
|
|||
|
|
return img
|
|||
|
|
|
|||
|
|
# ------------------- RKNN 模型初始化 -------------------
|
|||
|
|
def init_rknn_model(model_path):
|
|||
|
|
global _global_rknn_instance
|
|||
|
|
if _global_rknn_instance is None:
|
|||
|
|
rknn_lite = RKNNLite(verbose=False)
|
|||
|
|
ret = rknn_lite.load_rknn(model_path)
|
|||
|
|
if ret != 0:
|
|||
|
|
print(f'[ERROR] Load model failed (code: {ret})')
|
|||
|
|
exit(ret)
|
|||
|
|
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
|
|||
|
|
if ret != 0:
|
|||
|
|
print(f'[ERROR] Init runtime failed (code: {ret})')
|
|||
|
|
exit(ret)
|
|||
|
|
_global_rknn_instance = rknn_lite
|
|||
|
|
print(f'[INFO] Model loaded successfully: {model_path}')
|
|||
|
|
return _global_rknn_instance
|
|||
|
|
|
|||
|
|
# ------------------- 推理 -------------------
|
|||
|
|
def yolov11_cls_inference(model_path, raw_image, target_size=(640, 640)):
|
|||
|
|
"""
|
|||
|
|
返回：(class_id, boolean)
|
|||
|
|
类别 0 -> False
|
|||
|
|
类别 1 -> True
|
|||
|
|
"""
|
|||
|
|
img = preprocess(raw_image, target_size)
|
|||
|
|
rknn = init_rknn_model(model_path)
|
|||
|
|
outputs = rknn.inference([img])
|
|||
|
|
|
|||
|
|
# 获取类别 ID
|
|||
|
|
output = outputs[0].reshape(-1)
|
|||
|
|
class_id = int(np.argmax(output))
|
|||
|
|
bool_value = True if class_id == 1 else False
|
|||
|
|
|
|||
|
|
return class_id, bool_value
|
|||
|
|
|
|||
|
|
# ------------------- 测试 -------------------
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
image_path = "12.png"
|
|||
|
|
bgr_image = cv2.imread(image_path)
|
|||
|
|
if bgr_image is None:
|
|||
|
|
print(f"Failed to read image from {image_path}")
|
|||
|
|
exit(-1)
|
|||
|
|
|
|||
|
|
rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
|
|||
|
|
class_id, bool_value = yolov11_cls_inference(
|
|||
|
|
model_path="yolov11_cls.rknn",
|
|||
|
|
raw_image=rgb_image,
|
|||
|
|
target_size=(640, 640)
|
|||
|
|
)
|
|||
|
|
print(f"类别ID: {class_id}, 布尔值: {bool_value}")
|