feeding

vision/overflow_model/README.md

@@ -0,0 +1,78 @@
+# RKNN 堆料分类推理系统 README
+
+本项目用于在 RK3588 平台上运行 RKNN 分类模型，对多个 ROI 区域进行堆料状态分类，包括：
+
+未堆料     0
+小堆料     1
+大堆料     2 
+未浇筑满   3
+浇筑满     4
+
+项目中支持 多 ROI 裁剪、模型推理、加权判断（小/大堆料） 和分类结果输出。
+
+## 目录结构
+
+project/
+│── yiliao_cls.rknn                # RKNN 模型
+│── best.pt                        # pt 模型
+│── roi_coordinates/         # ROI 坐标文件目录
+│     └── 1_rois.txt
+│── test_image/              # 测试图片目录
+│     └── 1.jpg
+      └── 2.jpg
+      └── 3.jpg
+│── yiliao_main_rknn.py        # RKNN主推理脚本
+│── yiliao_main_pc.py          # PC推理脚本 
+│── README.md
+
+
+## 配置（略）
+## 安装依赖（略）
+
+
+## 调用示例
+单张图片推理调用示例
+
+```bash
+
+from yiliao_main_rknn import classify_frame_with_rois
+
+if __name__ == "__main__":
+    model_path = "yiliao_cls.rknn"
+    roi_file = "./roi_coordinates/1_rois.txt"
+
+    frame = cv2.imread("./test_image/2.jpg")
+
+    outputs = classify_frame_with_rois(model_path, frame, roi_file)
+
+    for res in outputs:
+        print(res)
+
+
+```
+
+##小堆料 / 大堆料加权判定说明
+
+模型原始输出中，小堆料（class 1）与大堆料（class 2）相比时容易出现概率接近的情况。
+
+通过加权机制：
+
+✔ 可以避免因整体概率偏低导致分类不稳定
+✔ 优先放大“大堆料 的可能性”（因为 w2 > w1）
+✔ score 更能反映堆料大小的趋势，而不是绝对概率
+
+为提高判断稳定性，采用了加权评分方式：（这些参数都可以根据实际情况在文件中对weighted_small_large中参数进行修改）
+score = (0.3 * p1 + 0.7 * p2) / (p1 + p2)
+score ≥ 0.4 → 大堆料
+score < 0.4 → 小堆料
+
+p1：小堆料概率
+p2：大堆料概率
+score 越接近 1 越倾向于大堆料
+score 越接近 0 越倾向于小堆料
+
+
+
+
+
+

vision/overflow_model/roi_coordinates/1_rois.txt

@@ -0,0 +1 @@
+859,810,696,328

vision/overflow_model/yiliao_main_pc.py

