172 lines
7.2 KiB
Python
172 lines
7.2 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.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:, :]).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_ = 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_idx = min(index + (h * model_w) + w, len(angle_feature) - 1)
|
|
angle = (angle_feature[angle_idx] - 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_boxes_angle_rknn(model_path, image_path, save_path=None):
|
|
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])
|
|
detect_boxes = []
|
|
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)
|
|
detect_boxes += process(feature, x.shape[3], x.shape[2], stride, results[-1], index)
|
|
|
|
detect_boxes = NMS(detect_boxes)
|
|
|
|
# 输出每个检测框角度
|
|
for i, box in enumerate(detect_boxes):
|
|
print(f"框 {i+1}: angle = {box.angle:.4f} rad ({np.degrees(box.angle):.2f}°)")
|
|
if save_path:
|
|
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)
|
|
cv2.putText(img, f"{np.degrees(box.angle):.1f}°", (xmin, ymin-5),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 1)
|
|
|
|
if save_path:
|
|
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
|
cv2.imwrite(save_path, img)
|
|
print(f"✅ 带角度的检测结果已保存到 {save_path}")
|
|
|
|
rknn_lite.release()
|
|
return detect_boxes, img
|
|
|
|
# ------------------- 使用示例 -------------------
|
|
if __name__ == "__main__":
|
|
model_path = "obb.rknn"
|
|
image_path = "2.jpg"
|
|
save_path = "./inference_results/boxes_with_angle.jpg"
|
|
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
|
detect_boxes_angle_rknn(model_path, image_path, save_path)
|
|
|