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

ailai_pc/rename_image.py

@@ -52,5 +52,5 @@ def rename_images_in_folder(folder_path):
 # 🔧 使用这里：设置你的文件夹路径
 # ===========================
 if __name__ == "__main__":
-    folder = r"/home/hx/下载/2025-09-24"  # <-- 修改为你的图片文件夹路径
+    folder = r"/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/ailai_cls/train/class0"  # <-- 修改为你的图片文件夹路径
     rename_images_in_folder(folder)

cls_main/main_cls.py

@@ -5,11 +5,15 @@ from rknnlite.api import RKNNLite
 
 # ------------------- 全局变量 -------------------
 _global_rknn_instance = None
+labels = {0: '夹具夹紧', 1: '夹具打开'}
 
-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()
@@ -33,59 +37,63 @@ def get_host():
         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 模型初始化 -------------------
+# ------------------- 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)
+            raise RuntimeError(f'Load model failed: {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)
+            raise RuntimeError(f'Init runtime failed: {ret}')
         _global_rknn_instance = rknn_lite
-        print(f'[INFO] Model loaded successfully: {model_path}')
+        print(f'[INFO] RKNN model loaded: {model_path}')
     return _global_rknn_instance
 
-# ------------------- 推理 -------------------
-def yolov11_cls_inference(model_path, raw_image, target_size=(640, 640)):
+
+# ------------------- 图像预处理 + ROI 裁剪 -------------------
+def preprocess(raw_image, target_size=(640, 640)):
     """
-    返回：(class_id, boolean)
-    类别 0 -> False
-    类别 1 -> True
+    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)
-    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
-
+    bool_value = class_id == 1
     return class_id, bool_value
 
+
 # ------------------- 测试 -------------------
 if __name__ == '__main__':
-    image_path = "12.png"
+    image_path = "./test_image/class1/2.jpg"
+    model_path = "cls_rk3588.rknn"
+
     bgr_image = cv2.imread(image_path)
     if bgr_image is None:
-        print(f"Failed to read image from {image_path}")
-        exit(-1)
-
+        raise RuntimeError(f"Failed to read image: {image_path}")
     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)
-    )
+
+    # 只初始化一次模型
+    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}")

cls_main/readme.md

@@ -5,8 +5,8 @@
 
 类别定义：
 
-0 → 夹具未夹紧 (False)
-1 → 夹具夹紧 (True)
+0 → 夹具夹紧 (False)
+1 → 夹具打开 (True)
 
 rknn模型只加载一次，复用全局实例，提高推理效率。
 
@@ -14,10 +14,10 @@ rknn模型只加载一次，复用全局实例，提高推理效率。
 
 您可以直接调用 yolov11_cls_inference 函数，以便集成到其他项目中：
 
-示例 1: 测试仅获取单张图片的类别和布尔值
+示例 1: 单张图片推理
 
 ```bash
-from yolov11_cls_inference import yolov11_cls_inference
+from main_cls import yolov11_cls_inference
 import cv2
 
 # 读取图像
@@ -32,17 +32,27 @@ class_id, bool_value = yolov11_cls_inference(
 )
 
 print(f"类别ID: {class_id}, 布尔值: {bool_value}")
+
 ```
 
-示例 2: 直接在其他项目中集成使用
+示例 2: 多次推理（复用模型）
 ```bash
 
-from yolov11_cls_inference import yolov11_cls_inference
+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)
 
-# raw_image 已经读取或处理好的图像
-class_id, bool_value = yolov11_cls_inference(model_path="yolov11_cls.rknn", raw_image=raw_image)
 if bool_value:
     print("夹具夹紧")
 else:
-    print("夹具未夹紧")
+    print("夹具打开")
 ```