@@ -0,0 +1,168 @@
+import os
+from pathlib import Path
+import cv2
+import numpy as np
+from ultralytics import YOLO
+
+# ---------------------------
+# 类别映射
+# ---------------------------
+CLASS_NAMES = {
+    0: "未堆料",
+    1: "小堆料",
+    2: "大堆料",
+    3: "未浇筑满",
+    4: "浇筑满"
+}
+
+# ---------------------------
+# 加载 ROI 列表
+# ---------------------------
+def load_global_rois(txt_path):
+    rois = []
+    if not os.path.exists(txt_path):
+        print(f"❌ ROI 文件不存在: {txt_path}")
+        return rois
+    with open(txt_path, 'r') as f:
+        for line in f:
+            s = line.strip()
+            if s:
+                try:
+                    x, y, w, h = map(int, s.split(','))
+                    rois.append((x, y, w, h))
+                except Exception as e:
+                    print(f"无法解析 ROI 行 '{s}': {e}")
+    return rois
+
+# ---------------------------
+# 裁剪并 resize ROI
+# ---------------------------
+def crop_and_resize(img, rois, target_size=640):
+    crops = []
+    h_img, w_img = img.shape[:2]
+    for i, (x, y, w, h) in enumerate(rois):
+        if x < 0 or y < 0 or x + w > w_img or y + h > h_img:
+            continue
+        roi = img[y:y+h, x:x+w]
+        roi_resized = cv2.resize(roi, (target_size, target_size), interpolation=cv2.INTER_AREA)
+        crops.append((roi_resized, i))
+    return crops
+
+# ---------------------------
+# class1/class2 加权判断
+# ---------------------------
+def weighted_small_large(pred_probs, threshold=0.4, w1=0.3, w2=0.7):
+    p1 = float(pred_probs[1])
+    p2 = float(pred_probs[2])
+    total = p1 + p2
+    if total > 0:
+        score = (w1 * p1 + w2 * p2) / total
+    else:
+        score = 0.0
+    final_class = "大堆料" if score >= threshold else "小堆料"
+    return final_class, score, p1, p2
+
+# ---------------------------
+# 单张图片推理函数
+# ---------------------------
+def classify_image_weighted(image, model, threshold=0.4):
+    results = model(image)
+    pred_probs = results[0].probs.data.cpu().numpy().flatten()
+    class_id = int(pred_probs.argmax())
+    confidence = float(pred_probs[class_id])
+    class_name = CLASS_NAMES.get(class_id, f"未知类别({class_id})")
+
+    # class1/class2 使用加权得分
+    if class_id in [1, 2]:
+        final_class, score, p1, p2 = weighted_small_large(pred_probs, threshold=threshold)
+    else:
+        final_class = class_name
+        score = confidence
+        p1 = float(pred_probs[1])
+        p2 = float(pred_probs[2])
+
+    return final_class, score, p1, p2
+
+# ---------------------------
+# 批量推理主函数
+# ---------------------------
+def batch_classify_images(model_path, input_folder, output_root, roi_file, target_size=640, threshold=0.5):
+    # 加载模型
+    model = YOLO(model_path)
+
+    # 确保输出根目录存在
+    output_root = Path(output_root)
+    output_root.mkdir(parents=True, exist_ok=True)
+
+    # 为所有类别创建目录
+    class_dirs = {}
+    for name in CLASS_NAMES.values():
+        d = output_root / name
+        d.mkdir(exist_ok=True)
+        class_dirs[name] = d
+
+    rois = load_global_rois(roi_file)
+    if not rois:
+        print("❌ 没有有效 ROI，退出")
+        return
+
+    # 遍历图片
+    for img_path in Path(input_folder).glob("*.*"):
+        if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tif']:
+            continue
+        try:
+            img = cv2.imread(str(img_path))
+            if img is None:
+                continue
+
+            crops = crop_and_resize(img, rois, target_size)
+
+            for roi_resized, roi_idx in crops:
+                final_class, score, p1, p2 = classify_image_weighted(roi_resized, model, threshold=threshold)
+
+                # 文件名中保存 ROI、类别、加权分数、class1/class2 置信度
+                suffix = f"_roi{roi_idx}_{final_class}_score{score:.2f}_p1{p1:.2f}_p2{p2:.2f}"
+                dst_path = class_dirs[final_class] / f"{img_path.stem}{suffix}{img_path.suffix}"
+                cv2.imwrite(dst_path, roi_resized)
+                print(f"{img_path.name}{suffix} -> {final_class} (score={score:.2f}, p1={p1:.2f}, p2={p2:.2f})")
+
+        except Exception as e:
+            print(f"处理失败 {img_path.name}: {e}")
+
+
+# ---------------------------
+# 单张图片使用示例（保留 ROI，不保存文件）
+# ---------------------------
+if __name__ == "__main__":
+    model_path = r"best.pt"
+    image_path = r"./test_image/2.jpg"  # 单张图片路径
+    roi_file = r"./roi_coordinates/1_rois.txt"
+    target_size = 640
+    threshold = 0.4  #加权得分阈值可以根据大小堆料分类结果进行调整
+
+    # 加载模型
+    model = YOLO(model_path)
+
+    # 读取 ROI
+    rois = load_global_rois(roi_file)
+    if not rois:
+        print("❌ 没有有效 ROI，退出")
+        exit(1)
+
+    # 读取图片
+    img = cv2.imread(image_path)
+    if img is None:
+        print(f"❌ 无法读取图片: {image_path}")
+        exit(1)
+
+    # 注意：必须裁剪 ROI 并推理，因为训练的时候输入的图像是经过resize的
+    crops = crop_and_resize(img, rois, target_size)
+    for roi_resized, roi_idx in crops:
+        #final_class, score, p1, p2 = classify_image_weighted(roi_resized, model, threshold=threshold)
+        final_class,_,_,_ = classify_image_weighted(roi_resized, model, threshold=threshold)
+        # 只输出信息，不保存文件
+        #print(f"ROI {roi_idx} -> 类别: {final_class}, 加权分数: {score:.2f}, "
+              #f"class1 置信度: {p1:.2f}, class2 置信度: {p2:.2f}")
+        print(f"类别: {final_class}")
+
+

vision/overflow_model/yiliao_main_rknn.py

