使用 Python 和 matplotlib 在 2D 数据上拟合椭圆

Fit Ellipse over 2D data using Python and matplotlib

我正在尝试使用 matplotlib 在不同的 2D 点集上拟合一个椭圆。我正在使用 here 中的椭圆拟合函数,这是代码。

from matplotlib.patches import Ellipse
import matplotlib.pyplot as plt
import numpy

def fitEllipse(x,y):
    x = x[:,np.newaxis]
    y = y[:,np.newaxis]
    D =  np.hstack((x*x, x*y, y*y, x, y, np.ones_like(x)))
    S = np.dot(D.T,D)
    C = np.zeros([6,6])
    C[0,2] = C[2,0] = 2; C[1,1] = -1
    E, V =  eig(np.dot(inv(S), C))
    n = np.argmax(np.abs(E))
    a = V[:,n]
    return a

def ellipse_center(a):
    b,c,d,f,g,a = a[1]/2, a[2], a[3]/2, a[4]/2, a[5], a[0]
    num = b*b-a*c
    x0=(c*d-b*f)/num
    y0=(a*f-b*d)/num
    return np.array([x0,y0])


def ellipse_axis_length( a ):
    b,c,d,f,g,a = a[1]/2, a[2], a[3]/2, a[4]/2, a[5], a[0]
    up = 2*(a*f*f+c*d*d+g*b*b-2*b*d*f-a*c*g)
    down1=(b*b-a*c)*( (c-a)*np.sqrt(1+4*b*b/((a-c)*(a-c)))-(c+a))
    down2=(b*b-a*c)*( (a-c)*np.sqrt(1+4*b*b/((a-c)*(a-c)))-(c+a))
    res1=np.sqrt(up/down1)
    res2=np.sqrt(up/down2)
    return np.array([res1, res2])

def ellipse_angle_of_rotation2( a ):
    b,c,d,f,g,a = a[1]/2, a[2], a[3]/2, a[4]/2, a[5], a[0]
    if b == 0:
        if a > c:
            return 0
        else:
            return np.pi/2
    else:
        if a > c:
            return np.arctan(2*b/(a-c))/2
        else:
            return np.pi/2 + np.arctan(2*b/(a-c))/2

但是,当我使用matplotlib绘制椭圆时,有时椭圆拟合得很好,有时我需要将椭圆旋转90度才能拟合(蓝色椭圆旋转90,红色椭圆没有额外旋转)这是我的代码。

def plot_ellipse(x, y):
    a = fitEllipse(x, y)

    center = ellipse_center(a)

    phi = ellipse_angle_of_rotation2(a)

    axes = ellipse_axis_length(a)

    a, b = axes

    ell = Ellipse(center, 2*a, 2*b, phi*180 / np.pi, facecolor=(1,0,0,0.2), edgecolor=(0,0,0,0.5))

    ell_rotated = Ellipse(center, 2*a, 2*b, phi*180 / np.pi + 90, facecolor=(0,0,1,0.2), edgecolor=(0,0,0,0.5))

    fig, ax = plt.subplots()
    ax.add_patch(ell)
    ax.add_patch(ell_rotated)

    plt.scatter(x, y)

    plt.show()

x1 = np.array([238, 238, 238, 237, 237, 237, 237, 237, 236, 236, 236, 236, 237, 238,
     239, 240, 240, 241, 242, 243, 243, 244, 245, 246, 247, 248, 249, 250,
     251, 252, 253, 254, 255, 255, 256, 257, 258, 259, 260, 261, 262, 263,
     264, 265, 266, 266, 267, 267, 268, 268, 269, 269, 270, 270, 271, 271,
     271, 271, 271, 272, 272, 272, 272, 272, 273, 273, 273, 273, 273, 273,
     274, 274, 274, 274, 274, 274, 275, 275, 275, 275, 275, 275, 275, 275,
     275, 275, 274, 274, 274, 274, 274, 274, 274, 274, 274, 273, 273, 273,
     272, 272, 272, 272, 271, 271, 271, 270, 270, 269, 268, 268, 267, 266,
     266, 265, 265, 264, 263, 262, 261, 260, 259, 258, 257, 256, 256, 255,
     254, 253, 252, 251, 250, 249, 248, 247, 246, 245, 245, 244, 243, 242,
     241, 240, 239, 238, 237, 236, 235, 235, 235, 234, 234, 233, 233, 232,
     232, 232, 231, 231, 230, 230, 229, 229, 229, 229, 229, 229, 229, 229,
     228, 228, 228, 228, 228, 228, 228, 228, 228, 228, 228, 228, 228, 228,
     228, 228, 228, 228, 228, 228, 228, 228, 228, 228, 228, 228, 228, 228,
     228, 229, 229, 229, 229, 229, 229, 229, 229, 229, 230, 230, 230, 230,
     231, 231, 231, 232, 232, 232, 232, 233, 233, 233, 234, 235, 236, 237,
     238, 239, 240, 241, 242])
