228 lines
8.0 KiB
Python
228 lines
8.0 KiB
Python
|
# ===================================================
|
|||
|
|
# final_compare_corner.py
|
|||
|
|
# 同时显示 Canny 物理边缘（红线）和 YOLO 预测左边缘（绿线）
|
|||
|
|
# 基于角点拟合直线，剔除离群点
|
|||
|
|
# ===================================================
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
from pathlib import Path
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
|
|||
|
|
# ============================
|
|||
|
|
# 参数
|
|||
|
|
# ============================
|
|||
|
|
TARGET_SIZE = 640
|
|||
|
|
MAX_CORNERS = 200
|
|||
|
|
QUALITY_LEVEL = 0.01
|
|||
|
|
MIN_DISTANCE = 5
|
|||
|
|
DIST_THRESH = 15
|
|||
|
|
ROIS = [(859, 810, 696, 328)] # 全局 ROI，可按需修改
|
|||
|
|
OUTPUT_DIR = "./final_output"
|
|||
|
|
|
|||
|
|
# ============================
|
|||
|
|
# Canny 边缘部分（保持原逻辑）
|
|||
|
|
# ============================
|
|||
|
|
def load_global_rois(txt_path):
|
|||
|
|
rois = []
|
|||
|
|
if not os.path.exists(txt_path):
|
|||
|
|
print(f"❌ ROI 文件不存在: {txt_path}")
|
|||
|
|
return rois
|
|||
|
|
with open(txt_path, 'r') as f:
|
|||
|
|
for line in f:
|
|||
|
|
s = line.strip()
|
|||
|
|
if s:
|
|||
|
|
try:
|
|||
|
|
x, y, w, h = map(int, s.split(','))
|
|||
|
|
rois.append((x, y, w, h))
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"⚠️ 无法解析 ROI 行 '{s}': {e}")
|
|||
|
|
return rois
|
|||
|
|
|
|||
|
|
def fit_line_best(points, distance_thresh=5, max_iter=5):
|
|||
|
|
if len(points) < 2:
|
|||
|
|
return None
|
|||
|
|
points = points.astype(np.float32)
|
|||
|
|
for _ in range(max_iter):
|
|||
|
|
mean = np.mean(points, axis=0)
|
|||
|
|
cov = np.cov(points.T)
|
|||
|
|
eigvals, eigvecs = np.linalg.eig(cov)
|
|||
|
|
idx = np.argmax(eigvals)
|
|||
|
|
direction = eigvecs[:, idx]
|
|||
|
|
vx, vy = direction
|
|||
|
|
x0, y0 = mean
|
|||
|
|
dists = np.abs(vy*(points[:,0]-x0) - vx*(points[:,1]-y0)) / np.hypot(vx, vy)
|
|||
|
|
inliers = points[dists <= distance_thresh]
|
|||
|
|
if len(inliers) == len(points) or len(inliers) < 2:
|
|||
|
|
break
|
|||
|
|
points = inliers
|
|||
|
|
if len(points) < 2:
|
|||
|
|
return None
|
|||
|
|
X = points[:, 0].reshape(-1, 1)
|
|||
|
|
y = points[:, 1]
|
|||
|
|
try:
|
|||
|
|
from sklearn.linear_model import RANSACRegressor
|
|||
|
|
ransac = RANSACRegressor(residual_threshold=distance_thresh)
|
|||
|
|
ransac.fit(X, y)
|
|||
|
|
k = ransac.estimator_.coef_[0]
|
|||
|
|
b = ransac.estimator_.intercept_
|
|||
|
|
vx = 1 / np.sqrt(1 + k**2)
|
|||
|
|
vy = k / np.sqrt(1 + k**2)
|
|||
|
|
x0 = np.mean(points[:,0])
|
|||
|
|
y0 = k*x0 + b
|
|||
|
|
except:
|
|||
|
|
mean = np.mean(points, axis=0)
|
|||
|
|
cov = np.cov(points.T)
|
|||
|
|
eigvals, eigvecs = np.linalg.eig(cov)
|
|||
|
|
idx = np.argmax(eigvals)
|
|||
|
|
direction = eigvecs[:, idx]
|
|||
|
|
vx, vy = direction
|
|||
|
|
x0, y0 = mean
|
|||
|
|
return vx, vy, x0, y0
|
|||
|
|
|
|||
|
|
def extract_canny_overlay(image_path, roi_file, distance_thresh=3):
|
|||
|
|
img = cv2.imread(image_path)
|
|||
|
|
if img is None:
|
|||
|
|
print(f"❌ 无法读取图片: {image_path}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
overlay_img = img.copy()
|
|||
|
|
rois = load_global_rois(roi_file)
|
|||
|
|
if not rois:
|
|||
|
|
print("❌ 没有有效 ROI")
|
|||
|
|
return overlay_img
|
|||
|
|
|
|||
|
|
for idx, (x, y, w, h) in enumerate(rois):
|
|||
|
|
roi = img[y:y+h, x:x+w]
|
|||
|
|
gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
|
|||
|
|
edges = cv2.Canny(gray, 100, 200)
|
|||
|
|
contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
|||
|
|
longest_contour = max(contours, key=lambda c: cv2.arcLength(c, closed=False), default=None)
|
|||
|
|
|
|||
|
|
if longest_contour is not None and len(longest_contour) >= 2:
|
|||
|
|
points = longest_contour.reshape(-1, 2)
|
|||
|
|
line = fit_line_best(points, distance_thresh=distance_thresh, max_iter=5)
|
|||
|
|
if line is not None:
|
|||
|
|
vx, vy, x0, y0 = line
|
|||
|
|
cols = w
|
|||
|
|
lefty = int(y0 - vy/vx * x0)
|
|||
|
|
righty = int(y0 + vy/vx * (cols - x0))
|
|||
|
|
pt1 = (x, y + lefty)
|
|||
|
|
pt2 = (x + cols - 1, y + righty)
|
|||
|
|
cv2.line(overlay_img, pt1, pt2, (0, 0, 255), 2) # 红色
|
|||
|
|
print(f"✅ ROI {idx} Canny 边缘拟合完成")
|
|||
|
|
return overlay_img
|
|||
|
|
|
|||
|
|
# ============================
|
|||
|
|
# YOLO 角点 + 拟合直线
|
|||
|
|
# ============================
|
|||
|
|
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)
|
|||
|
|
left_thresh = int(w*left_ratio)
|
|||
|
|
right_thresh = w - 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)
|
|||
|
|
_, idx = np.unique(selected.reshape(-1,2), axis=0, return_index=True)
|
|||
|
|
selected = selected[np.sort(idx)]
|
|||
|
|
return selected.astype(np.int32)
|
|||
|
|
|
|||
|
|
return filter_by_y_variation(left_candidates), filter_by_y_variation(right_candidates)
|
|||
|
|
|
|||
|
|
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, y = pts[:,0], pts[:,1]
|
|||
|
|
try:
|
|||
|
|
m, b = np.polyfit(y, x, 1)
|
|||
|
|
except:
|
|||
|
|
return None, pts.astype(np.int32)
|
|||
|
|
x_fit = m*y + b
|
|||
|
|
mask = np.abs(x-x_fit)<dist_thresh
|
|||
|
|
if np.sum(mask)<2:
|
|||
|
|
return (m,b), pts.astype(np.int32)
|
|||
|
|
x2, y2 = x[mask], y[mask]
|
|||
|
|
m2, b2 = np.polyfit(y2, x2, 1)
|
|||
|
|
inliers = np.stack([x2,y2],axis=1).astype(np.int32)
|
|||
|
|
return (m2,b2), inliers
|
|||
|
|
|
|||
|
|
def get_yolo_left_edge_lines(image_path, model_path, rois=ROIS, imgsz=TARGET_SIZE):
|
|||
|
|
model = YOLO(model_path)
|
|||
|
|
img = cv2.imread(image_path)
|
|||
|
|
if img is None:
|
|||
|
|
print(f"❌ 无法读取图片: {image_path}")
|
|||
|
|
return []
|
|||
|
|
lines = []
|
|||
|
|
|
|||
|
|
for (x, y, w, h) in rois:
|
|||
|
|
roi_img = img[y:y+h, x:x+w]
|
|||
|
|
resized = cv2.resize(roi_img, (imgsz, imgsz))
|
|||
|
|
results = model(resized, imgsz=imgsz, verbose=False)
|
|||
|
|
result = results[0]
|
|||
|
|
if result.masks is None or len(result.masks.data)==0:
|
|||
|
|
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=MAX_CORNERS,
|
|||
|
|
qualityLevel=QUALITY_LEVEL, minDistance=MIN_DISTANCE)
|
|||
|
|
left_pts, _ = select_edge_corners(corners, w)
|
|||
|
|
line_params, inliers = fit_line_with_outlier_removal(left_pts)
|
|||
|
|
if line_params is None:
|
|||
|
|
continue
|
|||
|
|
m,b = line_params
|
|||
|
|
y1, y2 = 0, h-1
|
|||
|
|
x1 = int(m*y1 + b)
|
|||
|
|
x2 = int(m*y2 + b)
|
|||
|
|
lines.append((x+x1, y+y1, x+x2, y+y2))
|
|||
|
|
return lines
|
|||
|
|
|
|||
|
|
# ============================
|
|||
|
|
# 对比融合
|
|||
|
|
# ============================
|
|||
|
|
def compare_canny_vs_yolo(image_path, canny_roi_file, model_path, output_dir=OUTPUT_DIR):
|
|||
|
|
output_dir = Path(output_dir)
|
|||
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
canny_img = extract_canny_overlay(image_path, canny_roi_file, distance_thresh=6)
|
|||
|
|
if canny_img is None:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
yolo_lines = get_yolo_left_edge_lines(image_path, model_path)
|
|||
|
|
result_img = canny_img.copy()
|
|||
|
|
for x1,y1,x2,y2 in yolo_lines:
|
|||
|
|
cv2.line(result_img, (x1,y1), (x2,y2), (0,255,0), 2) # 绿色
|
|||
|
|
cv2.circle(result_img, (x1,y1), 4, (255,0,0), -1) # 蓝色起点
|
|||
|
|
|
|||
|
|
output_path = output_dir / f"compare_{Path(image_path).stem}.jpg"
|
|||
|
|
cv2.imwrite(str(output_path), result_img)
|
|||
|
|
print(f"✅ 对比图已保存: {output_path}")
|
|||
|
|
|
|||
|
|
# ============================
|
|||
|
|
# 使用示例
|
|||
|
|
# ============================
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
IMAGE_PATH = "../test_image/2.jpg"
|
|||
|
|
CANNY_ROI_FILE = "../roi_coordinates/1_rois1.txt"
|
|||
|
|
MODEL_PATH = "best.pt"
|
|||
|
|
|
|||
|
|
compare_canny_vs_yolo(IMAGE_PATH, CANNY_ROI_FILE, MODEL_PATH)
|