完成夹具开合判断的代码,更新了rknn3588上完成测试,3568在生产下次去测试的时候再进行测试

This commit is contained in:
琉璃月光
2025-11-16 22:02:11 +08:00
parent 514ed6f1fd
commit 267aac1605
12 changed files with 973 additions and 0 deletions

Binary file not shown.

BIN
cls_main/cls_rk3588.rknn Normal file

Binary file not shown.

99
cls_main/main_cls.py Normal file
View File

@ -0,0 +1,99 @@
import cv2
import numpy as np
import platform
from rknnlite.api import RKNNLite
# ------------------- 全局变量 -------------------
_global_rknn_instance = None
labels = {0: '夹具夹紧', 1: '夹具打开'}
# ROI: x, y, w, h
ROI = (818, 175, 1381, 1271) # 示例
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
# ------------------- 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:
raise RuntimeError(f'Load model failed: {ret}')
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
if ret != 0:
raise RuntimeError(f'Init runtime failed: {ret}')
_global_rknn_instance = rknn_lite
print(f'[INFO] RKNN model loaded: {model_path}')
return _global_rknn_instance
# ------------------- 图像预处理 + ROI 裁剪 -------------------
def preprocess(raw_image, target_size=(640, 640)):
"""
ROI 裁剪 + resize + batch 维度
"""
global ROI
x, y, w, h = ROI
roi_img = raw_image[y:y+h, x:x+w]
img_resized = cv2.resize(roi_img, target_size)
img_batch = np.expand_dims(img_resized, 0) # 添加 batch 维度
return img_batch
# ------------------- 推理函数 -------------------
def yolov11_cls_inference_once(rknn, raw_image, target_size=(640, 640)):
"""
使用已加载的 rknn 实例进行推理
返回: (class_id, boolean)
"""
img = preprocess(raw_image, target_size)
outputs = rknn.inference([img])
output = outputs[0].reshape(-1)
class_id = int(np.argmax(output))
bool_value = class_id == 1
return class_id, bool_value
# ------------------- 测试 -------------------
if __name__ == '__main__':
image_path = "./test_image/class1/2.jpg"
model_path = "cls_rk3588.rknn"
bgr_image = cv2.imread(image_path)
if bgr_image is None:
raise RuntimeError(f"Failed to read image: {image_path}")
rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
# 只初始化一次模型
rknn_model = init_rknn_model(model_path)
# 多次调用都用同一个 rknn_model
class_id, bool_value = yolov11_cls_inference_once(rknn_model, rgb_image)
print(f"类别ID: {class_id}, 布尔值: {bool_value}")

58
cls_main/readme.md Normal file
View File

@ -0,0 +1,58 @@
# yolov11_cls_inference README
## 概述
该模块用于对米厂输入图像执行二分类推理,用于判断机械臂夹爪是否夹紧。
类别定义:
0 → 夹具夹紧 (False)
1 → 夹具打开 (True)
rknn模型只加载一次复用全局实例提高推理效率。
## 调用示例
您可以直接调用 yolov11_cls_inference 函数,以便集成到其他项目中:
示例 1: 单张图片推理
```bash
from main_cls import yolov11_cls_inference
import cv2
# 读取图像
bgr_image = cv2.imread("11.jpg")
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}")
```
示例 2: 多次推理(复用模型)
```bash
from main_cls import init_rknn_model, yolov11_cls_inference_once
import cv2
# 初始化一次模型
rknn_model = init_rknn_model("cls_rk3568.rknn")
# 读取图像
bgr_image = cv2.imread("12.jpg")
rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
# 使用已加载模型进行推理
class_id, bool_value = yolov11_cls_inference_once(rknn_model, rgb_image)
if bool_value:
print("夹具夹紧")
else:
print("夹具打开")
```

Binary file not shown.

After

Width:  |  Height:  |  Size: 952 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 403 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 872 KiB