y1 = np.array([283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 293, 294, 295,
     296, 297, 298, 299, 300, 301, 301, 301, 301, 301, 301, 301, 301, 301,
     301, 301, 301, 301, 301, 301, 301, 300, 300, 299, 299, 298, 298, 297,
     297, 296, 296, 296, 295, 294, 293, 292, 291, 290, 289, 288, 287, 286,
     286, 285, 284, 283, 282, 281, 280, 279, 278, 277, 276, 276, 275, 274,
     273, 272, 271, 270, 269, 268, 267, 266, 265, 264, 264, 263, 262, 261,
     260, 259, 258, 257, 256, 255, 254, 253, 252, 252, 251, 250, 249, 248,
     247, 246, 245, 244, 243, 242, 242, 241, 240, 239, 238, 237, 236, 235,
     234, 233, 233, 232, 232, 231, 230, 230, 229, 228, 228, 227, 227, 227,
     227, 226, 226, 226, 226, 226, 226, 225, 225, 225, 225, 226, 226, 227,
     228, 229, 229, 230, 231, 231, 232, 232, 233, 234, 235, 236, 237, 238,
     239, 240, 241, 242, 243, 244, 245, 246, 246, 247, 248, 249, 250, 251,
     252, 253, 254, 255, 256, 257, 257, 258, 259, 260, 261, 262, 263, 264,
     265, 266, 267, 268, 269, 270, 271, 272, 273, 273, 274, 275, 276, 277,
     278, 279, 280, 281, 282, 283, 284, 285, 285, 286, 287, 288, 289, 290,
     291, 292, 293, 294, 295, 296, 297, 298, 299, 299, 300, 301, 301, 302,
     303, 304, 304, 305, 306])

x2 = np.array(    [235, 236, 237, 238, 239, 240, 241, 242, 243, 243, 244, 245, 246, 247,
     248, 249, 250, 251, 252, 253, 254, 255, 255, 256, 257, 258, 259, 260,
     261, 262, 263, 264, 265, 266, 266, 267, 268, 269, 270, 270, 271, 272,
     273, 274, 274, 274, 275, 275, 276, 276, 277, 277, 278, 278, 279, 279,
     279, 279, 280, 280, 280, 280, 281, 281, 281, 281, 282, 282, 282, 282,
     282, 282, 282, 282, 282, 281, 281, 281, 281, 281, 281, 281, 281, 281,
     280, 280, 280, 280, 280, 279, 279, 279, 279, 278, 278, 277, 277, 276,
     276, 275, 275, 274, 274, 274, 273, 272, 271, 270, 269, 268, 267, 266,
     265, 264, 263, 263, 262, 261, 260, 259, 258, 257, 256, 255, 254, 253,
     252, 252, 251, 250, 249, 248, 247, 246, 245, 244, 243, 242, 241, 240,
     240, 239, 238, 237, 236, 235, 234, 233, 232, 232, 231, 231, 230, 230,
     229, 229, 228, 228, 227, 227, 227, 227, 227, 227, 227, 227, 227, 227,
     226, 226, 226, 226, 226, 226, 226, 226, 226, 226, 226, 226, 226, 227,
     227, 227, 227, 227, 227, 227, 227, 228, 228, 228, 228, 228, 228, 229,
     229, 230, 230, 231, 231, 232, 232, 233, 233, 234, 234, 235, 235, 235,
     236, 237, 238, 239, 240, 241, 242, 243, 244])

