import os import cv2 import numpy as np from pathlib import Path from ultralytics import YOLO import matplotlib.pyplot as plt TARGET_SIZE = 640 # 模型输入尺寸 IMAGE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/611/train/class0" # 👈 图片文件夹路径 MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/61seg/exp5/weights/best.pt" OUTPUT_DIR = "./output" # -------------------- # 全局 ROI 定义 # -------------------- ROIS = [ (670,623,465,178), # (x, y, w, h) ] # -------------------- # 根据角点分布，选取左右边缘角点 # -------------------- def select_edge_corners(corners, w, left_ratio=0.2, right_ratio=0.2, y_var_thresh=5): if corners is None: return [], [] corners = np.int32(corners).reshape(-1, 2) x_min, x_max = 0, w left_thresh = x_min + int(w * left_ratio) right_thresh = x_max - int(w * right_ratio) # 左右候选角点 left_candidates = corners[corners[:, 0] <= left_thresh] right_candidates = corners[corners[:, 0] >= right_thresh] # 按 y 变化筛选 def filter_by_y_variation(pts): if len(pts) < 2: return pts pts_sorted = pts[np.argsort(pts[:, 1])] diffs = np.abs(np.diff(pts_sorted[:, 1])) keep_idx = np.where(diffs > y_var_thresh)[0] selected = [pts_sorted[i] for i in keep_idx] + [pts_sorted[i + 1] for i in keep_idx] return np.array(selected) if len(selected) > 0 else pts_sorted left_final = filter_by_y_variation(left_candidates) right_final = filter_by_y_variation(right_candidates) return left_final, right_final # -------------------- # 拟合直线并剔除离散点 # -------------------- def fit_line_with_outlier_removal(pts, dist_thresh=10): if pts is None or len(pts) < 2: return None, pts pts = np.array(pts) x, y = pts[:, 0], pts[:, 1] # 第一次拟合 m, b = np.polyfit(y, x, 1) # x = m*y + b x_fit = m * y + b dists = np.abs(x - x_fit) # 剔除离群点 mask = dists < dist_thresh x2, y2 = x[mask], y[mask] if len(x2) < 2: return (m, b), pts # 二次拟合 m, b = np.polyfit(y2, x2, 1) return (m, b), np.stack([x2, y2], axis=1) # -------------------- # 推理 ROI 并计算左右直线 x 差 # -------------------- def infer_mask_and_get_xdiff(image_path, model, output_dir="./output"): image_path = Path(image_path) orig_img = cv2.imread(str(image_path)) overlay_img = orig_img.copy() x_diff_value = None intersections = [] # 👈 保存交点 (全图坐标) for idx, (x, y, w, h) in enumerate(ROIS): roi_img = orig_img[y:y+h, x:x+w] resized_img = cv2.resize(roi_img, (TARGET_SIZE, TARGET_SIZE)) results = model(source=resized_img, imgsz=TARGET_SIZE, verbose=False) result = results[0] if result.masks is None or len(result.masks.data) == 0: print(f"❌ {image_path.name} 未检测到 mask") continue mask = result.masks.data[0].cpu().numpy() mask_bin = (mask > 0.5).astype(np.uint8) mask_bin = cv2.resize(mask_bin, (w, h), interpolation=cv2.INTER_NEAREST) mask_gray = (mask_bin * 255).astype(np.uint8) corners = cv2.goodFeaturesToTrack(mask_gray, maxCorners=200, qualityLevel=0.01, minDistance=5) left_pts, right_pts = select_edge_corners(corners, w) left_line, _ = fit_line_with_outlier_removal(left_pts) right_line, _ = fit_line_with_outlier_removal(right_pts) if left_line and right_line: # ---------- 黄色参考线 ---------- y_ref = int(h * 0.6) # ---------- 左右交点 ---------- m1, b1 = left_line m2, b2 = right_line x_left = int(m1 * y_ref + b1) x_right = int(m2 * y_ref + b2) x_diff_value = abs(x_right - x_left) # ---------- ROI → 原图坐标 ---------- left_pt_global = (x + x_left, y + y_ref) right_pt_global = (x + x_right, y + y_ref) intersections.append({ "left": left_pt_global, "right": right_pt_global, "x_diff": x_diff_value }) # ---------- 画黄色参考线 ---------- cv2.line( overlay_img[y:y+h, x:x+w], (0, y_ref), (w, y_ref), (0, 255, 255), 2 ) # ---------- 画红 / 蓝拟合线 ---------- for (m, b), color in [(left_line, (0, 0, 255)), (right_line, (255, 0, 0))]: y1, y2 = 0, h x1_line, x2_line = int(m * y1 + b), int(m * y2 + b) cv2.line( overlay_img[y:y+h, x:x+w], (x1_line, y1), (x2_line, y2), color, 3 ) # ---------- 画交点 ---------- cv2.circle(overlay_img, left_pt_global, 6, (0, 0, 255), -1) cv2.circle(overlay_img, right_pt_global, 6, (255, 0, 0), -1) cv2.putText( overlay_img, f"L{left_pt_global} R{right_pt_global}", (x + 10, y + 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2 ) save_path = Path(output_dir) / f"mask_edge_{image_path.name}" cv2.imwrite(str(save_path), overlay_img) print(f"✅ 保存结果: {save_path}") return x_diff_value, intersections # ===================== # 批量处理文件夹中的图片 # ===================== def process_image_folder(image_dir, model_path): model = YOLO(model_path) output_dir = OUTPUT_DIR Path(output_dir).mkdir(parents=True, exist_ok=True) xdiff_results = {} for img_name in sorted(os.listdir(image_dir)): if not img_name.lower().endswith((".jpg", ".png", ".jpeg")): continue img_path = os.path.join(image_dir, img_name) x_diff = infer_mask_and_get_xdiff(img_path, model, output_dir) if x_diff is not None: xdiff_results[img_name] = x_diff # 结果展示 if xdiff_results: plt.figure(figsize=(8, 4)) plt.plot(list(xdiff_results.keys()), list(xdiff_results.values()), '-o') plt.xticks(rotation=45) plt.title("同一y坐标下左右液面直线x差变化趋势") plt.ylabel("x差 (像素)") plt.xlabel("图片序号") plt.tight_layout() plt.show() print("\n=== 液面变化结果 ===") for k, v in xdiff_results.items(): print(f"{k}: x差 = {v:.2f}") else: print("⚠️ 未检测到有效结果") # ===================== 使用示例 ===================== if __name__ == "__main__": process_image_folder(IMAGE_DIR, MODEL_PATH)