import cv2 import numpy as np import math from shapely.geometry import Polygon from rknnlite.api import RKNNLite import os CLASSES = ['clamp'] nmsThresh = 0.4 objectThresh = 0.5 # ------------------- 工具函数 ------------------- def letterbox_resize(image, size, bg_color=114): target_width, target_height = size image_height, image_width, _ = image.shape scale = min(target_width / image_width, target_height / image_height) new_width, new_height = int(image_width * scale), int(image_height * scale) image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA) canvas = np.ones((target_height, target_width, 3), dtype=np.uint8) * bg_color offset_x, offset_y = (target_width - new_width) // 2, (target_height - new_height) // 2 canvas[offset_y:offset_y + new_height, offset_x:offset_x + new_width] = image return canvas, scale, offset_x, offset_y class DetectBox: def __init__(self, classId, score, xmin, ymin, xmax, ymax, angle): self.classId = classId self.score = score self.xmin = xmin self.ymin = ymin self.xmax = xmax self.ymax = ymax self.angle = angle def rotate_rectangle(x1, y1, x2, y2, a): cx, cy = (x1 + x2) / 2, (y1 + y2) / 2 x1_new = int((x1 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx) y1_new = int((x1 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy) x2_new = int((x2 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx) y2_new = int((x2 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy) x3_new = int((x1 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx) y3_new = int((x1 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy) x4_new = int((x2 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx) y4_new = int((x2 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy) return [(x1_new, y1_new), (x3_new, y3_new), (x2_new, y2_new), (x4_new, y4_new)] def intersection(g, p): g = Polygon(np.array(g).reshape(-1,2)) p = Polygon(np.array(p).reshape(-1,2)) if not g.is_valid or not p.is_valid: return 0 inter = g.intersection(p).area union = g.area + p.area - inter return 0 if union == 0 else inter / union def NMS(detectResult): predBoxs = [] sort_detectboxs = sorted(detectResult, key=lambda x: x.score, reverse=True) for i in range(len(sort_detectboxs)): if sort_detectboxs[i].classId == -1: continue p1 = rotate_rectangle(sort_detectboxs[i].xmin, sort_detectboxs[i].ymin, sort_detectboxs[i].xmax, sort_detectboxs[i].ymax, sort_detectboxs[i].angle) predBoxs.append(sort_detectboxs[i]) for j in range(i + 1, len(sort_detectboxs)): if sort_detectboxs[j].classId == sort_detectboxs[i].classId: p2 = rotate_rectangle(sort_detectboxs[j].xmin, sort_detectboxs[j].ymin, sort_detectboxs[j].xmax, sort_detectboxs[j].ymax, sort_detectboxs[j].angle) if intersection(p1, p2) > nmsThresh: sort_detectboxs[j].classId = -1 return predBoxs def sigmoid(x): return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x))) def softmax(x, axis=-1): exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True)) return exp_x / np.sum(exp_x, axis=axis, keepdims=True) def process(out, model_w, model_h, stride, angle_feature, index, scale_w=1, scale_h=1): class_num = len(CLASSES) angle_feature = angle_feature.reshape(-1) xywh = out[:, :64, :] conf = sigmoid(out[:, 64:, :]) conf = conf.reshape(-1) boxes = [] for ik in range(model_h * model_w * class_num): if conf[ik] > objectThresh: w = ik % model_w h = (ik % (model_w * model_h)) // model_w c = ik // (model_w * model_h) xywh_ = xywh[0, :, (h * model_w) + w].reshape(1, 4, 16, 1) data = np.arange(16).reshape(1, 1, 16, 1) xywh_ = softmax(xywh_, 2) xywh_ = np.sum(xywh_ * data, axis=2).reshape(-1) xywh_add = xywh_[:2] + xywh_[2:] xywh_sub = (xywh_[2:] - xywh_[:2]) / 2 angle = (angle_feature[index + (h * model_w) + w] - 0.25) * math.pi cos_a, sin_a = math.cos(angle), math.sin(angle) xy = xywh_sub[0] * cos_a - xywh_sub[1] * sin_a, xywh_sub[0] * sin_a + xywh_sub[1] * cos_a xywh1 = np.array([xy[0] + w + 0.5, xy[1] + h + 0.5, xywh_add[0], xywh_add[1]]) xywh1 *= stride xmin = (xywh1[0] - xywh1[2] / 2) * scale_w ymin = (xywh1[1] - xywh1[3] / 2) * scale_h xmax = (xywh1[0] + xywh1[2] / 2) * scale_w ymax = (xywh1[1] + xywh1[3] / 2) * scale_h boxes.append(DetectBox(c, conf[ik], xmin, ymin, xmax, ymax, angle)) return boxes # ------------------- 主函数 ------------------- def detect_two_box_angle(model_path, image_path, save_path="result.jpg"): img = cv2.imread(image_path) if img is None: print(f"❌ 无法读取图像: {image_path}") return None, None img_resized, scale, offset_x, offset_y = letterbox_resize(img, (640, 640)) infer_img = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB) infer_img = np.expand_dims(infer_img, 0) rknn_lite = RKNNLite(verbose=False) rknn_lite.load_rknn(model_path) rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0) results = rknn_lite.inference([infer_img]) outputs = [] for x in results[:-1]: index, stride = 0, 0 if x.shape[2] == 20: stride, index = 32, 20*4*20*4 + 20*2*20*2 elif x.shape[2] == 40: stride, index = 16, 20*4*20*4 elif x.shape[2] == 80: stride, index = 8, 0 feature = x.reshape(1, 65, -1) outputs += process(feature, x.shape[3], x.shape[2], stride, results[-1], index) predbox = NMS(outputs) if len(predbox) < 2: print("⚠️ 检测少于两个目标，无法计算夹角。") return None, img predbox = sorted(predbox, key=lambda x: x.score, reverse=True) box1, box2 = predbox[:2] # 可视化 for box in [box1, box2]: xmin = int((box.xmin - offset_x) / scale) ymin = int((box.ymin - offset_y) / scale) xmax = int((box.xmax - offset_x) / scale) ymax = int((box.ymax - offset_y) / scale) points = rotate_rectangle(xmin, ymin, xmax, ymax, box.angle) cv2.polylines(img, [np.array(points, np.int32)], True, (0, 255, 0), 1) def main_direction(box): w, h = (box.xmax - box.xmin)/scale, (box.ymax - box.ymin)/scale direction = box.angle if w >= h else box.angle + np.pi/2 return direction % np.pi dir1 = main_direction(box1) dir2 = main_direction(box2) diff = abs(dir1 - dir2) diff = min(diff, np.pi - diff) angle_deg = np.degrees(diff) os.makedirs(os.path.dirname(save_path), exist_ok=True) cv2.imwrite(save_path, img) rknn_lite.release() return angle_deg, img # ------------------- 使用示例 ------------------- if __name__ == "__main__": angle_deg, image = detect_two_box_angle("obb.rknn", "2.jpg", "./inference_results/detected_2.jpg") if angle_deg is not None: print(f"两个置信度最高框的夹角（度）：{angle_deg:.2f}")