Initial commit

README.md

@@ -0,0 +1,112 @@
+# 木条检测与分类 Python 接口示例
+
+本项目提供完整的 Python 示例，用于图像中木条数量检测、存在性判定以及 NG 木条异常分类。
+它支持：
+- 基于 RKNNLite 的木条检测
+- NG/OK 木条判定
+- NG 木条异常类型分类
+- 简单接口调用，支持本地图片推理
+
+---
+
+## 目录结构
+
+wood_detection/  
+├── wood_detect/  
+│   ├── wood_detect.py        # 木条数量检测接口  
+│   └── wood_detect.rknn      # RKNN模型文件  
+├── wood_exist/  
+│   ├── wood_exist.py         # 木条存在性判定接口  
+│   └── wood_exist_cls.rknn   # RKNN模型文件  
+├── wood_ng/  
+│   ├── wood_ng.py            # NG木条异常分类接口  
+│   └── wood_ng_cls.rknn      # RKNN模型文件  
+└── README.md                 # 说明文档  
+---
+
+## 配置
+
+### 安装依赖
+```bash
+pip install opencv-python numpy rknnlite
+
+```
+
+## 接口说明
+
+### 1. 木条数量检测
+
+函数：detect_wood(img: np.ndarray) -> int
+
+参数：
+
+img：BGR格式图像 (np.ndarray)
+
+返回值：
+
+检测到的木条数量 (int)
+
+示例：
+
+#### 函数调用
+
+```python
+import cv2
+from wood_detect.wood_detect import detect_wood
+
+img = cv2.imread("1.jpg")
+result = detect_wood(img)
+print(f"检测到木条数量: {result}")
+
+```
+
+### 2. 木条存在性判定
+
+函数：classify_wood_exist(img: np.ndarray) -> int
+
+参数：
+
+img：BGR格式图像 (np.ndarray)
+
+返回值：
+
+整数结果，对应木条状态（int: 存在结果（0 / 1））
+
+示例：
+
+```python
+import cv2
+from wood_exist.wood_exist import classify_wood_exist, CLASS_NAMES
+
+img = cv2.imread("1.png")
+result = classify_wood_exist(img)
+print(f"木条存在性判定: {result}")
+
+```
+
+### 3. NG 木条异常分类
+
+函数：classify_wood_ng(img: np.ndarray) -> int
+
+参数：
+
+img：BGR格式图像 (np.ndarray)
+
+返回值：
+
+整数结果，对应 NG木条（int: NG结果（0 / 1））
+
+示例：
+
+```python
+import cv2
+from wood_ng.wood_ng import classify_wood_ng, CLASS_NAMES
+
+img = cv2.imread("1.png")
+result = classify_wood_ng(img)
+print(f"NG 木条: {result}")
+
+```
+
+
+

wood_detect/wood_detect.py

