完成夹具开合判断的代码,更新了rknn3588上完成测试,3568在生产下次去测试的时候再进行测试
This commit is contained in:
Binary file not shown.
BIN
cls_main/cls_rk3588.rknn
Normal file
BIN
cls_main/cls_rk3588.rknn
Normal file
Binary file not shown.
99
cls_main/main_cls.py
Normal file
99
cls_main/main_cls.py
Normal 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
58
cls_main/readme.md
Normal 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("夹具打开")
|
||||
```
|
||||
BIN
cls_main/test_image/class0/1.jpg
Normal file
BIN
cls_main/test_image/class0/1.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 952 KiB |
BIN
cls_main/test_image/class1/1.jpg
Normal file
BIN
cls_main/test_image/class1/1.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 403 KiB |
BIN
cls_main/test_image/class1/2.jpg
Normal file
BIN
cls_main/test_image/class1/2.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 872 KiB |
Reference in New Issue
Block a user