增加部署代码

angle_base_seg/bushu.py

@@ -1,206 +1,93 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
 import os
 import cv2
 import numpy as np
-import argparse
-import torch
-import torch.nn.functional as F
-import torchvision
+from rknnlite.api import RKNNLite
 
-# ---------------- 配置 ----------------
-OBJ_THRESH = 0.25
-NMS_THRESH = 0.45
-MAX_DETECT = 300
-IMG_SIZE = (640, 640)  # (W,H)
+# ------------------- 参数 -------------------
+objectThresh = 0.7  # 置信度阈值，可调整
 
-OUTPUT_DIR = "result"
-os.makedirs(OUTPUT_DIR, exist_ok=True)
+# ------------------- 工具函数 -------------------
+def letterbox_resize(image, size, bg_color=114):
+    target_width, target_height = size
+    h, w, _ = image.shape
+    scale = min(target_width / w, target_height / h)
+    new_w, new_h = int(w * scale), int(h * scale)
+    resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
+    canvas = np.ones((target_height, target_width, 3), dtype=np.uint8) * bg_color
+    offset_x, offset_y = (target_width - new_w) // 2, (target_height - new_h) // 2
+    canvas[offset_y:offset_y+new_h, offset_x:offset_x+new_w] = resized
+    return canvas, scale, offset_x, offset_y
 
-
-# ---------------- 工具函数 ----------------
 def sigmoid(x):
     return 1 / (1 + np.exp(-x))
 
-def dfl(position):
-    x = torch.tensor(position)
-    n, c, h, w = x.shape
-    y = x.reshape(n, 4, c // 4, h, w)
-    y = y.softmax(2)
-    acc_metrix = torch.arange(c // 4).float().reshape(1, 1, c // 4, 1, 1)
-    y = (y * acc_metrix).sum(2)
-    return y.numpy()
-
-def box_process(position):
-    grid_h, grid_w = position.shape[2:4]
-    col, row = np.meshgrid(np.arange(0, grid_w), np.arange(0, grid_h))
-    col, row = col.reshape(1, 1, grid_h, grid_w), row.reshape(1, 1, grid_h, grid_w)
-    grid = np.concatenate((col, row), axis=1)
-    stride = np.array([IMG_SIZE[1] // grid_h, IMG_SIZE[0] // grid_w]).reshape(1, 2, 1, 1)
-
-    position = dfl(position)
-    box_xy = grid + 0.5 - position[:, 0:2, :, :]
-    box_xy2 = grid + 0.5 + position[:, 2:4, :, :]
-    xyxy = np.concatenate((box_xy * stride, box_xy2 * stride), axis=1)
-    return xyxy
-
-def _crop_mask(masks, boxes):
-    n, h, w = masks.shape
-    x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1)
-    r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :]
-    c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None]
-    return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
-
-def post_process(input_data):
-    proto = input_data[-1]
-    boxes, scores, seg_part = [], [], []
-    default_branch = 3
-    pair_per_branch = len(input_data) // default_branch
-
-    for i in range(default_branch):
-        boxes.append(box_process(input_data[pair_per_branch * i]))
-        scores.append(np.ones_like(input_data[pair_per_branch * i + 1][:, :1, :, :], dtype=np.float32))
-        seg_part.append(input_data[pair_per_branch * i + 3])
-
-    def sp_flatten(_in):
-        ch = _in.shape[1]
-        _in = _in.transpose(0, 2, 3, 1)
-        return _in.reshape(-1, ch)
-
-    boxes = np.concatenate([sp_flatten(v) for v in boxes])
-    scores = np.concatenate([sp_flatten(v) for v in scores])
-    seg_part = np.concatenate([sp_flatten(v) for v in seg_part])
-
-    # 阈值过滤
-    keep = np.where(scores.reshape(-1) >= OBJ_THRESH)
-    boxes, scores, seg_part = boxes[keep], scores[keep], seg_part[keep]
-
-    # NMS
-    ids = torchvision.ops.nms(torch.tensor(boxes, dtype=torch.float32),
-                              torch.tensor(scores, dtype=torch.float32), NMS_THRESH)
-    ids = ids.tolist()[:MAX_DETECT]
-    boxes, scores, seg_part = boxes[ids], scores[ids], seg_part[ids]
-
-    # mask decode
-    ph, pw = proto.shape[-2:]
-    proto = proto.reshape(seg_part.shape[-1], -1)
-    seg_img = np.matmul(seg_part, proto)
-    seg_img = sigmoid(seg_img)
-    seg_img = seg_img.reshape(-1, ph, pw)
-    seg_img = F.interpolate(torch.tensor(seg_img)[None], torch.Size([640, 640]), mode='bilinear', align_corners=False)[0]
-    seg_img_t = _crop_mask(seg_img, torch.tensor(boxes))
-    seg_img = seg_img_t.numpy() > 0.5
-
-    return boxes, scores, seg_img
-
-
-# ---------------- 角度计算 ----------------
-def compute_angle(boxes, seg_img, h, w, filename, mode="show"):
-    composite_mask = np.zeros((h, w), dtype=np.uint8)
-    jaws = []
-
-    for i, box in enumerate(boxes):
-        x1, y1, x2, y2 = map(int, box)
-        x1, y1 = max(0, x1), max(0, y1)
-        x2, y2 = min(w, x2), min(h, y2)
-
-        obj_mask = np.zeros((h, w), dtype=np.uint8)
-        mask_resized = cv2.resize(seg_img[i].astype(np.uint8), (w, h))
-        obj_mask[y1:y2, x1:x2] = mask_resized[y1:y2, x1:x2] * 255
-
-        contours, _ = cv2.findContours(obj_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
-        if len(contours) == 0:
-            continue
-
-        largest_contour = max(contours, key=cv2.contourArea)
-        area = cv2.contourArea(largest_contour)
-        if area < 100:
-            continue
-
-        rect = cv2.minAreaRect(largest_contour)
-        jaws.append({'contour': largest_contour, 'rect': rect, 'area': area})
-        composite_mask = np.maximum(composite_mask, obj_mask)
-
-    if len(jaws) < 2:
-        print(f"❌ 检测到的夹具少于2个（共{len(jaws)}个）")
+# ------------------- 单目标分割函数 -------------------
+def detect_single_mask(model_path, image_path):
+    img = cv2.imread(image_path)
+    if img is None:
+        print(f"❌ 错误：无法读取图像 {image_path}")
         return None
 
-    jaws.sort(key=lambda x: x['area'], reverse=True)
-    jaw1, jaw2 = jaws[0], jaws[1]
+    img_resized, scale, offset_x, offset_y = letterbox_resize(img, (640, 640))
+    infer_img = img_resized[..., ::-1]  # BGR->RGB
+    infer_img = np.expand_dims(infer_img, 0)
 
-    def get_long_edge_vector(rect):
-        center, (w_, h_), angle = rect
-        rad = np.radians(angle + (0 if w_ >= h_ else 90))
-        return np.array([np.cos(rad), np.sin(rad)])
+    # ------------------- RKNN 推理 -------------------
+    rknn = RKNNLite(verbose=True)
+    print('--> Load RKNN model')
+    rknn.load_rknn(model_path)
+    print('done')
 
-    def get_center(contour):
-        M = cv2.moments(contour)
-        return np.array([M['m10']/M['m00'], M['m01']/M['m00']]) if M['m00'] != 0 else np.array([0, 0])
+    print('--> Init runtime environment')
+    rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
+    print('done')
 
-    dir1, dir2 = get_long_edge_vector(jaw1['rect']), get_long_edge_vector(jaw2['rect'])
-    center1, center2 = get_center(jaw1['contour']), get_center(jaw2['contour'])
-    fixture_center = (center1 + center2) / 2
+    print('--> Running model')
+    outputs = rknn.inference([infer_img])
+    rknn.release()
 
-    to_center1, to_center2 = fixture_center - center1, fixture_center - center2
-    if np.linalg.norm(to_center1) > 1e-6 and np.dot(dir1, to_center1) < 0:
-        dir1 = -dir1
-    if np.linalg.norm(to_center2) > 1e-6 and np.dot(dir2, to_center2) < 0:
-        dir2 = -dir2
+    # ------------------- 处理输出 -------------------
+    # outputs顺序参考你提供：
+    # [output1, output2, output3, output4, ... , output12, output13]
+    # 我们只取置信度最高点的 mask
 
-    cos_angle = np.clip(np.dot(dir1, dir2), -1.0, 1.0)
-    angle = np.degrees(np.arccos(cos_angle))
-    opening_angle = min(angle, 180 - angle)
+    mask_coeffs_list = [outputs[3], outputs[7], outputs[11]]  # mask coefficients
+    conf_list = [outputs[1], outputs[5], outputs[9]]            # object confidence
+    proto = outputs[12][0]                                      # proto [32,160,160]
 
-    if mode == "show":
-        vis_img = np.stack([composite_mask]*3, axis=-1)
-        vis_img[composite_mask > 0] = [255, 255, 255]
+    # 找到所有尺度的最大置信度位置
+    best_idx = None
+    best_conf = -1
+    best_scale = None
+    for i, conf_map in enumerate(conf_list):
+        conf_map_flat = conf_map.flatten()
+        idx = np.argmax(conf_map_flat)
+        if conf_map_flat[idx] > best_conf:
+            best_conf = conf_map_flat[idx]
+            best_idx = idx
+            best_scale = i
 
-        box1, box2 = np.int32(cv2.boxPoints(jaw1['rect'])), np.int32(cv2.boxPoints(jaw2['rect']))
-        cv2.drawContours(vis_img, [box1], 0, (0, 0, 255), 2)
-        cv2.drawContours(vis_img, [box2], 0, (255, 0, 0), 2)
+    if best_conf < objectThresh:
+        print(f"⚠️ 置信度低于阈值 {objectThresh}，未检测到目标")
+        return None
 
-        scale = 60
-        c1, c2 = tuple(np.int32(center1)), tuple(np.int32(center2))
-        end1, end2 = tuple(np.int32(center1 + scale * dir1)), tuple(np.int32(center2 + scale * dir2))
-        cv2.arrowedLine(vis_img, c1, end1, (0, 255, 0), 2, tipLength=0.3)
-        cv2.arrowedLine(vis_img, c2, end2, (0, 255, 0), 2, tipLength=0.3)
+    # ------------------- 构建 mask -------------------
+    coeff = mask_coeffs_list[best_scale].reshape(mask_coeffs_list[best_scale].shape[1], -1)
+    mask_flat = np.matmul(coeff[:, best_idx], proto.reshape(proto.shape[0], -1))
+    mask = sigmoid(mask_flat).reshape(proto.shape[1], proto.shape[2])
 
-        cv2.putText(vis_img, f"Angle: {opening_angle:.2f}°", (20, 50),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
+    # resize 回原图
+    mask_resized = cv2.resize(mask, (img.shape[1], img.shape[0]))
+    mask_bin = (mask_resized > 0.5).astype(np.uint8) * 255
 
-        save_path = os.path.join(OUTPUT_DIR, f'angle_{filename}')
-        cv2.imwrite(save_path, vis_img)
-        print(f"✅ 结果已保存: {save_path}")
-
-    return round(opening_angle, 2)
-
-
-# ---------------- 主程序 ----------------
-def main():
-    # 固定路径（写死）
-    MODEL_PATH = "/userdata/bushu/seg.rknn"
-    IMG_PATH   = "/userdata/bushu/test.jpg"
-
-    from py_utils.rknn_executor import RKNN_model_container
-    model = RKNN_model_container(MODEL_PATH, target='rk3588', device_id=None)
-
-    img_src = cv2.imread(IMG_PATH)
-    if img_src is None:
-        print("❌ 图片路径错误:", IMG_PATH)
-        return
-    h, w = img_src.shape[:2]
-    img = cv2.resize(img_src, IMG_SIZE)
-
-    outputs = model.run([img])
-    boxes, scores, seg_img = post_process(outputs)
-
-    filename = os.path.basename(IMG_PATH)
-    angle = compute_angle(boxes, seg_img, h, w, filename, mode="show")
-    if angle is not None:
-        print(f"🎉 检测到的夹具开合角度: {angle}°")
-
-    model.release()
+    # 保存或显示
+    cv2.imwrite("mask_result.png", mask_bin)
+    print("✅ 单目标 mask 已保存: mask_result.png")
 
+    return mask_bin
 
+# ------------------- 调用示例 -------------------
 if __name__ == "__main__":
-    main()
+    model_path = "seg.rknn"
+    image_path = "2.jpg"
+    detect_single_mask(model_path, image_path)

angle_base_seg/test_seg_angle_f.py

@@ -6,7 +6,7 @@ import os
 # ------------------ 配置 ------------------
 model_path = '../ultralytics_yolo11-main/runs/train/exp4/weights/best.pt'
 #img_folder = '/home/hx/yolo/ultralytics_yolo11-main/dataset1/test'
-img_folder = '/home/hx/yolo/test_image'
+img_folder = '/home/hx/yolo/output_masks'
 output_mask_dir = 'output_masks1'
 os.makedirs(output_mask_dir, exist_ok=True)
 

angle_base_seg/test_seg_angle_f1.py

@@ -4,8 +4,8 @@ import numpy as np
 import os
 
 # ------------------ 配置 ------------------
-model_path = '/home/hx/yolo/ultralytics_yolo11-main/runs/train/seg/exp3/weights/best.pt'
-img_folder = '/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/dataset2/test'
+model_path = '/home/hx/yolo/ultralytics_yolo11-main/runs/train/seg_j/exp/weights/best.pt'
+img_folder = '/home/hx/yolo/output_masks'
 output_mask_dir = 'output_masks1'
 os.makedirs(output_mask_dir, exist_ok=True)
 
@@ -13,7 +13,7 @@ SUPPORTED_FORMATS = ('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff')
 
 # ------------------ 加载模型 ------------------
 model = YOLO(model_path)
-model.to('cuda')  # 使用 GPU 加速（如有）
+model.to('cuda')  # 使用 GPU 加速（如有））
 
 def get_orientation_vector(rect):
     """从 minAreaRect 中提取主方向向量（单位向量）"""