@@ -0,0 +1,212 @@
+# -*- coding: utf-8 -*-
+"""
+木条检测模块（基于RKNNLite + ROI + 单类别NMS）
+功能：
+- 检测ROI区域内木条数量
+- 输出实际木条数量
+"""
+
+import os
+from typing import Optional
+
+import cv2
+import numpy as np
+from rknnlite.api import RKNNLite
+
+# =====================================================
+# 常量配置
+# =====================================================
+RKNN_MODEL_PATH: str = "wood_detect.rknn"
+ROI: tuple[int, int, int, int] = (1, 1, 1000, 1000)  # ROI 坐标(x1, y1, x2, y2)
+IMG_SIZE: tuple[int, int] = (640, 640)  # 模型输入大小
+OBJ_THRESH: float = 0.25  # 置信度阈值
+NMS_THRESH: float = 0.45  # NMS阈值
+CLASS_NAME: list[str] = ["bag"]  # 单类别名称，后续改成wood
+
+# =====================================================
+# 私有工具函数
+# =====================================================
+def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
+    """计算softmax概率"""
+    x = x - np.max(x, axis=axis, keepdims=True)
+    exp_x = np.exp(x)
+    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
+
+
+def _letterbox_resize(image: np.ndarray,
+                      size: tuple[int, int],
+                      bg_color: int = 114) -> tuple[np.ndarray, float, int, int]:
+    """保持长宽比缩放并填充到目标尺寸"""
+    target_w, target_h = size
+    h, w = image.shape[:2]
+    scale = min(target_w / w, target_h / h)
+    new_w, new_h = int(w * scale), int(h * scale)
+    resized = cv2.resize(image, (new_w, new_h))
+    canvas = np.full((target_h, target_w, 3), bg_color, dtype=np.uint8)
+    dx = (target_w - new_w) // 2
+    dy = (target_h - new_h) // 2
+    canvas[dy:dy + new_h, dx:dx + new_w] = resized
+    return canvas, scale, dx, dy
+
+def _nms(boxes: np.ndarray, scores: np.ndarray, thresh: float) -> list[int]:
+    """非极大值抑制"""
+    x1, y1, x2, y2 = boxes.T
+    areas = (x2 - x1) * (y2 - y1)
+    order = scores.argsort()[::-1]
+
+    keep: list[int] = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+        inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
+        iou = inter / (areas[i] + areas[order[1:]] - inter)
+        order = order[1:][iou <= thresh]
+    return keep
+
+def _post_process(outputs: list[np.ndarray],
+                  scale: float,
+                  dx: int,
+                  dy: int) -> Optional[np.ndarray]:
+    """RKNN模型输出后处理：解码坐标 + NMS，返回有效boxes"""
+    boxes_list, scores_list = [], []
+    strides = [8, 16, 32]
+
+    for i, stride in enumerate(strides):
+        reg = outputs[i * 3 + 0][0]
+        cls = outputs[i * 3 + 1][0]
+        obj = outputs[i * 3 + 2][0]
+
+        num_classes, H, W = cls.shape
+        reg_max = reg.shape[0] // 4
+
+        grid_y, grid_x = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
+        grid_x = grid_x.astype(np.float32).ravel()
+        grid_y = grid_y.astype(np.float32).ravel()
+
+        cls_flat = cls.reshape(num_classes, -1).T
+        obj_flat = obj.ravel()
+        max_cls_score = np.max(cls_flat, axis=1)
+        scores = max_cls_score * obj_flat
+
+        valid_mask = scores >= OBJ_THRESH
+        if not np.any(valid_mask):
+            continue
+
+        valid_idx = np.where(valid_mask)[0]
+        scores_v = scores[valid_idx]
+        gx = grid_x[valid_idx]
+        gy = grid_y[valid_idx]
+
+        reg_valid = reg.reshape(4, reg_max, -1)[:, :, valid_idx]
+        reg_softmax = _softmax(reg_valid, axis=1)
+        acc = np.arange(reg_max, dtype=np.float32).reshape(1, -1, 1)
+        distance = np.sum(reg_softmax * acc, axis=1)
+
+        cx = (gx + 0.5) * stride
+        cy = (gy + 0.5) * stride
+        l, t, r, b = distance[0], distance[1], distance[2], distance[3]
+        x1 = cx - l * stride
+        y1 = cy - t * stride
+        x2 = cx + r * stride
+        y2 = cy + b * stride
+
+        boxes = np.stack([x1, y1, x2, y2], axis=1)
+        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - dx) / scale
+        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - dy) / scale
+
+        boxes_list.append(boxes)
+        scores_list.append(scores_v)
+
+    if not boxes_list:
+        return None
+
+    boxes_all = np.concatenate(boxes_list, axis=0)
+    scores_all = np.concatenate(scores_list, axis=0)
+    keep_idx = _nms(boxes_all, scores_all, NMS_THRESH)
+    return boxes_all[keep_idx]
+
+# =====================================================
+# RKNN模型全局初始化
+# =====================================================
+_global_rknn = RKNNLite()
+_global_rknn.load_rknn(RKNN_MODEL_PATH)
+_global_rknn.init_runtime()
+
+# =====================================================
+# 木条检测类
+# =====================================================
+class WoodDetectorRKNN:
+    """基于RKNN的木条检测器"""
+
+    def __init__(self):
+        """初始化木条检测器"""
+        self.rknn = _global_rknn
+
+    def detect(self, img_np: np.ndarray) -> int:
+        """
+        检测木条数量
+
+        Args:
+            img_np (np.ndarray): 原始BGR图像
+
+        Returns:
+            int: 检测到的木条数量
+        """
+        # ROI裁剪
+        h, w = img_np.shape[:2]
+        x1, y1, x2, y2 = ROI
+        x1 = max(0, min(x1, w - 1))
+        x2 = max(0, min(x2, w))
+        y1 = max(0, min(y1, h - 1))
+        y2 = max(0, min(y2, h))
+        roi_img = img_np[y1:y2, x1:x2]
+
+        # letterbox resize
+        resized_img, scale, dx, dy = _letterbox_resize(roi_img, IMG_SIZE)
+
+        # 推理
+        input_tensor = np.expand_dims(resized_img.astype(np.float32), axis=0)
+        outputs = self.rknn.inference([input_tensor])
+        boxes = _post_process(outputs, scale, dx, dy)
+
+        if boxes is None:
+            return 0
+
+        return len(boxes)
+
+# =====================================================
+# 对外统一接口
+# =====================================================
+_detector = WoodDetectorRKNN()
+
+def detect_wood(img_np: np.ndarray) -> int:
+    """
+    对外木条检测接口（返回木条数量）
+
+    Args:
+        img_np (np.ndarray): 原始BGR图像
+
+    Returns:
+        int: 检测到的木条数量
+    """
+    return _detector.detect(img_np)
+
+# =====================================================
+# 主程序入口
+# =====================================================
+if __name__ == "__main__":
+    img_path = "1.jpg"
+
+    if not os.path.exists(img_path):
+        raise FileNotFoundError(f"图片不存在: {img_path}")
+
+    img = cv2.imread(img_path)
+    if img is None:
+        raise ValueError("图像加载失败")
+
+    total_count = detect_wood(img)
+    print(f"检测到木条数量: {total_count}")