y2 = np.array(
    [279, 280, 281, 282, 283, 284, 285, 286, 287, 287, 287, 288, 288, 289,
     289, 290, 290, 290, 291, 291, 292, 292, 292, 292, 291, 291, 290, 290,
     289, 289, 288, 288, 287, 287, 287, 286, 285, 284, 283, 282, 281, 280,
     279, 278, 278, 277, 276, 275, 274, 273, 272, 271, 270, 269, 268, 267,
     267, 266, 265, 264, 263, 262, 261, 260, 259, 258, 257, 256, 255, 255,
     254, 253, 252, 251, 250, 249, 248, 247, 246, 245, 244, 244, 243, 242,
     241, 240, 239, 238, 237, 236, 235, 234, 234, 233, 232, 231, 230, 229,
     228, 227, 226, 225, 224, 224, 223, 222, 222, 221, 220, 219, 218, 217,
     217, 216, 215, 215, 215, 215, 215, 215, 215, 214, 214, 214, 214, 214,
     214, 214, 214, 215, 215, 216, 216, 217, 217, 217, 218, 218, 219, 219,
     219, 220, 221, 222, 223, 223, 224, 225, 226, 226, 227, 228, 229, 230,
     231, 232, 233, 234, 235, 236, 236, 237, 238, 239, 240, 241, 242, 243,
     244, 245, 246, 247, 248, 249, 250, 251, 251, 252, 253, 254, 255, 256,
     257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 268, 269,
     270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281, 282, 282,
     283, 284, 284, 285, 286, 287, 287, 288, 289])

plot_ellipse(x1, y1)
plot_ellipse(x2, y2)

下面是情节的屏幕截图:

x1, y1 plot

x2, y2 plot

如您所见,未旋转(红色)的椭圆很好地拟合了 x1,y1 数据,但旋转后的椭圆(蓝色)拟合了 x2,y2 数据。

我很困惑,如果我在这里遗漏了什么,什么时候需要将椭圆旋转 90º,什么时候不需要?

这是我的猜测,没有 100% 检查数学。查看定义和要解决的拉格朗日量,一切看起来都很好且合乎逻辑。我认为向量 a 是正确的,内部参数 af 也是如此。但是,代码提到要最小化的函数与比例因子无关。所以在计算轴角时,可能 运行 变成符号问题。我的猜测是,这发生在 arctan() 函数中。一种解决方案可能是在 a 向量的计算中添加符号检查,即 a -> -a if a[-1] < 0

更好并且可能也有效(我刚刚测试了 OP 的 2 个案例)是替换

if a > c:
    return np.arctan( 2 * b / ( a - c ) ) / 2
else:
    return np.pi / 2 + np.arctan( 2 * b / (a - c ) ) / 2

来自

if a > c:
    return np.arctan2( 2 * b, ( a - c ) ) / 2
else:
    return np.pi / 2 + np.arctan2( 2 * b, ( a - c) ) / 2

如果很明显 arctan 的论点来自一个部门,arctan2 就是一个选择。

补充说明:我认为在一个函数中包含多个 return 语句的代码应该省略。在这个简短的函数中它仍然很容易处理,但我不认为这是一个好的做法。

更新

在联系原始适配代码的作者时,他提到arctan2在他在GitHub

上提供的版本中使用

此代码看起来更简洁,我建议使用它而不是主页上的代码片段。

补充想法

实际上,我认为可以用更一致和更容易遵循的方式来完成,如何从 a 向量中提取椭圆的参数。所以我写下了这段代码

# -*- coding: utf-8 -*-
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
import numpy as np

RAD = 180. / np.pi
DEGREE = 1. / RAD

def rot( a ):
    """
    simple rotation matrix in 2D
    """
    return np.array( 
        [ [ +np.cos( a ), -np.sin( a ) ],
          [ +np.sin( a ), +np.cos( a ) ] ] 
    )

