import cv2
import numpy as np
from ultralytics import YOLO
# ================= 参数配置 =================
MODEL_PATH = 'runs/detect/yolo_swim_markers/weights/best.pt' # 训练好的模型
VIDEO_SOURCE = r"D:\test\video-test\test1.mp4" # 视频源
CONF_THRESHOLD = 0.5 # 置信度阈值
MIN_POINTS_FOR_LINE = 3 # 拟合所需最少点数
MAX_X_GAP = 300 # 相邻点最大允许X方向间距（像素）
# 加载模型
model = YOLO(MODEL_PATH)
cap = cv2.VideoCapture(VIDEO_SOURCE)
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# 推理
results = model(frame, conf=CONF_THRESHOLD)
boxes = results[0].boxes.xyxy.cpu().numpy() # [x1, y1, x2, y2]
# 提取右上角作为关键点
centers = []
for box in boxes:
x1, y1, x2, y2 = map(int, box[:4])
cx = x2 # 右上角 x 坐标
cy = y1 # 右上角 y 坐标
centers.append([cx, cy])
# 可视化检测框和所选点
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
cv2.circle(frame, (cx, cy), 5, (0, 255, 0), -1) # 绿点：使用的右上角
filtered_centers = []
if len(centers) > 0:
# 按 x 坐标排序（从左到右）
centers = sorted(centers, key=lambda p: p[0])
# 第一个点无条件加入
filtered_centers.append(centers[0])
# 遍历后续点，判断是否与前一个有效点足够接近
for i in range(1, len(centers)):
prev_cx, prev_cy = filtered_centers[-1]
curr_cx, curr_cy = centers[i]
# 计算欧氏距离（也可只用 x 差）
dist = np.sqrt((curr_cx - prev_cx)**2 + (curr_cy - prev_cy)**2)
if dist < MAX_X_GAP: # 控制密度和连续性
filtered_centers.append(centers[i])
# 重绘筛选后的点（蓝色表示最终用于拟合的点）
for cx, cy in filtered_centers:
cv2.circle(frame, (cx, cy), 7, (255, 0, 0), -1) # 蓝色大圆点
# 只有当筛选后至少有 MIN_POINTS_FOR_LINE 个点才拟合
if len(filtered_centers) >= MIN_POINTS_FOR_LINE:
filtered_centers = np.array(filtered_centers)
# 最小二乘拟合函数
def fit_line_least_squares(points):
x = points[:, 0].reshape(-1, 1)
y = points[:, 1]
A = np.hstack([x, np.ones_like(x)])
coeff, _, _, _ = np.linalg.lstsq(A, y, rcond=None)
a, b = coeff
return lambda x: a * x + b
try:
line_func = fit_line_least_squares(filtered_centers)
# 绘制整条线（跨越图像宽度）
h, w = frame.shape[:2]
x_start, x_end = 0, w
y_start, y_end = int(line_func(x_start)), int(line_func(x_end))
cv2.line(frame, (x_start, y_start), (x_end, y_end), (255, 0, 0), 3)
cv2.putText(frame, "Fitted Lane Line", (50, 50),
cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
except Exception as e:
print("❌ 直线拟合失败:", str(e))
# 显示结果
cv2.imshow("Lane Line Fitting from Red Markers", frame)
if cv2.waitKey(50) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
# 重绘筛选后的点（蓝色表示最终用于拟合的点）
for cx, cy in filtered_centers:
cv2.circle(frame, (cx, cy), 7, (255, 0, 0), -1) # 蓝色大圆点
保存蓝色原点，在蓝色原点出现在画面并且满足拟合直线条件时拟合直线
最新发布
本文介绍了如何利用Open3D库对二维平面上的点云数据进行最小二乘拟合，以估计圆的中心坐标和半径。首先，通过pip安装Open3D，接着加载点云数据，将其转换为二维平面坐标，然后使用Open3D的最小二乘拟合函数进行估计。最后，输出估计的圆心和半径，为处理和分析点云中的圆形结构提供便利。
订阅专栏 解锁全文
扫一扫
482
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