wood_exist/wood_exist.py

@@ -0,0 +1,207 @@
+# -*- coding: utf-8 -*-
+"""
+ROI RKNN 图像分类模块
+
+基于 RKNNLite 对输入图像的指定 ROI 区域进行分类，
+输出类别 ID（0: 异常，1: 正常）。
+
+支持：
+- RKNN 模型单例加载
+- ROI 裁剪与缩放
+- BGR → RGB 预处理
+- 主程序测试入口
+"""
+
+import os
+from typing import Dict
+
+import cv2
+import numpy as np
+from rknnlite.api import RKNNLite
+
+# =====================================================
+# 全局配置（常量）
+# =====================================================
+
+RKNN_MODEL_PATH: str = "wood_exist_cls.rknn"
+
+# ROI 坐标：x1, y1, x2, y2（像素坐标）
+ROI: tuple[int, int, int, int] = (3, 0, 694, 182)
+
+CLASS_NAMES: Dict[int, str] = {
+    0: "异常",
+    1: "正常",
+}
+
+# =====================================================
+# 全局 RKNN 实例（单例）
+# =====================================================
+
+_global_rknn: RKNNLite | None = None
+
+
+def _init_rknn_model(model_path: str) -> RKNNLite:
+    """
+    初始化并返回 RKNN 模型（单例模式）。
+
+    Args:
+        model_path (str): RKNN 模型路径
+
+    Returns:
+        RKNNLite: 已初始化的 RKNNLite 实例
+
+    Raises:
+        FileNotFoundError: 模型文件不存在
+        RuntimeError: RKNN 加载或运行时初始化失败
+    """
+    global _global_rknn
+
+    if _global_rknn is not None:
+        return _global_rknn
+
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"RKNN 模型不存在: {model_path}")
+
+    rknn = RKNNLite(verbose=False)
+
+    ret = rknn.load_rknn(model_path)
+    if ret != 0:
+        raise RuntimeError(f"Load RKNN failed: {ret}")
+
+    ret = rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
+    if ret != 0:
+        raise RuntimeError(f"Init runtime failed: {ret}")
+
+    _global_rknn = rknn
+    print(f"[INFO] RKNN 模型加载成功: {model_path}")
+
+    return rknn
+
+
+# =====================================================
+# 预处理函数
+# =====================================================
+
+def _preprocess_input(img: np.ndarray) -> np.ndarray:
+    """
+    对 ROI 图像进行模型输入预处理。
+
+    Args:
+        img (np.ndarray): BGR 格式图像，shape=(640, 640, 3)
+
+    Returns:
+        np.ndarray: NHWC 格式 float32 输入张量，shape=(1, 640, 640, 3)
+
+    Raises:
+        ValueError: 输入图像尺寸不符合要求
+    """
+    if img.shape[:2] != (640, 640):
+        raise ValueError("输入图像必须是 640x640")
+
+    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    input_tensor = np.expand_dims(img_rgb.astype(np.float32), axis=0)
+
+    return np.ascontiguousarray(input_tensor)
+
+
+# =====================================================
+# ROI + RKNN 分类器
+# =====================================================
+
+class ROIClassifierRKNN:
+    """
+    基于 RKNN 的 ROI 区域分类器。
+
+    功能：
+    - 加载 RKNN 分类模型
+    - 从原始图像中裁剪 ROI
+    - 执行分类推理并返回类别 ID
+    """
+
+    def __init__(self, model_path: str) -> None:
+        """
+        初始化分类器。
+
+        Args:
+            model_path (str): RKNN 模型路径
+        """
+        self.rknn = _init_rknn_model(model_path)
+
+    def classify(self, img_np: np.ndarray) -> int:
+        """
+        对输入图像进行 ROI 分类。
+
+        Args:
+            img_np (np.ndarray): 原始 BGR 图像
+
+        Returns:
+            int: 分类结果（0: 异常，1: 正常）
+
+        Raises:
+            ValueError: ROI 坐标非法
+        """
+        height, width = img_np.shape[:2]
+        x1, y1, x2, y2 = ROI
+
+        # -------- ROI 边界保护 --------
+        x1 = max(0, min(x1, width - 1))
+        x2 = max(0, min(x2, width))
+        y1 = max(0, min(y1, height - 1))
+        y2 = max(0, min(y2, height))
+
+        if x2 <= x1 or y2 <= y1:
+            raise ValueError(f"ROI 坐标无效: {(x1, y1, x2, y2)}")
+
+        # -------- 1. 裁剪 ROI --------
+        roi_img = img_np[y1:y2, x1:x2]
+
+        # -------- 2. resize 到 640×640 --------
+        roi_img = cv2.resize(roi_img, (640, 640))
+
+        # -------- 3. 预处理 --------
+        input_tensor = _preprocess_input(roi_img)
+
+        # -------- 4. RKNN 推理 --------
+        outputs = self.rknn.inference([input_tensor])
+        logits = outputs[0].reshape(-1).astype(float)
+
+        return int(np.argmax(logits))
+
+
+# =====================================================
+# 对外接口
+# =====================================================
+
+_classifier = ROIClassifierRKNN(RKNN_MODEL_PATH)
+
+def classify_wood_exist(img_np: np.ndarray) -> int:
+    """
+    线条是否存在图像分类接口函数。
+
+    Args:
+        img_np (np.ndarray): 原始 BGR 图像
+
+    Returns:
+        int: 分类结果（0 / 1）
+    """
+    return _classifier.classify(img_np)
+
+# =====================================================
+# 测试入口
+# =====================================================
+
+if __name__ == "__main__":
+    img_path = "1.png"
+
+    if not os.path.exists(img_path):
+        raise FileNotFoundError(f"图片不存在: {img_path}")
+
+    img = cv2.imread(img_path)
+    if img is None:
+        raise ValueError("图像加载失败")
+
+    result = classify_wood_exist(img)
+    print(
+        f"NG料结果：{result} "
+        f"({CLASS_NAMES[result]})"
+    )