def fit_ellipse( x, y ):
    """
    main fit from the original publication:
    http://nicky.vanforeest.com/misc/fitEllipse/fitEllipse.html
    """
    x = x[ :, np.newaxis ]
    y = y[ :, np.newaxis ]
    D =  np.hstack( ( x * x, x * y, y * y, x, y, np.ones_like( x ) ) )
    S = np.dot( D.T, D )
    C = np.zeros( [ 6, 6 ] )
    C[ 0, 2 ] = +2 
    C[ 2, 0 ] = +2
    C[ 1, 1 ] = -1
    E, V =  np.linalg.eig( np.dot( np.linalg.inv( S ), C ) )
    n = np.argmax( np.abs( E ) )
    a = V[ :, n ]
    return a

def ell_parameters( a ):
    """
    New function substituting the original 3 functions for 
    axis, centre and angle.
    We start by noting that the linear term is due to an offset. 
    Getting rid of it is equivalent to find the offset. 
    Starting with the Eq.
    xT A x + bT x + c = 0 and transforming x -> x - t 
    we get a new linear term. By demanding that this term vanishes
    we get the Eq.
    b = (AT + A ) t. 
    Hence, an easy way to write down how to get t
    """
    A = np.array( [ [ a[0], a[1]/2. ], [ a[1]/2., a[2] ] ] )
    b = np.array( [ a[3], a[4] ] )
    t = np.dot( np.linalg.inv( np.transpose( A ) + A ), b )
    """
    the transformation changes the constant term, which we need
    for proper scaling
    """
    c = a[5]
    cnew =  c - np.dot( t, b ) + np.dot( t, np.dot( A, t ) )
    Anew = A / (-cnew)
    # ~cnew = cnew / (-cnew) ### debug only
    """
    now it is in the form xT A x - 1 = 0
    and we know that A is a rotation of the matrix 
        ( 1 / a²   0 )
    B = (            )
        ( 0   1 / b² )
    where a and b are the semi axes of the ellipse
    it is hence A = ST B S
    We note that rotation does not change the eigenvalues, which are 
    the diagonal elements of matrix B. Moreover, we note that 
    the matrix of eigenvectors rotates B into A
    """
    E, V = np.linalg.eig( Anew )
    """
    so we have
    B = VT A V
    and consequently
    A = V B VT
    where V is of a form as given by the function rot() from above
    """
    # ~B = np.dot( np.transpose(V), np.dot( Anew, V ) ) ### debug only
    phi = np.arccos( V[ 0, 0 ] )
    """
    checking the sin for changes in sign to detect angles above 180°
    """
    if V[ 0, 1 ] < 0: 
        phi = 2 * np.pi - phi
    ### cw vs ccw and periodicity of pi
    phi = -phi % np.pi
    return np.sqrt( 1. / E ), phi * RAD, -t
    """
    That's it. One might put some additional work/thought in the 180° 
    and cw vs ccw thing, as it is a bit messy. 
    """

"""
creating some test data
"""
xl = np.linspace(-3,2.5, 10)
yl = np.fromiter( (2.0 * np.sqrt( 1 - ( x / 3. )**2 ) for x in xl ), np.float )
xl = np.append(xl,-xl)
yl = np.append(yl,-yl)
R = rot( -103.01 * DEGREE ) ### check different angles
# ~R = rot( 153 * DEGREE ) results in singular matrix !!!...strange
xyrot = np.array( [ np.dot(R, [ x, y ] )for x, y  in zip( xl, yl ) ] )
xl = xyrot[:,0] + 7
yl = xyrot[:,1] + 16.4

"""
fitting
"""
avec = fit_ellipse( xl, yl )
(a, b), phi, t = ell_parameters( avec )

ell = Ellipse(
    t, 2 * a, 2 * b, phi, 
    facecolor=( 1, 0, 0, 0.2 ), edgecolor=( 0, 0, 0, 0.5 )
)

"""
plotting
"""
fig = plt.figure()
ax = fig.add_subplot( 1, 1, 1 )
ax.add_patch( ell )
ax.scatter( xl ,yl )
plt.show()

我觉得这段代码也不应该出现90°的问题。如果删除所有注释,它非常紧凑并且(从数学的角度来看)可读。

备注

我在 inv( S ) 中遇到问题。这个矩阵变得奇异。一种解决方法可能是:将所有数据旋转一个小角度,然后将计算出的 t 旋转回来。同样从计算的角度中减去角度 phi.