基于OpenCV的轮廓物体定向拟合直线：现有方案失效求替代方法

我正在使用OpenCV（cv2）计算掩码图像中物体的拟合直线，数据集内物体的方向/形状（水平、垂直等）存在差异。目前采用的方法可靠性不足，仅在少量图像中有效，无法在其他掩码图像中准确绘制符合物体方向的直线，恳请各位提供可靠的替代实现方案。

原始掩码图像 #

预期直线绘制效果（贴合物体方向） #

当前实现代码 #

import numpy as np
import cv2
import matplotlib.pyplot as plt

image_bgr = cv2.imread(IMAGE_PATH)

mask = masks[2]

mask_uint8 = mask.astype(np.uint8) * 255

contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

for c in contours:
    # 计算轮廓的质心（中心点）
    M = cv2.moments(c)
    cx = int(M['m10'] / M['m00'])
    cy = int(M['m01'] / M['m00'])
    
    cv2.drawContours(image_bgr, [c], -1, (255, 0, 0), 3)
    cv2.circle(image_bgr, (cx, cy), 5, (0, 255, 0), -1)
    
    left_side_point = tuple(c[c[:, :, 0].argmin()][0])
    right_side_point = tuple(c[c[:, :, 0].argmax()][0])
    center_point = (cx, cy)
    
    # 代码存在错误：left_side_point仅含(x,y)两个元素，索引2会触发越界
    left_center_point = ((left_side_point[0] + center_point[0]) // 2, (left_side_point[2] + center_point[2]) // 2)
    right_center_point = ((right_side_point[0] + center_point[0]) // 2, (right_side_point[2] + center_point[2]) // 2)
    
    cv2.line(image_bgr, left_side_point, left_center_point, (0, 0, 255), 2)
    cv2.line(image_bgr, left_center_point, center_point, (0, 0, 255), 2)
    cv2.line(image_bgr, center_point, right_center_point, (0, 0, 255), 2)
    cv2.line(image_bgr, right_center_point, right_side_point, (0, 0, 255), 2)

plt.imshow(image_bgr)
plt.show()

针对不同形状的物体，以下三种方案可适配大部分场景：

方案1：最小二乘法拟合直线 #

提取掩码中所有前景像素坐标，用最小二乘法拟合全局最优直线，适用于任意方向的物体：

import numpy as np
import cv2
import matplotlib.pyplot as plt

image_bgr = cv2.imread(IMAGE_PATH)
mask = masks[2].astype(np.uint8) * 255

# 获取所有前景像素的坐标
y_coords, x_coords = np.where(mask == 255)
points = np.column_stack((x_coords, y_coords))

if len(points) > 1:
    # 计算均值
    mean_x = np.mean(x_coords)
    mean_y = np.mean(y_coords)
    
    # 计算直线斜率和截距
    numerator = np.sum((x_coords - mean_x) * (y_coords - mean_y))
    denominator = np.sum((x_coords - mean_x)**2)
    
    if denominator != 0:
        m = numerator / denominator  # 斜率
        b = mean_y - m * mean_x      # 截距
        
        # 计算直线在图像边界的两个端点
        x1 = 0
        y1 = int(m * x1 + b)
        x2 = image_bgr.shape[1] - 1
        y2 = int(m * x2 + b)
        
        # 绘制拟合直线
        cv2.line(image_bgr, (x1, y1), (x2, y2), (0, 0, 255), 2)

plt.imshow(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB))
plt.show()

方案2：椭圆拟合获取主方向 #

通过拟合物体外接椭圆，以椭圆长轴作为物体方向直线，适合轴对称性较强的物体：

import numpy as np
import cv2
import matplotlib.pyplot as plt

image_bgr = cv2.imread(IMAGE_PATH)
mask = masks[2].astype(np.uint8) * 255

contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

for c in contours:
    if len(c) >= 5:  # 椭圆拟合至少需要5个点
        # 拟合椭圆
        ellipse = cv2.fitEllipse(c)
        center, axes, angle = ellipse
        
        # 计算长轴的两个端点
        major_axis_length = max(axes)
        angle_rad = np.deg2rad(angle)
        dx = np.cos(angle_rad) * major_axis_length / 2
        dy = np.sin(angle_rad) * major_axis_length / 2
        
        pt1 = (int(center[0] - dx), int(center[1] - dy))
        pt2 = (int(center[0] + dx), int(center[1] + dy))
        
        # 绘制长轴直线
        cv2.line(image_bgr, pt1, pt2, (0, 0, 255), 2)
        # 可选：绘制拟合椭圆
        cv2.ellipse(image_bgr, ellipse, (255, 0, 0), 2)

plt.imshow(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB))
plt.show()

方案3：PCA主成分分析提取方向 #

通过PCA提取物体的主成分方向，该方向即为物体的延伸方向，适配任意形状的物体：

import numpy as np
import cv2
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA

image_bgr = cv2.imread(IMAGE_PATH)
mask = masks[2].astype(np.uint8) * 255

# 获取所有前景像素坐标
y, x = np.where(mask == 255)
points = np.column_stack((x, y))

if len(points) > 1:
    # 拟合PCA模型
    pca = PCA(n_components=2)
    pca.fit(points)
    
    # 获取主方向向量和质心
    direction = pca.components_[0]
    mean = pca.mean_
    
    # 计算直线方程：ax + by + c = 0
    a = direction[1]
    b = -direction[0]
    c = -(a * mean[0] + b * mean[1])
    
    # 找到直线与图像边界的交点
    def get_boundary_points(a, b, c, img_shape):
        h, w = img_shape[:2]
        boundary_pts = []
        # 左边界x=0
        if b != 0:
            y = int(-c / b)
            if 0 <= y < h:
                boundary_pts.append((0, y))
        # 右边界x=w-1
        if b != 0:
            y = int((-c - a*(w-1)) / b)
            if 0 <= y < h:
                boundary_pts.append((w-1, y))
        # 上边界y=0
        if a != 0:
            x = int(-c / a)
            if 0 <= x < w:
                boundary_pts.append((x, 0))
        # 下边界y=h-1
        if a != 0:
            x = int((-c - b*(h-1)) / a)
            if 0 <= x < w:
                boundary_pts.append((x, h-1))
        # 取距离质心最远的两个点
        if len(boundary_pts) >= 2:
            dists = [np.linalg.norm(np.array(p)-mean) for p in boundary_pts]
            idx1 = np.argmax(dists)
            dists[idx1] = -1
            idx2 = np.argmax(dists)
            return boundary_pts[idx1], boundary_pts[idx2]
        return None
    
    boundary_pts = get_boundary_points(a, b, c, image_bgr.shape)
    if boundary_pts:
        cv2.line(image_bgr, boundary_pts[0], boundary_pts[1], (0, 0, 255), 2)

plt.imshow(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB))
plt.show()

内容的提问来源于stack exchange，提问作者Alok Chauhan