Created
July 19, 2024 12:42
-
-
Save zach2good/9ea1b5c34697112e71fcd3c53b010b0d to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import numpy as np | |
import matplotlib.pyplot as plt | |
def fit_ellipse(x, y): | |
""" | |
Fit the coefficients a,b,c,d,e,f, representing an ellipse described by | |
the formula F(x,y) = ax^2 + bxy + cy^2 + dx + ey + f = 0 to the provided | |
arrays of data points x=[x1, x2, ..., xn] and y=[y1, y2, ..., yn]. | |
Based on the algorithm of Halir and Flusser, "Numerically stable direct | |
least squares fitting of ellipses'. | |
""" | |
D1 = np.vstack([x**2, x*y, y**2]).T | |
D2 = np.vstack([x, y, np.ones(len(x))]).T | |
S1 = D1.T @ D1 | |
S2 = D1.T @ D2 | |
S3 = D2.T @ D2 | |
T = -np.linalg.inv(S3) @ S2.T | |
M = S1 + S2 @ T | |
C = np.array(((0, 0, 2), (0, -1, 0), (2, 0, 0)), dtype=float) | |
M = np.linalg.inv(C) @ M | |
_, eigvec = np.linalg.eig(M) | |
con = 4 * eigvec[0]* eigvec[2] - eigvec[1]**2 | |
ak = eigvec[:, np.nonzero(con > 0)[0]] | |
return np.concatenate((ak, T @ ak)).ravel() | |
def cart_to_pol(coeffs): | |
""" | |
Convert the cartesian conic coefficients, (a, b, c, d, e, f), to the | |
ellipse parameters, where F(x, y) = ax^2 + bxy + cy^2 + dx + ey + f = 0. | |
The returned parameters are x0, y0, ap, bp, theta, where (x0, y0) is the | |
ellipse centre; (ap, bp) are the semi-major and semi-minor axes, | |
respectively; and theta is the rotation of the semi- | |
major axis from the x-axis. | |
""" | |
a = coeffs[0] | |
b = coeffs[1] / 2 | |
c = coeffs[2] | |
d = coeffs[3] / 2 | |
e = coeffs[4] / 2 | |
f = coeffs[5] | |
den = b**2 - a*c | |
if den > 0: | |
raise ValueError('coeffs do not represent an ellipse: b^2 - 4ac must' | |
' be negative!') | |
# The location of the ellipse centre. | |
x0, y0 = (c*d - b*e) / den, (a*e - b*d) / den | |
num = 2 * (a*e**2 + c*d**2 + f*b**2 - 2*b*d*e - a*c*f) | |
fac = np.sqrt((a - c)**2 + 4*b**2) | |
# The semi-major and semi-minor axis lengths (these are not sorted). | |
ap = np.sqrt(num / den / (fac - a - c)) | |
bp = np.sqrt(num / den / (-fac - a - c)) | |
# Sort the semi-major and semi-minor axis lengths but keep track of | |
# the original relative magnitudes of width and height. | |
width_gt_height = True | |
if ap < bp: | |
width_gt_height = False | |
ap, bp = bp, ap | |
# The angle of anticlockwise rotation of the major-axis from x-axis. | |
if b == 0: | |
theta = 0 if a < c else np.pi/2 | |
else: | |
theta = np.arctan((2.*b) / (a - c)) / 2 | |
if a > c: | |
theta += np.pi/2 | |
if not width_gt_height: | |
# Ensure that theta is the angle to rotate to the semi-major axis. | |
theta += np.pi/2 | |
theta = theta % np.pi | |
# Return ellipse parameters in the desired format: cX, cY, a, b, theta | |
return x0, y0, ap, bp, theta | |
def get_ellipse_pts(params, npts=100, tmin=0, tmax=2*np.pi): | |
""" | |
Return npts points on the ellipse described by the params = x0, y0, ap, | |
bp, theta for values of the parametric variable t between tmin and tmax. | |
""" | |
x0, y0, ap, bp, theta = params | |
# A grid of the parametric variable, t. | |
t = np.linspace(tmin, tmax, npts) | |
x = x0 + ap * np.cos(t) * np.cos(theta) - bp * np.sin(t) * np.sin(theta) | |
y = y0 + ap * np.cos(t) * np.sin(theta) + bp * np.sin(t) * np.cos(theta) | |
return x, y | |
def fit_and_draw_ellipse_from_points(points): | |
x, y = points.T | |
coeffs = fit_ellipse(x, y) | |
print("Ellipse coefficients:") | |
print(f"a = {coeffs[0]}, b = {coeffs[1]}, c = {coeffs[2]}, d = {coeffs[3]}, e = {coeffs[4]}, f = {coeffs[5]}") | |
x0, y0, a, b, theta = cart_to_pol(coeffs) | |
print("Ellipse parameters:") | |
print(f"cX = {x0}, cY = {y0}, a = {a}, b = {b}, theta = {theta}") | |
plt.plot(x, y, 'x') | |
x, y = get_ellipse_pts((x0, y0, a, b, theta)) | |
plt.plot(x, y) | |
plt.gca().set_aspect('equal') | |
plt.show() | |
fit_and_draw_ellipse_from_points(np.array([ | |
[3.5, 2.1], | |
[4.2, 1.0], | |
[2.0, -0.5], | |
[0.5, 1.0], | |
[1.0, 3.0], | |
[3.0, 4.0], | |
[4.0, 2.5], | |
[4.5, 1.5], | |
[2.5, 0.0], | |
[1.0, 0.5], | |
])) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment