commit beb8c8f1e69fa2eb8543d2531ae60cbe3c48822e Author: llyg777 Date: Wed Feb 25 14:24:05 2026 +0800 Initial commit diff --git a/README.md b/README.md new file mode 100644 index 0000000..5c3c026 --- /dev/null +++ b/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}") + +``` + + + diff --git a/wood_detect/1.jpg b/wood_detect/1.jpg new file mode 100644 index 0000000..e69faee Binary files /dev/null and b/wood_detect/1.jpg differ diff --git a/wood_detect/wood_detect.py b/wood_detect/wood_detect.py new file mode 100644 index 0000000..1fe03a7 --- /dev/null +++ b/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}") diff --git a/wood_detect/wood_detect.rknn b/wood_detect/wood_detect.rknn new file mode 100644 index 0000000..7687df8 Binary files /dev/null and b/wood_detect/wood_detect.rknn differ diff --git a/wood_exist/1.png b/wood_exist/1.png new file mode 100644 index 0000000..7395075 Binary files /dev/null and b/wood_exist/1.png differ diff --git a/wood_exist/wood_exist.py b/wood_exist/wood_exist.py new file mode 100644 index 0000000..9766600 --- /dev/null +++ b/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]})" + ) diff --git a/wood_exist/wood_exist_cls.rknn b/wood_exist/wood_exist_cls.rknn new file mode 100644 index 0000000..cfbdda8 Binary files /dev/null and b/wood_exist/wood_exist_cls.rknn differ diff --git a/wood_ng/1.png b/wood_ng/1.png new file mode 100644 index 0000000..7395075 Binary files /dev/null and b/wood_ng/1.png differ diff --git a/wood_ng/wood_ng.py b/wood_ng/wood_ng.py new file mode 100644 index 0000000..a854fce --- /dev/null +++ b/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]})" + ) diff --git a/wood_ng/wood_ng_cls.rknn b/wood_ng/wood_ng_cls.rknn new file mode 100644 index 0000000..cfbdda8 Binary files /dev/null and b/wood_ng/wood_ng_cls.rknn differ