#!/usr/bin/env python3 # -*- coding: utf-8 -*- import os import cv2 import numpy as np from pathlib import Path from ultralytics import YOLO # -------------------- # 参数设置（固定在脚本中） # -------------------- INPUT_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/test_l" # 图片文件夹 MODEL_PATH = "best.pt" # YOLO 模型 OUTPUT_DIR = "./output" # 保存结果 TARGET_SIZE = 640 # YOLO 输入尺寸 DIST_THRESH = 10 # 剔除离群点阈值 MAX_CORNERS = 200 # goodFeaturesToTrack 最大角点数 QUALITY_LEVEL = 0.01 # goodFeaturesToTrack qualityLevel MIN_DISTANCE = 5 # goodFeaturesToTrack minDistance # 全局 ROI 定义 ROIS = [ (859, 810, 696, 328), # (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 np.zeros((0,2), dtype=np.int32), np.zeros((0,2), dtype=np.int32) 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] def filter_by_y_variation(pts): if len(pts) < 2: return pts.astype(np.int32) pts_sorted = pts[np.argsort(pts[:,1])] diffs = np.abs(np.diff(pts_sorted[:,1])) keep_idx = np.where(diffs > y_var_thresh)[0] if len(keep_idx) == 0: return pts_sorted.astype(np.int32) selected = [pts_sorted[i] for i in keep_idx] + [pts_sorted[i+1] for i in keep_idx] selected = np.array(selected) selected = selected[np.argsort(selected[:,1])] _, idx = np.unique(selected.reshape(-1,2), axis=0, return_index=True) selected = selected[np.sort(idx)] return selected.astype(np.int32) 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=DIST_THRESH): if pts is None or len(pts) < 2: return None, np.zeros((0,2), dtype=np.int32) pts = np.array(pts, dtype=np.float64) x = pts[:,0] y = pts[:,1] try: m, b = np.polyfit(y, x, 1) except: return None, np.zeros((0,2), dtype=np.int32) x_fit = m*y + b dists = np.abs(x - x_fit) mask = dists < dist_thresh if np.sum(mask) < 2: return (m,b), pts.astype(np.int32) x2, y2 = x[mask], y[mask] try: m2, b2 = np.polyfit(y2, x2, 1) except: return (m,b), np.stack([x2,y2],axis=1).astype(np.int32) inliers = np.stack([x2,y2],axis=1).astype(np.int32) return (m2,b2), inliers # -------------------- # 单张图 ROI 处理 # -------------------- def process_roi_on_image(orig_img, roi): rx, ry, rw, rh = roi h_img, w_img = orig_img.shape[:2] rx = max(0, rx); ry = max(0, ry) rw = min(rw, w_img - rx); rh = min(rh, h_img - ry) roi_img = orig_img[ry:ry+rh, rx:rx+rw].copy() if roi_img.size == 0: return None resized = cv2.resize(roi_img, (TARGET_SIZE, TARGET_SIZE)) results = MODEL(resized, imgsz=TARGET_SIZE, verbose=False) result = results[0] overlay_roi = roi_img.copy() if result.masks is None or len(result.masks.data)==0: return overlay_roi mask = result.masks.data[0].cpu().numpy() mask_bin = (mask>0.5).astype(np.uint8) mask_bin = cv2.resize(mask_bin,(rw,rh), interpolation=cv2.INTER_NEAREST) # mask 半透明覆盖 color_mask = np.zeros_like(overlay_roi, dtype=np.uint8) color_mask[mask_bin==1] = (0,255,0) overlay_roi = cv2.addWeighted(overlay_roi,0.7,color_mask,0.3,0) # 角点检测 mask_gray = (mask_bin*255).astype(np.uint8) corners = cv2.goodFeaturesToTrack(mask_gray, maxCorners=MAX_CORNERS, qualityLevel=QUALITY_LEVEL, minDistance=MIN_DISTANCE) left_pts, right_pts = select_edge_corners(corners, rw) left_line, left_inliers = fit_line_with_outlier_removal(left_pts) right_line, right_inliers = fit_line_with_outlier_removal(right_pts) # 可视化 inliers for (cx,cy) in left_inliers: cv2.circle(overlay_roi,(int(cx),int(cy)),4,(0,0,255),-1) for (cx,cy) in right_inliers: cv2.circle(overlay_roi,(int(cx),int(cy)),4,(255,0,0),-1) # 拟合直线 if left_line is not None: m,b = left_line y1,y2 = 0, rh-1 x1 = int(m*y1+b); x2 = int(m*y2+b) cv2.line(overlay_roi,(x1,y1),(x2,y2),(0,0,200),3) if right_line is not None: m,b = right_line y1,y2 = 0, rh-1 x1 = int(m*y1+b); x2 = int(m*y2+b) cv2.line(overlay_roi,(x1,y1),(x2,y2),(200,0,0),3) return overlay_roi # -------------------- # 批量推理文件夹 # -------------------- def infer_folder_images(): input_dir = Path(INPUT_DIR) output_dir = Path(OUTPUT_DIR) output_dir.mkdir(parents=True, exist_ok=True) exts = ('*.jpg','*.jpeg','*.png','*.bmp','*.tif','*.tiff') files = [] for e in exts: files.extend(sorted(input_dir.glob(e))) if len(files)==0: print("未找到图片文件") return print(f"找到 {len(files)} 张图片，开始推理...") for img_path in files: print("-> 处理:", img_path.name) orig_img = cv2.imread(str(img_path)) if orig_img is None: print(" 无法读取，跳过") continue out_img = orig_img.copy() for roi in ROIS: overlay_roi = process_roi_on_image(orig_img, roi) if overlay_roi is not None: rx,ry,rw,rh = roi h,w = overlay_roi.shape[:2] out_img[ry:ry+h, rx:rx+w] = overlay_roi save_path = output_dir / f"mask_edge_corners_{img_path.name}" cv2.imwrite(str(save_path), out_img) print(" 已保存 ->", save_path.name) print("批量推理完成，结果保存在:", output_dir) # -------------------- # 主程序 # -------------------- if __name__ == "__main__": MODEL = YOLO(MODEL_PATH) infer_folder_images()