Merge pull request #221 from worldcoin/yichen1/aib-1661-rewrite-eye-center-calculation

Yichen1/aib 1661 rewrite eye center calculation

Author: ycbiometrics

src/iris/nodes/eye_properties_estimation/circle_fit_for_eye_center_method.py

@@ -0,0 +1,145 @@
+from typing import Tuple
+
+import cv2
+import numpy as np
+from pydantic import Field
+
+from iris.io.class_configs import Algorithm
+from iris.io.dataclasses import EyeCenters, GeometryPolygons
+from iris.io.errors import EyeCentersEstimationError
+
+
+class CircleFitEyeCenterMethod(Algorithm):
+    """Estimate pupil and iris centers using a robust circle-fitting approach.
+
+    This algorithm estimates the center of a pupil or iris polygon using:
+
+    1. an initial least-squares circle fit to all polygon points,
+    2. trimming of outlier points based on radial residuals using a dynamic threshold,
+    3. a second circle fit on the retained inliers.
+
+    LIMITATIONS:
+    This method assumes that the pupil and iris contours are reasonably well
+    approximated by circles in image space. It is therefore most appropriate
+    when off-gaze and strong perspective distortion have already been filtered out.
+    """
+
+    class Parameters(Algorithm.Parameters):
+        """Default parameters for circle-fit eye center algorithm."""
+
+        mad_scale: float = Field(..., gt=0.0)
+
+    __parameters_type__ = Parameters
+
+    def __init__(
+        self,
+        mad_scale: float = 3.0,
+    ) -> None:
+        """Assign parameters.
+
+        Args:
+            mad_scale (float, optional): Scale factor used in the dynamic inlier
+                threshold: median(residuals) + mad_scale * MAD(residuals).
+                Defaults to 3.0.
+        """
+        super().__init__(mad_scale=mad_scale)
+
+    def run(self, geometries: GeometryPolygons) -> EyeCenters:
+        """Estimate pupil and iris centers.
+
+        Args:
+            geometries (GeometryPolygons): Pupil and iris geometry polygons.
+
+        Returns:
+            EyeCenters: Estimated pupil and iris center coordinates.
+        """
+        pupil_center_x, pupil_center_y = self._calculate_circle_fit_center(geometries.pupil_array.astype(np.float64))
+        iris_center_x, iris_center_y = self._calculate_circle_fit_center(geometries.iris_array.astype(np.float64))
+
+        return EyeCenters(
+            pupil_x=pupil_center_x,
+            pupil_y=pupil_center_y,
+            iris_x=iris_center_x,
+            iris_y=iris_center_y,
+        )
+
+    def _calculate_circle_fit_center(self, polygon: np.ndarray) -> Tuple[float, float]:
+        """Estimate the center of a polygon using fit -> trim outliers -> refit.
+
+        Args:
+            polygon (np.ndarray): Polygon points of shape (N, 2) representing a
+                contour that is approximately circular.
+
+        Raises:
+            EyeCentersEstimationError: Raised if the polygon is invalid or if a
+                valid circle cannot be fit.
+
+        Returns:
+            Tuple[float, float]: Estimated center coordinates (x, y).
+        """
+        if polygon.ndim != 2 or polygon.shape[1] != 2 or polygon.shape[0] < 3:
+            raise EyeCentersEstimationError("Polygon must have shape (N, 2) with at least 3 points")
+
+        pts = polygon.astype(np.float64, copy=False)
+
+        # Initial least-squares circle fit
+        x = pts[:, 0]
+        y = pts[:, 1]
+
+        A = np.column_stack([x, y, np.ones_like(x)]).astype(np.float64)
+        b = (-(x * x + y * y)).reshape(-1, 1).astype(np.float64)
+
+        ok, coeffs = cv2.solve(A, b, flags=cv2.DECOMP_SVD)
+        if not ok:
+            raise EyeCentersEstimationError("Circle fit failed")
+
+        a, b_, c = coeffs.ravel()
+
+        cx = -a / 2.0
+        cy = -b_ / 2.0
+
+        r_sq = cx * cx + cy * cy - c
+        if r_sq <= 0:
+            raise EyeCentersEstimationError("Failed to fit a valid circle to the polygon")
+
+        r = np.sqrt(r_sq)
+        center = np.array([cx, cy], dtype=np.float64)
+
+        # Trim outliers using dynamic radial residual threshold
+        d = np.linalg.norm(pts - center, axis=1)
+        residuals = np.abs(d - r)
+
+        median_residual = np.median(residuals)
+        mad = np.median(np.abs(residuals - median_residual))
+
+        if mad == 0:
+            keep_mask = residuals <= median_residual
+        else:
+            threshold = median_residual + self.params.mad_scale * mad
+            keep_mask = residuals <= threshold
+
+        inliers = pts[keep_mask]
+
+        # Ensure enough points remain to refit
+        if len(inliers) < 3:
+            keep_idx = np.argsort(residuals)[:3]
+            inliers = pts[keep_idx]
+
+        # Refit circle on inliers
+        x = inliers[:, 0]
+        y = inliers[:, 1]
+
+        A = np.column_stack([x, y, np.ones_like(x)])
+        b = -(x * x + y * y)
+
+        coeffs, _, _, _ = np.linalg.lstsq(A, b, rcond=None)
+
+        a, b_, c = coeffs
+        cx = -a / 2.0
+        cy = -b_ / 2.0
+
+        r_sq = cx * cx + cy * cy - c
+        if r_sq <= 0:
+            raise EyeCentersEstimationError("Failed to refit a valid circle to the inlier polygon points")
+
+        return float(cx), float(cy)

