191 lines
6.3 KiB
Python
191 lines
6.3 KiB
Python
#!/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 = 15 # 剔除离群点阈值
|
|
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()
|