diff --git a/colabseg/parametrization.py b/colabseg/parametrization.py index 7cacc65..4f7d921 100644 --- a/colabseg/parametrization.py +++ b/colabseg/parametrization.py @@ -192,51 +192,41 @@ def fit(cls, positions) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: NotImplementedError If the points are not 3D. """ - positions = np.asarray(positions) + positions = np.asarray(positions, dtype=np.float64) if positions.shape[1] != 3 or len(positions.shape) != 2: raise NotImplementedError( "Only three-dimensional point clouds are supported." ) - positions_mean = positions.mean(axis=0) - positions -= positions_mean - - D = np.array( - [ - np.square(positions[:, 0]), - np.square(positions[:, 1]), - np.square(positions[:, 2]), - 2 * positions[:, 0] * positions[:, 1], - 2 * positions[:, 0] * positions[:, 2], - 2 * positions[:, 1] * positions[:, 2], - 2 * positions[:, 0], - 2 * positions[:, 1], - 2 * positions[:, 2], - np.ones_like(positions[:, 0]), - ] - ).T - v, *_ = np.linalg.lstsq(D, np.ones_like(positions[:, 0]), rcond=None) - - A = np.array( - [ - [v[0], v[3], v[4], v[6]], - [v[3], v[1], v[5], v[7]], - [v[4], v[5], v[2], v[8]], - [v[6], v[7], v[8], -1], - ] - ) + def ellipsoid_loss(radii, data_points, center, orientations): + transformed_points = np.dot(data_points - center, orientations) + + normalized_points = transformed_points / radii + + distances = np.sum(normalized_points**2, axis=1) - 1 + + loss = np.sum(distances**2) + return loss - center = np.linalg.solve(-A[:3, :3], v[6:9]) + center = positions.mean(axis=0) + positions_centered = positions - center - T = np.eye(4) - T[3, :3] = center - R = T.dot(A).dot(T.T) + cov_mat = np.cov(positions_centered, rowvar=False) + evals, evecs = np.linalg.eigh(cov_mat) - evals, evecs = np.linalg.eig(R[:3, :3] / -R[3, 3]) - radii = np.sqrt(1.0 / np.abs(evals)) - evecs = np.dot(evecs, np.diag(np.sign(evals))) + sort_indices = np.argsort(evals)[::-1] + evals = evals[sort_indices] + evecs = evecs[:, sort_indices] - center = np.add(positions_mean, center) + initial_radii = 2 * np.sqrt(evals) + + result = optimize.minimize( + ellipsoid_loss, + initial_radii, + args=(positions, center, evecs), + method="Nelder-Mead", + ) + radii = result.x return cls(radii=radii, center=center, orientations=evecs) @@ -271,18 +261,18 @@ def sample( orientations = self.orientations if orientations is None else orientations phi = np.random.uniform(0, 2 * np.pi, n_samples) - costheta = np.random.uniform(-1, 1, n_samples) - theta = np.arccos(costheta) + theta = np.random.uniform(0, np.pi, n_samples) + x = np.sin(theta) * np.cos(phi) y = np.sin(theta) * np.sin(phi) z = np.cos(theta) - points = np.array([x, y, z]).T * radii + samples = np.vstack((x, y, z)).T + samples = samples * radii + samples = samples.dot(orientations.T) + samples += center - points = np.dot(points, orientations.T) - np.add(points, center, out=points) - - return points + return samples class Cylinder(Parametrization): @@ -336,12 +326,26 @@ def fit( ------ ValueError If th number of initial parameters is not equal to five. + NotImplementedError + If the points are not 3D. References ---------- .. [1] Ting On Chan https://stackoverflow.com/posts/44164662/revisions """ + positions = np.asarray(positions, dtype=np.float64) + if positions.shape[1] != 3 or len(positions.shape) != 2: + raise NotImplementedError( + "Only three-dimensional point clouds are supported." + ) + + if initial_parameters is None: + initial_parameters = np.zeros(5) + + if len(initial_parameters) != 5: + raise ValueError(f"Expected 5 parameters, got {len(initial_parameters)}") + def fit_function(p: np.ndarray, positions: np.ndarray) -> np.ndarray: term1 = -np.cos(p[3]) * (p[0] - positions[:, 0]) term2 = positions[:, 2] * np.cos(p[2]) * np.sin(p[3]) @@ -353,12 +357,6 @@ def fit_function(p: np.ndarray, positions: np.ndarray) -> np.ndarray: def error_function(p: np.ndarray, positions: np.ndarray) -> np.ndarray: return fit_function(p, positions) - p[4] ** 2 - if initial_parameters is None: - initial_parameters = np.zeros(5) - - if len(initial_parameters) != 5: - raise ValueError(f"Expected 5 parameters, got {len(initial_parameters)}") - parameters, _ = optimize.leastsq( error_function, initial_parameters, args=(positions,), maxfev=10000 )