@@ -0,0 +1,185 @@
+import os
+from pathlib import Path
+import cv2
+import numpy as np
+import platform
+
+
+# ---------------------------
+# 类别映射
+# ---------------------------
+CLASS_NAMES = {
+    0: "未堆料",
+    1: "小堆料",
+    2: "大堆料",
+    3: "未浇筑满",
+    4: "浇筑满"
+}
+
+# ---------------------------
+# RKNN 全局实例（只加载一次）
+# ---------------------------
+_global_rknn = None
+DEVICE_COMPATIBLE_NODE = '/proc/device-tree/compatible'
+
+
+# =====================================================
+#                   RKNN MODEL
+# =====================================================
+def init_rknn_model(model_path):
+    from rknnlite.api import RKNNLite
+    
+    global _global_rknn
+    if _global_rknn is not None:
+        return _global_rknn
+
+    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
+
+
+# ---------------------------
+# 图像预处理（统一 640×640）
+# ---------------------------
+def preprocess(img, size=(640, 640)):
+    img = cv2.resize(img, size)
+    img = np.expand_dims(img, 0)
+    return img
+
+
+# ---------------------------
+# 单次 RKNN 分类
+# ---------------------------
+def rknn_classify(img_resized, model_path):
+    rknn = init_rknn_model(model_path)
+    input_tensor = preprocess(img_resized)
+    outs = rknn.inference([input_tensor])
+
+    pred = outs[0].reshape(-1)
+    class_id = int(np.argmax(pred))
+    return class_id, pred.astype(float)
+
+
+# =====================================================
+#                    ROI 逻辑
+# =====================================================
+def load_rois(txt_path):
+    rois = []
+    if not os.path.exists(txt_path):
+        print(f"❌ ROI 文件不存在: {txt_path}")
+        return rois
+
+    with open(txt_path) as f:
+        for line in f:
+            s = line.strip()
+            if s:
+                try:
+                    x, y, w, h = map(int, s.split(','))
+                    rois.append((x, y, w, h))
+                except:
+                    print("ROI 格式错误:", s)
+    return rois
+
+
+def crop_and_resize(img, rois, target_size=640):
+    crops = []
+    h_img, w_img = img.shape[:2]
+
+    for idx, (x, y, w, h) in enumerate(rois):
+        if x < 0 or y < 0 or x + w > w_img or y + h > h_img:
+            continue
+        roi = img[y:y + h, x:x + w]
+        roi_resized = cv2.resize(roi, (target_size, target_size), interpolation=cv2.INTER_AREA)
+        crops.append((roi_resized, idx))
+    return crops
+
+
+# =====================================================
+#            class1/class2 加权分类增强
+# =====================================================
+def weighted_small_large(pred, threshold=0.4, w1=0.3, w2=0.7):
+    p1 = float(pred[1])
+    p2 = float(pred[2])
+    total = p1 + p2
+
+    score = (w1 * p1 + w2 * p2) / total if total > 0 else 0.0
+    final_class = "大堆料" if score >= threshold else "小堆料"
+
+    return final_class, score, p1, p2
+
+
+# =====================================================
+#             ⭐ 高复用：一行完成 ROI + 推理 ⭐
+# =====================================================
+def classify_frame_with_rois(model_path, frame, roi_file, threshold=0.4):
+    """
+    输入:
+        - frame: BGR 图像 (numpy array)
+        - model_path: RKNN 模型路径
+        - roi_file: ROI 的 txt 文件
+        - threshold: class1/class2 小/大堆料判断阈值
+
+    输出:
+        [
+            { "roi": idx, "class": 类别, "score": 0.93, "p1": 0.22, "p2": 0.71 },
+            ...
+        ]
+    """
+
+    if frame is None or not isinstance(frame, np.ndarray):
+        raise RuntimeError("❌ classify_frame_with_rois 传入的 frame 无效")
+
+    rois = load_rois(roi_file)
+    if not rois:
+        raise RuntimeError("ROI 文件为空")
+
+    crops = crop_and_resize(frame, rois)
+
+    results = []
+
+    for roi_img, idx in crops:
+        class_id, pred = rknn_classify(roi_img, model_path)
+        class_name = CLASS_NAMES.get(class_id, f"未知类别({class_id})")
+
+        if class_id in [1, 2]:
+            final_class, score, p1, p2 = weighted_small_large(pred, threshold)
+        else:
+            final_class = class_name
+            score = float(pred[class_id])
+            p1, p2 = float(pred[1]), float(pred[2])
+
+        results.append({
+            "roi": idx,
+            "class": final_class,
+            "score": round(score, 4),
+            "p1": round(p1, 4),
+            "p2": round(p2, 4)
+        })
+
+    return results
+
+
+# =====================================================
+#                   示例调用
+# =====================================================
+if __name__ == "__main__":
+    model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "yiliao_cls.rknn")
+    roi_file = "./roi_coordinates/1_rois.txt"
+
+    frame = cv2.imread("./test_image/2.jpg")
+
+    outputs = classify_frame_with_rois(model_path, frame, roi_file)
+
+    for res in outputs:
+        print(res)
+