215 lines
6.6 KiB
Python
215 lines
6.6 KiB
Python
import cv2
|
||
import numpy as np
|
||
import math
|
||
|
||
# ---------- 配置 ----------
|
||
CLASSES = ['clamp']
|
||
nmsThresh = 0.4
|
||
objectThresh = 0.5
|
||
INPUT_DTYPE = np.uint8
|
||
|
||
DRAW_BOX = True
|
||
DRAW_SCORE = False
|
||
|
||
# ---------------- 工具函数 ----------------
|
||
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 letterbox_resize(image, size, bg_color=114):
|
||
tw, th = size
|
||
h, w = image.shape[:2]
|
||
scale = min(tw / w, th / h)
|
||
nw, nh = int(w * scale), int(h * scale)
|
||
img_resized = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_AREA)
|
||
canvas = np.full((th, tw, 3), bg_color, dtype=np.uint8)
|
||
dx, dy = (tw - nw) // 2, (th - nh) // 2
|
||
canvas[dy:dy + nh, dx:dx + nw] = img_resized
|
||
return canvas, scale, dx, dy
|
||
|
||
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 polygon_area(pts):
|
||
x, y = zip(*pts)
|
||
return 0.5 * abs(sum(x[i] * y[(i+1)%4] - x[(i+1)%4]*y[i] for i in range(4)))
|
||
|
||
def polygon_intersection_area(p1, p2):
|
||
try:
|
||
from shapely.geometry import Polygon
|
||
poly1, poly2 = Polygon(p1), Polygon(p2)
|
||
if not poly1.is_valid: poly1 = poly1.buffer(0)
|
||
if not poly2.is_valid: poly2 = poly2.buffer(0)
|
||
inter = poly1.intersection(poly2)
|
||
return inter.area if not inter.is_empty else 0
|
||
except:
|
||
return 0
|
||
|
||
def IoU(b1, b2):
|
||
p1 = rotate_rectangle(b1.xmin, b1.ymin, b1.xmax, b1.ymax, b1.angle)
|
||
p2 = rotate_rectangle(b2.xmin, b2.ymin, b2.xmax, b2.ymax, b2.angle)
|
||
inter = polygon_intersection_area(p1, p2)
|
||
area1 = polygon_area(p1)
|
||
area2 = polygon_area(p2)
|
||
union = area1 + area2 - inter
|
||
return inter / union if union > 1e-6 else 0
|
||
|
||
def NMS(boxes):
|
||
if not boxes: return []
|
||
boxes = sorted(boxes, key=lambda x: x.score, reverse=True)
|
||
keep = []
|
||
for i, b1 in enumerate(boxes):
|
||
if b1.classId == -1: continue
|
||
keep.append(b1)
|
||
for j in range(i+1, len(boxes)):
|
||
b2 = boxes[j]
|
||
if b2.classId == b1.classId and IoU(b1, b2) > nmsThresh:
|
||
b2.classId = -1
|
||
return keep
|
||
|
||
from scipy.special import expit
|
||
def sigmoid(x): return expit(x)
|
||
|
||
def softmax(x, axis=-1):
|
||
x = np.asarray(x)
|
||
x_max = np.max(x, axis=axis, keepdims=True)
|
||
e = np.exp(x - x_max)
|
||
return e / (e.sum(axis=axis, keepdims=True) + 1e-9)
|
||
|
||
ARANGE16 = np.arange(16).reshape(1,1,16,1)
|
||
|
||
def process(out, mw, mh, stride, angle_feature, index, scale_w=1, scale_h=1):
|
||
angle_feature = angle_feature.reshape(-1)
|
||
xywh = out[:, :64, :]
|
||
conf = sigmoid(out[:, 64:, :]).reshape(-1)
|
||
boxes = []
|
||
class_num = len(CLASSES)
|
||
|
||
total = mh * mw * class_num
|
||
for ik in range(total):
|
||
if conf[ik] <= objectThresh: continue
|
||
w = ik % mw
|
||
h = (ik // mw) % mh
|
||
c = ik // (mw * mh)
|
||
|
||
xywh_ = xywh[0, :, h * mw + w].reshape(1,4,16,1)
|
||
xywh_ = softmax(xywh_, axis=2)
|
||
xywh_ = np.sum(xywh_ * ARANGE16, axis=2).reshape(-1)
|
||
|
||
xy_add = xywh_[:2] + xywh_[2:]
|
||
xy_sub = (xywh_[2:] - xywh_[:2]) / 2
|
||
|
||
angle = (angle_feature[index + h*mw + w] - 0.25) * math.pi
|
||
|
||
cos_a, sin_a = math.cos(angle), math.sin(angle)
|
||
xy_rot = np.array([
|
||
xy_sub[0]*cos_a - xy_sub[1]*sin_a,
|
||
xy_sub[0]*sin_a + xy_sub[1]*cos_a
|
||
])
|
||
|
||
cx = (xy_rot[0] + w + 0.5) * stride
|
||
cy = (xy_rot[1] + h + 0.5) * stride
|
||
w_box = xy_add[0] * stride
|
||
h_box = xy_add[1] * stride
|
||
|
||
xmin = (cx - w_box/2) * scale_w
|
||
ymin = (cy - h_box/2) * scale_h
|
||
xmax = (cx + w_box/2) * scale_w
|
||
ymax = (cy + h_box/2) * scale_h
|
||
|
||
boxes.append(DetectBox(c, float(conf[ik]), xmin, ymin, xmax, ymax, float(angle)))
|
||
|
||
return boxes
|
||
|
||
# ---------------- RKNN 推理接口 ----------------
|
||
def myFunc(rknn, frame, lock=None):
|
||
try:
|
||
# --- 推理前 letterbox ---
|
||
img_resized, scale, offset_x, offset_y = letterbox_resize(frame, (640,640))
|
||
infer_img = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)
|
||
infer_input = np.expand_dims(infer_img.astype(INPUT_DTYPE), 0)
|
||
|
||
results = rknn.inference([infer_input])
|
||
if not results or len(results) < 1:
|
||
return frame
|
||
|
||
outputs = []
|
||
for x in results[:-1]:
|
||
if x is None: continue
|
||
# 决定 stride 与索引
|
||
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
|
||
else:
|
||
continue
|
||
|
||
feature = x.reshape(1,65,-1)
|
||
outputs += process(
|
||
feature,
|
||
x.shape[3],
|
||
x.shape[2],
|
||
stride,
|
||
results[-1],
|
||
index,
|
||
1.0, 1.0 # 输出坐标保持在640×640下
|
||
)
|
||
|
||
if not outputs:
|
||
return frame
|
||
|
||
predbox = NMS(outputs)
|
||
if len(predbox) < 2:
|
||
return frame
|
||
|
||
box1, box2 = sorted(predbox, key=lambda x: x.score, reverse=True)[:2]
|
||
|
||
out_frame = frame.copy()
|
||
|
||
# ========== 还原到原图坐标 ==========
|
||
def restore_to_original(b):
|
||
xmin = int((b.xmin - offset_x) / scale)
|
||
ymin = int((b.ymin - offset_y) / scale)
|
||
xmax = int((b.xmax - offset_x) / scale)
|
||
ymax = int((b.ymax - offset_y) / scale)
|
||
return xmin, ymin, xmax, ymax
|
||
|
||
for box in [box1, box2]:
|
||
xmin, ymin, xmax, ymax = restore_to_original(box)
|
||
|
||
if DRAW_BOX:
|
||
# 旋转框顶点（注意：旋转必须在原图坐标系）
|
||
pts = rotate_rectangle(xmin, ymin, xmax, ymax, box.angle)
|
||
cv2.polylines(out_frame, [np.array(pts, np.int32)], True, (0,255,0), 2)
|
||
|
||
if DRAW_SCORE:
|
||
cv2.putText(
|
||
out_frame, f"{box.score:.2f}",
|
||
(xmin, max(10, ymin - 6)),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2
|
||
)
|
||
|
||
return out_frame
|
||
|
||
except Exception as e:
|
||
print(f"[func ❌] 推理异常: {e}")
|
||
return frame
|
||
|