wood_ng/wood_ng.py

@@ -0,0 +1,207 @@
+# -*- coding: utf-8 -*-
+"""
+ROI RKNN 图像分类模块
+
+基于 RKNNLite 对输入图像的指定 ROI 区域进行分类，
+输出类别 ID（0: 异常，1: 正常）。
+
+支持：
+- RKNN 模型单例加载
+- ROI 裁剪与缩放
+- BGR → RGB 预处理
+- 主程序测试入口
+"""
+
+import os
+from typing import Dict
+
+import cv2
+import numpy as np
+from rknnlite.api import RKNNLite
+
+# =====================================================
+# 全局配置（常量）
+# =====================================================
+
+RKNN_MODEL_PATH: str = "wood_ng_cls.rknn"
+
+# ROI 坐标：x1, y1, x2, y2（像素坐标）
+ROI: tuple[int, int, int, int] = (3, 0, 694, 182)
+
+CLASS_NAMES: Dict[int, str] = {
+    0: "异常",
+    1: "正常",
+}
+
+# =====================================================
+# 全局 RKNN 实例（单例）
+# =====================================================
+
+_global_rknn: RKNNLite | None = None
+
+
+def _init_rknn_model(model_path: str) -> RKNNLite:
+    """
+    初始化并返回 RKNN 模型（单例模式）。
+
+    Args:
+        model_path (str): RKNN 模型路径
+
+    Returns:
+        RKNNLite: 已初始化的 RKNNLite 实例
+
+    Raises:
+        FileNotFoundError: 模型文件不存在
+        RuntimeError: RKNN 加载或运行时初始化失败
+    """
+    global _global_rknn
+
+    if _global_rknn is not None:
+        return _global_rknn
+
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"RKNN 模型不存在: {model_path}")
+
+    rknn = RKNNLite(verbose=False)
+
+    ret = rknn.load_rknn(model_path)
+    if ret != 0:
+        raise RuntimeError(f"Load RKNN failed: {ret}")
+
+    ret = rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
+    if ret != 0:
+        raise RuntimeError(f"Init runtime failed: {ret}")
+
+    _global_rknn = rknn
+    print(f"[INFO] RKNN 模型加载成功: {model_path}")
+
+    return rknn
+
+
+# =====================================================
+# 预处理函数
+# =====================================================
+
+def _preprocess_input(img: np.ndarray) -> np.ndarray:
+    """
+    对 ROI 图像进行模型输入预处理。
+
+    Args:
+        img (np.ndarray): BGR 格式图像，shape=(640, 640, 3)
+
+    Returns:
+        np.ndarray: NHWC 格式 float32 输入张量，shape=(1, 640, 640, 3)
+
+    Raises:
+        ValueError: 输入图像尺寸不符合要求
+    """
+    if img.shape[:2] != (640, 640):
+        raise ValueError("输入图像必须是 640x640")
+
+    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    input_tensor = np.expand_dims(img_rgb.astype(np.float32), axis=0)
+
+    return np.ascontiguousarray(input_tensor)
+
+
+# =====================================================
+# ROI + RKNN 分类器
+# =====================================================
+
+class ROIClassifierRKNN:
+    """
+    基于 RKNN 的 ROI 区域分类器。
+
+    功能：
+    - 加载 RKNN 分类模型
+    - 从原始图像中裁剪 ROI
+    - 执行分类推理并返回类别 ID
+    """
+
+    def __init__(self, model_path: str) -> None:
+        """
+        初始化分类器。
+
+        Args:
+            model_path (str): RKNN 模型路径
+        """
+        self.rknn = _init_rknn_model(model_path)
+
+    def classify(self, img_np: np.ndarray) -> int:
+        """
+        对输入图像进行 ROI 分类。
+
+        Args:
+            img_np (np.ndarray): 原始 BGR 图像
+
+        Returns:
+            int: 分类结果（0: 异常，1: 正常）
+
+        Raises:
+            ValueError: ROI 坐标非法
+        """
+        height, width = img_np.shape[:2]
+        x1, y1, x2, y2 = ROI
+
+        # -------- ROI 边界保护 --------
+        x1 = max(0, min(x1, width - 1))
+        x2 = max(0, min(x2, width))
+        y1 = max(0, min(y1, height - 1))
+        y2 = max(0, min(y2, height))
+
+        if x2 <= x1 or y2 <= y1:
+            raise ValueError(f"ROI 坐标无效: {(x1, y1, x2, y2)}")
+
+        # -------- 1. 裁剪 ROI --------
+        roi_img = img_np[y1:y2, x1:x2]
+
+        # -------- 2. resize 到 640×640 --------
+        roi_img = cv2.resize(roi_img, (640, 640))
+
+        # -------- 3. 预处理 --------
+        input_tensor = _preprocess_input(roi_img)
+
+        # -------- 4. RKNN 推理 --------
+        outputs = self.rknn.inference([input_tensor])
+        logits = outputs[0].reshape(-1).astype(float)
+
+        return int(np.argmax(logits))
+
+
+# =====================================================
+# 对外接口
+# =====================================================
+
+_classifier = ROIClassifierRKNN(RKNN_MODEL_PATH)
+
+def classify_wood_ng(img_np: np.ndarray) -> int:
+    """
+    线条是否为ng料图像分类接口函数。
+
+    Args:
+        img_np (np.ndarray): 原始 BGR 图像
+
+    Returns:
+        int: 分类结果（0 / 1）
+    """
+    return _classifier.classify(img_np)
+
+# =====================================================
+# 测试入口
+# =====================================================
+
+if __name__ == "__main__":
+    img_path = "1.png"
+
+    if not os.path.exists(img_path):
+        raise FileNotFoundError(f"图片不存在: {img_path}")
+
+    img = cv2.imread(img_path)
+    if img is None:
+        raise ValueError("图像加载失败")
+
+    result = classify_wood_ng(img)
+    print(
+        f"NG料结果：{result} "
+        f"({CLASS_NAMES[result]})"
+    )