…ectors_method/e2e_expected_result.pickle …sector_method_e2e_expected_result.pickletests/e2e_tests/nodes/eye_properties_estimation/mocks/bisectors_method/e2e_expected_result.pickle renamed to tests/e2e_tests/nodes/eye_properties_estimation/mocks/eye_center_method/bisector_method_e2e_expected_result.pickle tests/e2e_tests/nodes/eye_properties_estimation/mocks/bisectors_method/e2e_expected_result.pickle renamed to tests/e2e_tests/nodes/eye_properties_estimation/mocks/eye_center_method/bisector_method_e2e_expected_result.pickle

@@ -9,7 +9,7 @@
 
 
 def load_mock_pickle(name: str) -> Any:
-    testdir = os.path.join(os.path.dirname(__file__), "mocks", "bisectors_method")
+    testdir = os.path.join(os.path.dirname(__file__), "mocks", "eye_center_method")
 
     mock_path = os.path.join(testdir, f"{name}.pickle")
 
@@ -23,7 +23,7 @@ def algorithm() -> BisectorsMethod:
 
 def test_e2e_bisectors_method_algorithm(algorithm: BisectorsMethod) -> None:
     mock_polygons = load_mock_pickle(name="geometry_polygons")
-    expected_result = load_mock_pickle(name="e2e_expected_result")
+    expected_result = load_mock_pickle(name="bisector_method_e2e_expected_result")
 
     result = algorithm(geometries=mock_polygons)
 

tests/e2e_tests/nodes/eye_properties_estimation/mocks/eye_center_method/circle_fit_method_e2e_expected_result.pickle

@@ -0,0 +1,33 @@
+import math
+import os
+import pickle
+from typing import Any
+
+import pytest
+
+from iris.nodes.eye_properties_estimation.circle_fit_for_eye_center_method import CircleFitEyeCenterMethod
+
+
+def load_mock_pickle(name: str) -> Any:
+    testdir = os.path.join(os.path.dirname(__file__), "mocks", "eye_center_method")
+
+    mock_path = os.path.join(testdir, f"{name}.pickle")
+
+    return pickle.load(open(mock_path, "rb"))
+
+
+@pytest.fixture
+def algorithm() -> CircleFitEyeCenterMethod:
+    return CircleFitEyeCenterMethod(mad_scale=3.0)
+
+
+def test_e2e_bisectors_method_algorithm(algorithm: CircleFitEyeCenterMethod) -> None:
+    mock_polygons = load_mock_pickle(name="geometry_polygons")
+    expected_result = load_mock_pickle(name="circle_fit_method_e2e_expected_result")
+
+    result = algorithm(geometries=mock_polygons)
+
+    assert math.isclose(result.pupil_x, expected_result.pupil_x)
+    assert math.isclose(result.pupil_y, expected_result.pupil_y)
+    assert math.isclose(result.iris_x, expected_result.iris_x)
+    assert math.isclose(result.iris_y, expected_result.iris_y)

…isectors_method/geometry_polygons.pickle …e_center_method/geometry_polygons.pickletests/e2e_tests/nodes/eye_properties_estimation/mocks/bisectors_method/geometry_polygons.pickle renamed to tests/e2e_tests/nodes/eye_properties_estimation/mocks/eye_center_method/geometry_polygons.pickle tests/e2e_tests/nodes/eye_properties_estimation/mocks/bisectors_method/geometry_polygons.pickle renamed to tests/e2e_tests/nodes/eye_properties_estimation/mocks/eye_center_method/geometry_polygons.pickle

@@ -4,16 +4,17 @@
 from typing import Any
 
 from iris.nodes.eye_properties_estimation.bisectors_method import BisectorsMethod
+from iris.nodes.eye_properties_estimation.circle_fit_for_eye_center_method import CircleFitEyeCenterMethod
 from iris.nodes.eye_properties_estimation.pupil_iris_property_calculator import PupilIrisPropertyCalculator
 
 
 def load_mock_pickle(name: str) -> Any:
-    testdir = os.path.join(os.path.dirname(__file__), "mocks", "bisectors_method")
+    testdir = os.path.join(os.path.dirname(__file__), "mocks", "eye_center_method")
     mock_path = os.path.join(testdir, f"{name}.pickle")
     return pickle.load(open(mock_path, "rb"))
 
 
