150 lines
4.7 KiB
Python
150 lines
4.7 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from sklearn.linear_model import RANSACRegressor
|
|
|
|
# ---------------------------
|
|
# 读取 ROI 列表 txt
|
|
# ---------------------------
|
|
def load_rois_from_txt(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
|
|
|
|
# ---------------------------
|
|
# PCA + RANSAC + 迭代去离群点拟合直线
|
|
# ---------------------------
|
|
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
|
|
|
|
# RANSAC 拟合 y = kx + b
|
|
X = points[:, 0].reshape(-1, 1)
|
|
y = points[:, 1]
|
|
try:
|
|
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
|
|
|
|
# ---------------------------
|
|
# 封装函数：读取 ROI txt -> 拟合直线 -> 可视化
|
|
# ---------------------------
|
|
def fit_lines_from_image_txt(image_path, roi_txt_path, distance_thresh=5, draw_overlay=True):
|
|
"""
|
|
输入：
|
|
image_path: 原图路径
|
|
roi_txt_path: ROI txt 文件路径，每行 x,y,w,h
|
|
distance_thresh: 直线拟合残差阈值
|
|
draw_overlay: 是否在原图上叠加拟合直线
|
|
输出：
|
|
lines: 每个 ROI 的拟合直线 [(vx, vy, x0, y0), ...]
|
|
overlay_img: 可视化原图叠加拟合直线
|
|
"""
|
|
img = cv2.imread(image_path)
|
|
if img is None:
|
|
print(f"❌ 无法读取图片: {image_path}")
|
|
return [], None
|
|
|
|
rois = load_rois_from_txt(roi_txt_path)
|
|
if not rois:
|
|
print("❌ 没有有效 ROI")
|
|
return [], None
|
|
|
|
overlay_img = img.copy() if draw_overlay else None
|
|
lines = []
|
|
|
|
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:
|
|
vx, vy, x0, y0 = line
|
|
lines.append(line)
|
|
|
|
if draw_overlay:
|
|
cols = gray.shape[1]
|
|
lefty = int(y0 - vy/vx * x0)
|
|
righty = int(y0 + vy/vx * (cols - x0))
|
|
|
|
# 绘制在原图
|
|
cv2.line(overlay_img, (x, y + lefty), (x + cols - 1, y + righty), (0, 0, 255), 2)
|
|
cv2.drawContours(overlay_img, [longest_contour + np.array([x, y])], -1, (0, 255, 0), 1)
|
|
else:
|
|
lines.append(None)
|
|
else:
|
|
lines.append(None)
|
|
|
|
return lines, overlay_img
|
|
|
|
# ---------------------------
|
|
# 使用示例
|
|
# ---------------------------
|
|
if __name__ == "__main__":
|
|
image_path = "../test_image/1.jpg"
|
|
roi_txt_path = "../roi_coordinates/1_rois1.txt"
|
|
|
|
lines, overlay = fit_lines_from_image_txt(image_path, roi_txt_path, distance_thresh=5)
|
|
|
|
for idx, line in enumerate(lines):
|
|
print(f"ROI {idx} 拟合直线: {line}")
|
|
|
|
if overlay is not None:
|
|
cv2.imwrite("overlay_result.jpg", overlay)
|
|
print("✅ 原图叠加拟合直线已保存: overlay_result.jpg")
|