Source code for iris.nodes.geometry_estimation.lsq_ellipse_fit_with_refinement
from typing import List
import cv2
import numpy as np
from pydantic import Field
from iris.callbacks.callback_interface import Callback
from iris.io.class_configs import Algorithm
from iris.io.dataclasses import GeometryPolygons
[docs]
class LSQEllipseFitWithRefinement(Algorithm):
"""Algorithm that implements least square ellipse fit with iris polygon refinement by finding points to refine by computing euclidean distance.
Algorithm steps:
1) Use OpenCV's fitEllipse method to fit an ellipse to predicted iris and pupil polygons.
2) Refine predicted pupil polygons points to their original location to prevent location precision loss for those points which were predicted by semseg algorithm.
"""
[docs]
class Parameters(Algorithm.Parameters):
"""Parameters of least square ellipse fit extrapolation algorithm."""
dphi: float = Field(..., gt=0.0, lt=360.0)
__parameters_type__ = Parameters
def __init__(self, dphi: float = 1.0, callbacks: List[Callback] = []) -> None:
"""Assign parameters.
Args:
dphi (float, optional): Angle's delta. Defaults to 1.0.
callbacks (List[Callback], optional): List of callbacks. Defaults to [].
"""
super().__init__(dphi=dphi, callbacks=callbacks)
[docs]
def run(self, input_polygons: GeometryPolygons) -> GeometryPolygons:
"""Estimate extrapolated polygons with OpenCV's method fitEllipse.
Args:
input_polygons (GeometryPolygons): Smoothed polygons.
Returns:
GeometryPolygons: Extrapolated polygons.
"""
extrapolated_pupil = self._extrapolate(input_polygons.pupil_array)
extrapolated_iris = self._extrapolate(input_polygons.iris_array)
for point in input_polygons.pupil_array:
extrapolated_pupil[self._find_correspondence(point, extrapolated_pupil)] = point
return GeometryPolygons(
pupil_array=extrapolated_pupil, iris_array=extrapolated_iris, eyeball_array=input_polygons.eyeball_array
)
def _extrapolate(self, polygon_points: np.ndarray) -> np.ndarray:
"""Perform extrapolation for points in an array.
Args:
polygon_points (np.ndarray): Smoothed polygon ready for applying extrapolation algorithm on it.
Returns:
np.ndarray: Estimated extrapolated polygon.
"""
(x0, y0), (a, b), theta = cv2.fitEllipse(polygon_points)
extrapolated_polygon = LSQEllipseFitWithRefinement.parametric_ellipsis(
a / 2, b / 2, x0, y0, np.radians(theta), round(360 / self.params.dphi)
)
# Rotate such that 0 degree is parallel with x-axis and array is clockwise
roll_amount = round((-theta - 90) / self.params.dphi)
extrapolated_polygon = np.flip(np.roll(extrapolated_polygon, roll_amount, axis=0), axis=0)
return extrapolated_polygon
def _find_correspondence(self, src_point: np.ndarray, dst_points: np.ndarray) -> int:
"""Find correspondence with Euclidean distance.
Args:
src_point (np.ndarray): Source points.
dst_points (np.ndarray): Destination points.
Returns:
int: Source point index the closes one to the destination points.
"""
src_x, src_y = src_point
distance = (dst_points[:, 1] - src_y) ** 2 + (dst_points[:, 0] - src_x) ** 2
idx = np.where(distance == distance.min())[0]
return idx
[docs]
@staticmethod
def parametric_ellipsis(a: float, b: float, x0: float, y0: float, theta: float, nb_step: int = 100) -> np.ndarray:
"""Given the parameters of a general ellipsis, returns an array of points in this ellipsis.
Args:
a (float): Major axis length.
b (float): Minor axis length.
x0 (float): x offset.
y0 (float): y offset.
theta (float): rotation of the ellipsis.
nb_step (int): number of points in the ellipsis.
Returns:
np.ndarray: points within the ellipsis.
"""
t = np.linspace(0, 2 * np.pi, nb_step)
x_coords = x0 + b * np.cos(t) * np.sin(-theta) + a * np.sin(t) * np.cos(-theta)
y_coords = y0 + b * np.cos(t) * np.cos(-theta) - a * np.sin(t) * np.sin(-theta)
return np.array([x_coords, y_coords]).T