diff --git a/yemian_seg_diff_60/60seg.rknn b/yemian_seg_diff_60/60seg.rknn new file mode 100644 index 0000000..336ae0d Binary files /dev/null and b/yemian_seg_diff_60/60seg.rknn differ diff --git a/yemian_seg_diff_60/main.py b/yemian_seg_diff_60/main.py new file mode 100644 index 0000000..299356e --- /dev/null +++ b/yemian_seg_diff_60/main.py @@ -0,0 +1,291 @@ +import os +import cv2 +import numpy as np +from rknnlite.api import RKNNLite + +# --------------------------- +# 配置 +# --------------------------- +ROIS = [ + (604,182,594,252), +] + +IMG_SIZE = 640 +STRIDES = [8, 16, 32] +OBJ_THRESH = 0.25 +MASK_THRESH = 0.5 + +_global_rknn = None + +# --------------------------- +# RKNN 全局加载 +# --------------------------- +def init_rknn_model(model_path): + global _global_rknn + if _global_rknn is not None: + return _global_rknn + + 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 Seg 模型加载成功: {model_path}") + return rknn + +# --------------------------- +# 工具函数 +# --------------------------- +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 + +# --------------------------- +# RANSAC 直线拟合（核心新增） +# --------------------------- +def fit_line_ransac(pts, max_dist=2.5, min_inliers_ratio=0.6, iters=100): + """ + 拟合 x = m*y + b + pts: Nx2 -> [x,y] + """ + if len(pts) < 10: + return None + + xs = pts[:, 0] + ys = pts[:, 1] + + best_m, best_b = None, None + best_inliers = 0 + + for _ in range(iters): + idx = np.random.choice(len(pts), 2, replace=False) + y1, y2 = ys[idx] + x1, x2 = xs[idx] + if abs(y2 - y1) < 1e-3: + continue + + m = (x2 - x1) / (y2 - y1) + b = x1 - m * y1 + + x_pred = m * ys + b + dist = np.abs(xs - x_pred) + inliers = dist < max_dist + cnt = np.sum(inliers) + + if cnt > best_inliers: + best_inliers = cnt + best_m, best_b = m, b + + if best_m is None: + return None + + if best_inliers / len(pts) < min_inliers_ratio: + return None + + return best_m, best_b + +# --------------------------- +# Seg 推理 +# --------------------------- +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, ...] # BGR -> RGB + 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) + +# --------------------------- +# PC 后处理 +# --------------------------- +def extract_left_right_edge_points(mask_bin): + h, w = mask_bin.shape + left_pts, right_pts = [], [] + for y in range(h): + xs = np.where(mask_bin[y] > 0)[0] + if len(xs) >= 2: + left_pts.append([xs.min(), y]) + right_pts.append([xs.max(), y]) + return np.array(left_pts), np.array(right_pts) + +def filter_by_seg_y_ratio(pts, y_start=0.35, y_end=0.85): + if len(pts) < 2: + return pts + y_min, y_max = pts[:, 1].min(), pts[:, 1].max() + h = y_max - y_min + if h < 10: + return pts + y0 = y_min + int(h * y_start) + y1 = y_min + int(h * y_end) + return pts[(pts[:, 1] >= y0) & (pts[:, 1] <= y1)] + +def get_y_ref(mask_bin): + h, w = mask_bin.shape + ys = [] + for x in range(int(w * 0.2), int(w * 0.8)): + y = np.where(mask_bin[:, x] > 0)[0] + if len(y): + ys.append(y.max()) + return int(np.mean(ys)) if ys else h // 2 + +# --------------------------- +# 单张图计算函数 +# --------------------------- +def caculate_yemian_diff(img, return_vis=True): + if _global_rknn is None: + raise RuntimeError("请先 init_rknn_model()") + + vis = img.copy() if return_vis else None + result_data = None + + for rx, ry, rw, rh in ROIS: + roi = img[ry:ry + rh, rx:rx + rw] + mask_bin = seg_infer(roi) // 255 + + if return_vis: + green = np.zeros_like(roi) + green[mask_bin == 1] = (0, 255, 0) + vis[ry:ry + rh, rx:rx + rw] = cv2.addWeighted( + roi, 0.7, green, 0.3, 0 + ) + + left_pts, right_pts = extract_left_right_edge_points(mask_bin) + left_pts = filter_by_seg_y_ratio(left_pts) + right_pts = filter_by_seg_y_ratio(right_pts) + + left_line = fit_line_ransac(left_pts) + right_line = fit_line_ransac(right_pts) + if left_line is None or right_line is None: + continue + + m1, b1 = left_line + m2, b2 = right_line + + y_ref = get_y_ref(mask_bin) + x_left = int(m1 * y_ref + b1) + x_right = int(m2 * y_ref + b2) + + X_L, X_R, Y = rx + x_left, rx + x_right, ry + y_ref + diff = X_R - X_L + result_data = (X_L, Y, X_R, Y, diff) + + if return_vis: + roi_vis = vis[ry:ry + rh, rx:rx + rw] + cv2.line(roi_vis, (int(b1), 0), (int(m1 * rh + b1), rh), (0, 0, 255), 3) + cv2.line(roi_vis, (int(b2), 0), (int(m2 * rh + b2), rh), (255, 0, 0), 3) + cv2.line(roi_vis, (0, y_ref), (rw, y_ref), (0, 255, 255), 2) + cv2.circle(roi_vis, (x_left, y_ref), 6, (0, 0, 255), -1) + cv2.circle(roi_vis, (x_right, y_ref), 6, (255, 0, 0), -1) + cv2.putText( + roi_vis, f"diff={diff}px", + (10, 40), + cv2.FONT_HERSHEY_SIMPLEX, + 1, (0, 255, 255), 2 + ) + + return result_data, vis + +# --------------------------- +# main +# --------------------------- +if __name__ == "__main__": + RKNN_MODEL_PATH = "seg700.rknn" + IMAGE_PATH = "7.png" + + init_rknn_model(RKNN_MODEL_PATH) + + img = cv2.imread(IMAGE_PATH) + if img is None: + raise FileNotFoundError(IMAGE_PATH) + + result_data, vis_img = caculate_yemian_diff(img, return_vis=True) + + if result_data: + XL, YL, XR, YR, diff = result_data + print(f"左交点: ({XL},{YL}) 右交点: ({XR},{YR}) diff={diff}px") + + if vis_img is not None: + cv2.imwrite("vis_output.png", vis_img) + print("可视化结果保存到 vis_output.png") + diff --git a/yemian_seg_diff_60/test_image/1.png b/yemian_seg_diff_60/test_image/1.png new file mode 100644 index 0000000..debb19b Binary files /dev/null and b/yemian_seg_diff_60/test_image/1.png differ diff --git a/yemian_seg_diff_60/test_image/2.png b/yemian_seg_diff_60/test_image/2.png new file mode 100644 index 0000000..e6fc403 Binary files /dev/null and b/yemian_seg_diff_60/test_image/2.png differ diff --git a/yemian_seg_diff_60/test_image/3.png b/yemian_seg_diff_60/test_image/3.png new file mode 100644 index 0000000..5e6e1c1 Binary files /dev/null and b/yemian_seg_diff_60/test_image/3.png differ diff --git a/yemian_seg_diff/debug_mid/111.png b/yemian_seg_diff_61/debug_mid/111.png similarity index 100% rename from yemian_seg_diff/debug_mid/111.png rename to yemian_seg_diff_61/debug_mid/111.png diff --git a/yemian_seg_diff/debug_mid/roi0_input_640.png b/yemian_seg_diff_61/debug_mid/roi0_input_640.png similarity index 100% rename from yemian_seg_diff/debug_mid/roi0_input_640.png rename to yemian_seg_diff_61/debug_mid/roi0_input_640.png diff --git a/yemian_seg_diff/debug_mid/roi0_proto_mean.png b/yemian_seg_diff_61/debug_mid/roi0_proto_mean.png similarity index 100% rename from yemian_seg_diff/debug_mid/roi0_proto_mean.png rename to yemian_seg_diff_61/debug_mid/roi0_proto_mean.png diff --git a/yemian_seg_diff/debug_mid/roi0_stride16_3.png b/yemian_seg_diff_61/debug_mid/roi0_stride16_3.png similarity index 100% rename from yemian_seg_diff/debug_mid/roi0_stride16_3.png rename to yemian_seg_diff_61/debug_mid/roi0_stride16_3.png diff --git a/yemian_seg_diff/debug_mid/roi0_stride16_4.png b/yemian_seg_diff_61/debug_mid/roi0_stride16_4.png similarity index 100% rename from yemian_seg_diff/debug_mid/roi0_stride16_4.png rename to yemian_seg_diff_61/debug_mid/roi0_stride16_4.png diff --git a/yemian_seg_diff/debug_mid/roi0_stride16_5.png b/yemian_seg_diff_61/debug_mid/roi0_stride16_5.png similarity index 100% rename from yemian_seg_diff/debug_mid/roi0_stride16_5.png rename to yemian_seg_diff_61/debug_mid/roi0_stride16_5.png diff --git a/yemian_seg_diff/debug_mid/roi0_stride32_6.png b/yemian_seg_diff_61/debug_mid/roi0_stride32_6.png similarity index 100% rename from yemian_seg_diff/debug_mid/roi0_stride32_6.png rename to yemian_seg_diff_61/debug_mid/roi0_stride32_6.png diff --git a/yemian_seg_diff/debug_mid/roi0_stride32_7.png b/yemian_seg_diff_61/debug_mid/roi0_stride32_7.png similarity index 100% rename from yemian_seg_diff/debug_mid/roi0_stride32_7.png rename to yemian_seg_diff_61/debug_mid/roi0_stride32_7.png diff --git a/yemian_seg_diff/debug_mid/roi0_stride32_8.png b/yemian_seg_diff_61/debug_mid/roi0_stride32_8.png similarity index 100% rename from yemian_seg_diff/debug_mid/roi0_stride32_8.png rename to yemian_seg_diff_61/debug_mid/roi0_stride32_8.png diff --git a/yemian_seg_diff/debug_mid/roi0_stride8_0.png b/yemian_seg_diff_61/debug_mid/roi0_stride8_0.png similarity index 100% rename from yemian_seg_diff/debug_mid/roi0_stride8_0.png rename to yemian_seg_diff_61/debug_mid/roi0_stride8_0.png diff --git a/yemian_seg_diff/debug_mid/roi0_stride8_1.png b/yemian_seg_diff_61/debug_mid/roi0_stride8_1.png similarity index 100% rename from yemian_seg_diff/debug_mid/roi0_stride8_1.png rename to yemian_seg_diff_61/debug_mid/roi0_stride8_1.png diff --git a/yemian_seg_diff/debug_mid/roi0_stride8_2.png b/yemian_seg_diff_61/debug_mid/roi0_stride8_2.png similarity index 100% rename from yemian_seg_diff/debug_mid/roi0_stride8_2.png rename to yemian_seg_diff_61/debug_mid/roi0_stride8_2.png diff --git a/yemian_seg_diff/debug_mid/zhongjianjieguo.py b/yemian_seg_diff_61/debug_mid/zhongjianjieguo.py similarity index 100% rename from yemian_seg_diff/debug_mid/zhongjianjieguo.py rename to yemian_seg_diff_61/debug_mid/zhongjianjieguo.py diff --git a/yemian_seg_diff/main.py b/yemian_seg_diff_61/main.py similarity index 100% rename from yemian_seg_diff/main.py rename to yemian_seg_diff_61/main.py diff --git a/yemian_seg_diff/seg500.rknn b/yemian_seg_diff_61/seg500.rknn similarity index 100% rename from yemian_seg_diff/seg500.rknn rename to yemian_seg_diff_61/seg500.rknn diff --git a/yemian_seg_diff/seg700.rknn b/yemian_seg_diff_61/seg700.rknn similarity index 100% rename from yemian_seg_diff/seg700.rknn rename to yemian_seg_diff_61/seg700.rknn diff --git a/yemian_seg_diff/seg_old.rknn b/yemian_seg_diff_61/seg_old.rknn similarity index 100% rename from yemian_seg_diff/seg_old.rknn rename to yemian_seg_diff_61/seg_old.rknn diff --git a/yemian_seg_diff/test_image/1 (copy 1).png b/yemian_seg_diff_61/test_image/1 (copy 1).png similarity index 100% rename from yemian_seg_diff/test_image/1 (copy 1).png rename to yemian_seg_diff_61/test_image/1 (copy 1).png diff --git a/yemian_seg_diff/test_image/1.png b/yemian_seg_diff_61/test_image/1.png similarity index 100% rename from yemian_seg_diff/test_image/1.png rename to yemian_seg_diff_61/test_image/1.png diff --git a/yemian_seg_diff/test_image/2 (copy 1).png b/yemian_seg_diff_61/test_image/2 (copy 1).png similarity index 100% rename from yemian_seg_diff/test_image/2 (copy 1).png rename to yemian_seg_diff_61/test_image/2 (copy 1).png diff --git a/yemian_seg_diff/test_image/2.png b/yemian_seg_diff_61/test_image/2.png similarity index 100% rename from yemian_seg_diff/test_image/2.png rename to yemian_seg_diff_61/test_image/2.png diff --git a/yemian_seg_diff/test_image/3 (copy 1).png b/yemian_seg_diff_61/test_image/3 (copy 1).png similarity index 100% rename from yemian_seg_diff/test_image/3 (copy 1).png rename to yemian_seg_diff_61/test_image/3 (copy 1).png diff --git a/yemian_seg_diff/test_image/3.png b/yemian_seg_diff_61/test_image/3.png similarity index 100% rename from yemian_seg_diff/test_image/3.png rename to yemian_seg_diff_61/test_image/3.png diff --git a/yemian_seg_diff/test_image/33.png b/yemian_seg_diff_61/test_image/33.png similarity index 100% rename from yemian_seg_diff/test_image/33.png rename to yemian_seg_diff_61/test_image/33.png diff --git a/yemian_seg_diff/test_image/4.png b/yemian_seg_diff_61/test_image/4.png similarity index 100% rename from yemian_seg_diff/test_image/4.png rename to yemian_seg_diff_61/test_image/4.png diff --git a/yemian_seg_diff/test_image/5.png b/yemian_seg_diff_61/test_image/5.png similarity index 100% rename from yemian_seg_diff/test_image/5.png rename to yemian_seg_diff_61/test_image/5.png diff --git a/yemian_seg_diff/test_image/6.png b/yemian_seg_diff_61/test_image/6.png similarity index 100% rename from yemian_seg_diff/test_image/6.png rename to yemian_seg_diff_61/test_image/6.png diff --git a/yemian_seg_diff/test_image/7.png b/yemian_seg_diff_61/test_image/7.png similarity index 100% rename from yemian_seg_diff/test_image/7.png rename to yemian_seg_diff_61/test_image/7.png