first commit

rknn-multi-threaded-nosigmoid/func_cls.py

@@ -0,0 +1,70 @@
+import cv2
+import numpy as np
+
+# ------------------ 手动定义分类标签 ------------------
+labels = ["noready", "ready"]  # 0: 盖板不到位, 1: 盖板到位
+
+# ------------------ 配置 ------------------
+IMG_SIZE = (640, 640)
+
+
+def preprocess_cls(bgr_image_corrected, target_size=IMG_SIZE):
+    """
+    对已经校正方向（翻转后）的 BGR 图像进行预处理：
+      - 转 RGB
+      - Resize
+      - 添加 batch 维度
+    注意：此函数不再负责翻转！翻转应在调用前完成。
+    """
+    rgb = cv2.cvtColor(bgr_image_corrected, cv2.COLOR_BGR2RGB)
+    resized = cv2.resize(rgb, target_size)
+    input_tensor = np.expand_dims(resized, axis=0).astype(np.uint8)
+    return input_tensor
+
+
+def get_top1_class_str_from_output(outputs):
+    if outputs is None or len(outputs) == 0:
+        return "Error: No output"
+    logits = outputs[0].flatten()
+    top1_idx = int(np.argmax(logits))
+    if 0 <= top1_idx < len(labels):
+        return labels[top1_idx]
+    else:
+        return "Unknown"
+
+
+def myFunc(rknn_lite, IMG):
+    """
+    统一推理接口（分类任务）
+    - 输入原始 BGR 图像（可能倒置）
+    - 内部先上下翻转（校正方向）
+    - 模型在翻转后的图像上推理
+    - 返回带结果的翻转后图像（即“正”的图像）
+    """
+    # 1. 上下翻转原始图像（校正摄像头倒装）
+    corrected_img = cv2.flip(IMG, 0)  # 现在这是“正”的图像
+
+    # 2. 预处理（使用校正后的图像）
+    input_data = preprocess_cls(corrected_img, target_size=IMG_SIZE)
+
+    # 3. 推理
+    outputs = rknn_lite.inference(inputs=[input_data])
+
+    # 4. 解析结果
+    class_name = get_top1_class_str_from_output(outputs)
+
+    # 5. 在校正后的图像上绘制结果
+    vis_img = corrected_img.copy()
+    text = f"Class: {class_name}"
+    cv2.putText(
+        vis_img,
+        text,
+        (10, 30),
+        cv2.FONT_HERSHEY_SIMPLEX,
+        1.0,
+        (0, 255, 0),
+        2,
+        cv2.LINE_AA
+    )
+
+    return vis_img  # 返回的是上下翻转后的 BGR 图像（即视觉上“正”的图）