215 lines
7.1 KiB
Python
215 lines
7.1 KiB
Python
import cv2
|
|
import numpy as np
|
|
from rknnlite.api import RKNNLite
|
|
|
|
# ---------------------------
|
|
# 配置
|
|
# ---------------------------
|
|
ROIS = [(604, 182, 594, 252)] # (x, y, w, h)
|
|
|
|
IMG_SIZE = 640
|
|
STRIDES = [8, 16, 32]
|
|
OBJ_THRESH = 0.25
|
|
MASK_THRESH = 0.5
|
|
|
|
_global_rknn = None
|
|
|
|
|
|
def init_rknn_model(model_path):
|
|
global _global_rknn
|
|
if _global_rknn is not None:
|
|
return
|
|
rknn = RKNNLite(verbose=False)
|
|
ret = rknn.load_rknn(model_path)
|
|
if ret != 0:
|
|
raise RuntimeError(f"Load RKNN failed: {ret}")
|
|
ret = rknn.init_runtime()
|
|
if ret != 0:
|
|
raise RuntimeError(f"Init runtime failed: {ret}")
|
|
_global_rknn = rknn
|
|
print(f"[INFO] RKNN model loaded: {model_path}")
|
|
|
|
|
|
def sigmoid(x):
|
|
return 1 / (1 + np.exp(-x))
|
|
|
|
|
|
def dfl_decode(dfl):
|
|
bins = np.arange(16)
|
|
dfl = sigmoid(dfl)
|
|
dfl /= np.sum(dfl, axis=1, keepdims=True)
|
|
return np.sum(dfl * bins, axis=1)
|
|
|
|
|
|
def largest_intersect_cc(mask_bin, bbox):
|
|
x1, y1, x2, y2 = bbox
|
|
contours, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
if not contours:
|
|
return np.zeros_like(mask_bin, dtype=np.uint8)
|
|
max_inter = 0
|
|
best = np.zeros_like(mask_bin, dtype=np.uint8)
|
|
for cnt in contours:
|
|
tmp = np.zeros_like(mask_bin, dtype=np.uint8)
|
|
cv2.drawContours(tmp, [cnt], -1, 1, -1)
|
|
cx, cy, cw, ch = cv2.boundingRect(cnt)
|
|
ix1 = max(cx, x1)
|
|
iy1 = max(cy, y1)
|
|
ix2 = min(cx + cw, x2)
|
|
iy2 = min(cy + ch, y2)
|
|
area = max(0, ix2 - ix1) * max(0, iy2 - iy1)
|
|
if area > max_inter:
|
|
max_inter = area
|
|
best = tmp
|
|
return best
|
|
|
|
|
|
def seg_infer(roi):
|
|
rknn = _global_rknn
|
|
h0, w0 = roi.shape[:2]
|
|
inp_img = cv2.resize(roi, (IMG_SIZE, IMG_SIZE))
|
|
inp = inp_img[..., ::-1][None, ...]
|
|
outputs = rknn.inference([inp])
|
|
proto = outputs[12][0]
|
|
proto_h, proto_w = proto.shape[1:]
|
|
best_score = -1
|
|
best_coef = None
|
|
best_bbox = None
|
|
out_i = 0
|
|
for stride in STRIDES:
|
|
reg = outputs[out_i][0]
|
|
cls = outputs[out_i + 1][0, 0]
|
|
obj = outputs[out_i + 2][0, 0]
|
|
coef = outputs[out_i + 3][0]
|
|
out_i += 4
|
|
score_map = sigmoid(cls) * sigmoid(obj)
|
|
y, x = np.unravel_index(np.argmax(score_map), score_map.shape)
|
|
score = score_map[y, x]
|
|
if score > best_score and score > OBJ_THRESH:
|
|
best_score = score
|
|
best_coef = coef[:, y, x]
|
|
dfl = reg[:, y, x].reshape(4, 16)
|
|
l, t, r, b = dfl_decode(dfl)
|
|
cx = (x + 0.5) * stride
|
|
cy = (y + 0.5) * stride
|
|
scale = proto_w / IMG_SIZE
|
|
x1 = int((cx - l) * scale)
|
|
y1 = int((cy - t) * scale)
|
|
x2 = int((cx + r) * scale)
|
|
y2 = int((cy + b) * scale)
|
|
best_bbox = (max(0, x1), max(0, y1), min(proto_w, x2), min(proto_h, y2))
|
|
if best_coef is None:
|
|
return np.zeros((h0, w0), dtype=np.uint8)
|
|
proto_mask = sigmoid(np.tensordot(best_coef, proto, axes=1)) > MASK_THRESH
|
|
proto_mask = proto_mask.astype(np.uint8)
|
|
mask_final = largest_intersect_cc(proto_mask, best_bbox)
|
|
mask_roi = cv2.resize(mask_final, (w0, h0), interpolation=cv2.INTER_NEAREST) * 255
|
|
return mask_roi.astype(np.uint8)
|
|
|
|
|
|
# ---------------------------
|
|
# 主函数：支持可选可视化
|
|
# ---------------------------
|
|
def caculate_yemian_diff(img, visualize=False):
|
|
"""
|
|
输入:
|
|
img: BGR 图像 (H, W, 3) np.ndarray
|
|
visualize: bool, 是否生成可视化结果
|
|
|
|
输出:
|
|
若 visualize=False:
|
|
(diff14: float, diff43: float, mask_area: int)
|
|
若 visualize=True:
|
|
(diff14: float, diff43: float, mask_area: int, vis_img: np.ndarray)
|
|
|
|
失败时返回 (0.0, 0.0, 0) 或 (0.0, 0.0, 0, original_img)
|
|
"""
|
|
if _global_rknn is None:
|
|
raise RuntimeError("RKNN model not initialized. Call init_rknn_model() first.")
|
|
|
|
vis_img = img.copy() if visualize else None
|
|
|
|
for (rx, ry, rw, rh) in ROIS:
|
|
roi = img[ry:ry + rh, rx:rx + rw]
|
|
mask_full = seg_infer(roi)
|
|
mask_bin = mask_full // 255
|
|
mask_area = int(np.sum(mask_bin))
|
|
|
|
if visualize:
|
|
green = np.zeros_like(roi)
|
|
green[mask_bin == 1] = (0, 255, 0)
|
|
vis_img[ry:ry + rh, rx:rx + rw] = cv2.addWeighted(roi, 0.7, green, 0.3, 0)
|
|
|
|
contours, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
|
if not contours:
|
|
continue
|
|
cnt = max(contours, key=cv2.contourArea)
|
|
if len(cnt) < 20:
|
|
continue
|
|
|
|
pts_all = cnt.reshape(-1, 2)
|
|
P1 = pts_all[np.argmin(pts_all[:, 0])] # x_min
|
|
P4 = pts_all[np.argmin(pts_all[:, 1])] # y_min
|
|
P3 = pts_all[np.argmax(pts_all[:, 0])] # x_max
|
|
|
|
# 全局坐标
|
|
P1_g = (int(P1[0] + rx), int(P1[1] + ry))
|
|
P4_g = (int(P4[0] + rx), int(P4[1] + ry))
|
|
P3_g = (int(P3[0] + rx), int(P3[1] + ry))
|
|
|
|
diff14 = float(np.linalg.norm(np.array(P1_g) - np.array(P4_g)))
|
|
diff43 = float(np.linalg.norm(np.array(P4_g) - np.array(P3_g)))
|
|
|
|
if visualize:
|
|
roi_vis = vis_img[ry:ry + rh, rx:rx + rw]
|
|
local_pts = [P1.astype(int), P4.astype(int), P3.astype(int)]
|
|
colors = [(255, 0, 0), (0, 255, 0)]
|
|
lengths = [diff14, diff43]
|
|
# 画线
|
|
cv2.line(vis_img, P1_g, P4_g, colors[0], 2)
|
|
cv2.line(vis_img, P4_g, P3_g, colors[1], 2)
|
|
# 标长度
|
|
mid14 = ((P1_g[0] + P4_g[0]) // 2, (P1_g[1] + P4_g[1]) // 2 - 10)
|
|
mid43 = ((P4_g[0] + P3_g[0]) // 2, (P4_g[1] + P3_g[1]) // 2 - 10)
|
|
cv2.putText(vis_img, f"{diff14:.1f}", mid14, cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors[0], 1)
|
|
cv2.putText(vis_img, f"{diff43:.1f}", mid43, cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors[1], 1)
|
|
# 标点名
|
|
labels = ["P1", "P4", "P3"]
|
|
points = [P1_g, P4_g, P3_g]
|
|
offsets = [(-25, -10), (-25, -10), (10, -10)]
|
|
for lab, pt, off in zip(labels, points, offsets):
|
|
x = max(10, min(pt[0] + off[0], img.shape[1] - 50))
|
|
y = max(20, min(pt[1] + off[1], img.shape[0] - 10))
|
|
cv2.putText(vis_img, lab, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
|
|
# 面积
|
|
cv2.putText(vis_img, f"Area: {mask_area}", (rx + 10, ry + 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
|
|
|
|
if visualize:
|
|
return diff14, diff43, mask_area, vis_img
|
|
else:
|
|
return diff14, diff43, mask_area
|
|
|
|
# 检测失败
|
|
if visualize:
|
|
return 0.0, 0.0, 0, img.copy()
|
|
else:
|
|
return 0.0, 0.0, 0
|
|
|
|
|
|
# ---------------------------
|
|
# 示例用法
|
|
# ---------------------------
|
|
if __name__ == "__main__":
|
|
init_rknn_model("60seg.rknn")
|
|
img = cv2.imread("1.png")
|
|
if img is None:
|
|
raise FileNotFoundError("1.png")
|
|
|
|
# 不可视化
|
|
d14, d43, area = caculate_yemian_diff(img)
|
|
print(f"Without vis: {d14:.2f}, {d43:.2f}, {area}")
|
|
|
|
# 可视化
|
|
d14, d43, area, vis = caculate_yemian_diff(img, visualize=True)
|
|
cv2.imwrite("output_vis.png", vis)
|
|
print(f"With vis: saved to output_vis.png") |