235 lines
9.1 KiB
Python
235 lines
9.1 KiB
Python
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.35
|
||
|
||
# 可视化与保存控制（全局变量，可外部修改）
|
||
DRAW_RESULT = True # 是否在输出图像上绘制旋转框
|
||
SAVE_PATH = None # 保存路径，如 "./result.jpg"；设为 None 则不保存
|
||
|
||
# RKNN 单例
|
||
_rknn_instance = None
|
||
|
||
# ------------------- RKNN 管理函数 -------------------
|
||
def init_rknn(model_path):
|
||
"""只加载一次 RKNN 模型"""
|
||
global _rknn_instance
|
||
if _rknn_instance is None:
|
||
_rknn_instance = RKNNLite(verbose=False)
|
||
ret = _rknn_instance.load_rknn(model_path)
|
||
if ret != 0:
|
||
print(f"[ERROR] Failed to load RKNN model: {ret}")
|
||
return None
|
||
ret = _rknn_instance.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
|
||
if ret != 0:
|
||
print(f"[ERROR] Failed to init runtime: {ret}")
|
||
return None
|
||
return _rknn_instance
|
||
|
||
def release_rknn():
|
||
"""释放 RKNN 对象"""
|
||
global _rknn_instance
|
||
if _rknn_instance:
|
||
_rknn_instance.release()
|
||
_rknn_instance = None
|
||
|
||
# ------------------- 工具函数 -------------------
|
||
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_resized = 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_resized
|
||
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
|
||
cos_a, sin_a = math.cos(a), math.sin(a)
|
||
pts = [(x1, y1), (x1, y2), (x2, y2), (x2, y1)]
|
||
return [[int(cx + (xx - cx) * cos_a - (yy - cy) * sin_a),
|
||
int(cy + (xx - cx) * sin_a + (yy - cy) * cos_a)] for xx, yy in pts]
|
||
|
||
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):
|
||
x = np.clip(x, -709, 709) # 防止 exp 溢出
|
||
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, rgb_frame):
|
||
"""
|
||
输入模型路径和 RGB 图像（numpy array），输出夹角和结果图像。
|
||
可视化和保存由全局变量 DRAW_RESULT 和 SAVE_PATH 控制。
|
||
"""
|
||
global _rknn_instance, DRAW_RESULT, SAVE_PATH
|
||
|
||
if not isinstance(rgb_frame, np.ndarray) or rgb_frame is None:
|
||
print(f"[ERROR] detect_two_box_angle 接收到错误类型: {type(rgb_frame)}")
|
||
return None, np.zeros((640, 640, 3), np.uint8)
|
||
|
||
# 注意：输入是 BGR（因为 cv2.imread 返回 BGR），但内部会转为 RGB 给模型
|
||
img = rgb_frame.copy()
|
||
img_resized, scale, offset_x, offset_y = letterbox_resize(img, (640, 640))
|
||
infer_img = np.expand_dims(cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB), 0)
|
||
|
||
try:
|
||
rknn = init_rknn(model_path)
|
||
if rknn is None:
|
||
return None, img
|
||
results = rknn.inference([infer_img])
|
||
except Exception as e:
|
||
print(f"[ERROR] RKNN 推理失败: {e}")
|
||
return None, 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)
|
||
print(f"[DEBUG] 检测到 {len(predbox)} 个框")
|
||
|
||
if len(predbox) < 2:
|
||
print("检测少于两个目标，无法计算夹角。")
|
||
return None, img
|
||
|
||
predbox = sorted(predbox, key=lambda x: x.score, reverse=True)
|
||
box1, box2 = predbox[:2]
|
||
|
||
output_img = img.copy() if DRAW_RESULT else img # 若不绘制，则直接用原图
|
||
|
||
if DRAW_RESULT:
|
||
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(output_img, [np.array(points, np.int32)], True, (0, 255, 0), 2)
|
||
|
||
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)
|
||
|
||
# 保存结果（如果需要）
|
||
if SAVE_PATH:
|
||
save_dir = os.path.dirname(SAVE_PATH)
|
||
if save_dir: # 非空目录才创建
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
cv2.imwrite(SAVE_PATH, output_img)
|
||
|
||
return angle_deg, output_img
|
||
|
||
# ------------------- 示例调用 -------------------
|
||
if __name__ == "__main__":
|
||
MODEL_PATH = "./obb.rknn"
|
||
IMAGE_PATH = "./11.jpg"
|
||
|
||
# === 全局控制开关 ===
|
||
DRAW_RESULT = True # 是否绘制框
|
||
SAVE_PATH = "./result11.jpg" # 保存路径，设为 None 则不保存
|
||
|
||
frame = cv2.imread(IMAGE_PATH)
|
||
if frame is None:
|
||
print(f"[ERROR] 无法读取图像: {IMAGE_PATH}")
|
||
else:
|
||
angle_deg, output_image = detect_two_box_angle(MODEL_PATH, frame)
|
||
if angle_deg is not None:
|
||
print(f"检测到的角度差: {angle_deg:.2f}°")
|
||
else:
|
||
print("未能成功检测到目标或计算角度差") |