-def test_precomputed_pupil_iris_property() -> None:
+def test_precomputed_pupil_iris_property_bisector_method() -> None:
     mock_polygons = load_mock_pickle(name="geometry_polygons")
 
     eye_center_obj = BisectorsMethod()
@@ -24,3 +25,15 @@ def test_precomputed_pupil_iris_property() -> None:
 
     assert math.isclose(p2i_property.pupil_to_iris_diameter_ratio, 0.543019583685283)
     assert math.isclose(p2i_property.pupil_to_iris_center_dist_ratio, 0.032786957796171405)
+
+def test_precomputed_pupil_iris_property_circle_fit_method() -> None:
+    mock_polygons = load_mock_pickle(name="geometry_polygons")
+
+    eye_center_obj = CircleFitEyeCenterMethod()
+    eye_center = eye_center_obj(mock_polygons)
+
+    pupil_iris_property_obj = PupilIrisPropertyCalculator()
+    p2i_property = pupil_iris_property_obj(mock_polygons, eye_center)
+
+    assert math.isclose(p2i_property.pupil_to_iris_diameter_ratio, 0.543019583685283)
+    assert math.isclose(p2i_property.pupil_to_iris_center_dist_ratio, 0.03771351406478552)

tests/e2e_tests/nodes/eye_properties_estimation/test_e2e_bisectors_method.py

@@ -0,0 +1,34 @@
+import math
+
+import numpy as np
+import pytest
+
+from iris.io.dataclasses import GeometryPolygons
+from iris.nodes.eye_properties_estimation.circle_fit_for_eye_center_method import CircleFitEyeCenterMethod
+from tests.unit_tests.utils import generate_arc
+
+
+@pytest.fixture
+def algorithm() -> CircleFitEyeCenterMethod:
+    return CircleFitEyeCenterMethod(mad_scale=3.0)
+
+
+def test_estimation_on_mock_example(algorithm: CircleFitEyeCenterMethod) -> None:
+    pupil_radius = 25.0
+    iris_radius = 100.0
+    eyeball_radius = 400.0
+    pupil_center_x, pupil_center_y = 95.0, 145.0
+    iris_center_x, iris_center_y = 100.0, 155.0
+
+    mock_polygons = GeometryPolygons(
+        pupil_array=generate_arc(pupil_radius, pupil_center_x, pupil_center_y, 0.0, 2 * np.pi),
+        iris_array=generate_arc(iris_radius, iris_center_x, iris_center_y, 0.0, 2 * np.pi),
+        eyeball_array=generate_arc(eyeball_radius, iris_center_x, iris_center_y, 0.0, 2 * np.pi),
+    )
+
+    result = algorithm(mock_polygons)
+
+    assert math.isclose(result.pupil_x, pupil_center_x, rel_tol=0.1)
+    assert math.isclose(result.pupil_y, pupil_center_y, rel_tol=0.1)
+    assert math.isclose(result.iris_x, iris_center_x, rel_tol=0.1)
+    assert math.isclose(result.iris_y, iris_center_y, rel_tol=0.1)

tests/e2e_tests/nodes/eye_properties_estimation/test_e2e_circle_fit_for_eye_center_method.py

@@ -6,6 +6,7 @@
 
 from iris.io.dataclasses import GeometryPolygons
 from iris.nodes.eye_properties_estimation.bisectors_method import BisectorsMethod
+from iris.nodes.eye_properties_estimation.circle_fit_for_eye_center_method import CircleFitEyeCenterMethod
 from iris.nodes.eye_properties_estimation.pupil_iris_property_calculator import (
     PupilIrisPropertyCalculator,
     PupilIrisPropertyEstimationError,
@@ -133,3 +134,12 @@ def test_pupil_iris_property(
     p2i_property = pupil_iris_property_obj(mock_polygons, eye_center)
     assert math.isclose(p2i_property.pupil_to_iris_diameter_ratio, expected_diameter_ratio, rel_tol=1e-03)
     assert math.isclose(p2i_property.pupil_to_iris_center_dist_ratio, expected_center_dist_ratio, rel_tol=1e-03)
+
+    eye_center_obj = CircleFitEyeCenterMethod()
+    eye_center = eye_center_obj(mock_polygons)
+    pupil_iris_property_obj = PupilIrisPropertyCalculator(
+        min_pupil_diameter=min_pupil_diameter, min_iris_diameter=min_iris_diameter
+    )
+    p2i_property = pupil_iris_property_obj(mock_polygons, eye_center)
+    assert math.isclose(p2i_property.pupil_to_iris_diameter_ratio, expected_diameter_ratio, rel_tol=1e-03)
+    assert math.isclose(p2i_property.pupil_to_iris_center_dist_ratio, expected_center_dist_ratio, rel_tol=1e-03)