Ready I think

iris/Dockerfile

@@ -23,6 +23,9 @@ ENV IRIS_ENV="DEV"
 RUN uv pip install notebook git+https://github.com/worldcoin/open-iris.git@v1.11.0
 RUN patchelf --clear-execstack /src/.venv/lib/python3.11/site-packages/onnxruntime/capi/*.so
 
+# Run initialize to download the model
+RUN python ./initialize.py
+
 # Launch Jupyter Notebook
 EXPOSE 8888
 CMD ["python", "-m", "jupyter", "notebook", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--IdentityProvider.token=''"]

iris/src/initialize.py

@@ -0,0 +1,3 @@
+import iris
+
+iris_pipeline = iris.IRISPipeline()

iris/src/pipeline.py

@@ -3,30 +3,45 @@ pipeline_conf = {
     "pipeline": [
         {
             "name": "segmentation",
-            "algorithm": {"class_name": "iris.MultilabelSegmentation.create_from_hugging_face", "params": {}},
+            "algorithm": {
+                "class_name": "iris.MultilabelSegmentation.create_from_hugging_face",
+                "params": {},
+            },
             "inputs": [{"name": "image", "source_node": "input"}],
             "callbacks": None,
         },
         {
             "name": "segmentation_binarization",
-            "algorithm": {"class_name": "iris.MultilabelSegmentationBinarization", "params": {
-                "eyeball_threshold": 0.1,
-                "iris_threshold": 0.1,
-                "pupil_threshold": 0.1,
-                "eyelashes_threshold": 0.1
-            }},
+            "algorithm": {
+                "class_name": "iris.MultilabelSegmentationBinarization",
+                "params": {
+                    "eyeball_threshold": 0.5,
+                    "iris_threshold": 0.5,
+                    "pupil_threshold": 0.5,
+                    "eyelashes_threshold": 0.5,
+                },
+            },
             "inputs": [{"name": "segmentation_map", "source_node": "segmentation"}],
             "callbacks": None,
         },
         {
             "name": "vectorization",
             "algorithm": {"class_name": "iris.ContouringAlgorithm", "params": {}},
-            "inputs": [{"name": "geometry_mask", "source_node": "segmentation_binarization", "index": 0}],
+            "inputs": [
+                {
+                    "name": "geometry_mask",
+                    "source_node": "segmentation_binarization",
+                    "index": 0,
+                }
+            ],
             "callbacks": None,
         },
         {
             "name": "specular_reflection_detection",
-            "algorithm": {"class_name": "iris.SpecularReflectionDetection", "params": {}},
+            "algorithm": {
+                "class_name": "iris.SpecularReflectionDetection",
+                "params": {},
+            },
             "inputs": [{"name": "ir_image", "source_node": "input"}],
             "callbacks": None,
         },
@@ -38,18 +53,26 @@ pipeline_conf = {
         },
         {
             "name": "distance_filter",
-            "algorithm": {"class_name": "iris.ContourPointNoiseEyeballDistanceFilter", "params": {}},
+            "algorithm": {
+                "class_name": "iris.ContourPointNoiseEyeballDistanceFilter",
+                "params": {},
+            },
             "inputs": [
                 {"name": "polygons", "source_node": "interpolation"},
-                {"name": "geometry_mask", "source_node": "segmentation_binarization", "index": 1},
+                {
+                    "name": "geometry_mask",
+                    "source_node": "segmentation_binarization",
+                    "index": 1,
+                },
             ],
             "callbacks": None,
         },
         {
             "name": "eye_orientation",
-            "algorithm": {"class_name": "iris.MomentOfArea", "params": {
-                "eccentricity_threshold": 0.05
-            }},
+            "algorithm": {
+                "class_name": "iris.MomentOfArea",
+                "params": {"eccentricity_threshold": 0.05},
+            },
             "inputs": [{"name": "geometries", "source_node": "distance_filter"}],
             "callbacks": None,
         },
@@ -73,8 +96,14 @@ pipeline_conf = {
             "algorithm": {
                 "class_name": "iris.FusionExtrapolation",
                 "params": {
-                    "circle_extrapolation": {"class_name": "iris.LinearExtrapolation", "params": {"dphi": 0.703125}},
-                    "ellipse_fit": {"class_name": "iris.LSQEllipseFitWithRefinement", "params": {"dphi": 0.703125}},
+                    "circle_extrapolation": {
+                        "class_name": "iris.LinearExtrapolation",
+                        "params": {"dphi": 0.703125},
+                    },
+                    "ellipse_fit": {
+                        "class_name": "iris.LSQEllipseFitWithRefinement",
+                        "params": {"dphi": 0.703125},
+                    },
                     "algorithm_switch_std_threshold": 3.5,
                 },
             },
@@ -86,7 +115,10 @@ pipeline_conf = {
         },
         {
             "name": "pupil_to_iris_property_estimation",
-            "algorithm": {"class_name": "iris.PupilIrisPropertyCalculator", "params": {}},
+            "algorithm": {
+                "class_name": "iris.PupilIrisPropertyCalculator",
+                "params": {},
+            },
             "inputs": [
                 {"name": "geometries", "source_node": "geometry_estimation"},
                 {"name": "eye_centers", "source_node": "eye_center_estimation"},
@@ -95,15 +127,25 @@ pipeline_conf = {
         },
         {
             "name": "offgaze_estimation",
-            "algorithm": {"class_name": "iris.EccentricityOffgazeEstimation", "params": {}},
+            "algorithm": {
+                "class_name": "iris.EccentricityOffgazeEstimation",
+                "params": {},
+            },
             "inputs": [{"name": "geometries", "source_node": "geometry_estimation"}],
             "callbacks": None,
         },
         {
             "name": "occlusion90_calculator",
-            "algorithm": {"class_name": "iris.OcclusionCalculator", "params": {"quantile_angle": 90.0}},
+            "algorithm": {
+                "class_name": "iris.OcclusionCalculator",
+                "params": {"quantile_angle": 90.0},
+            },
             "inputs": [
-                {"name": "noise_mask", "source_node": "segmentation_binarization", "index": 1},
+                {
+                    "name": "noise_mask",
+                    "source_node": "segmentation_binarization",
+                    "index": 1,
+                },
                 {"name": "extrapolated_polygons", "source_node": "geometry_estimation"},
                 {"name": "eye_orientation", "source_node": "eye_orientation"},
                 {"name": "eye_centers", "source_node": "eye_center_estimation"},
@@ -112,9 +154,16 @@ pipeline_conf = {
         },
         {
             "name": "occlusion30_calculator",
-            "algorithm": {"class_name": "iris.OcclusionCalculator", "params": {"quantile_angle": 30.0}},
+            "algorithm": {
+                "class_name": "iris.OcclusionCalculator",
+                "params": {"quantile_angle": 30.0},
+            },
             "inputs": [
-                {"name": "noise_mask", "source_node": "segmentation_binarization", "index": 1},
+                {
+                    "name": "noise_mask",
+                    "source_node": "segmentation_binarization",
+                    "index": 1,
+                },
                 {"name": "extrapolated_polygons", "source_node": "geometry_estimation"},
                 {"name": "eye_orientation", "source_node": "eye_orientation"},
                 {"name": "eye_centers", "source_node": "eye_center_estimation"},
@@ -178,12 +227,20 @@ pipeline_conf = {
                         },
                     ],
                     "probe_schemas": [
-                        {"class_name": "iris.RegularProbeSchema", "params": {"n_rows": 16, "n_cols": 256}},
-                        {"class_name": "iris.RegularProbeSchema", "params": {"n_rows": 16, "n_cols": 256}},
+                        {
+                            "class_name": "iris.RegularProbeSchema",
+                            "params": {"n_rows": 16, "n_cols": 256},
+                        },
+                        {
+                            "class_name": "iris.RegularProbeSchema",
+                            "params": {"n_rows": 16, "n_cols": 256},
+                        },
                     ],
                 },
             },
-            "inputs": [{"name": "normalization_output", "source_node": "normalization"}],
+            "inputs": [
+                {"name": "normalization_output", "source_node": "normalization"}
+            ],
             "callbacks": None,
         },
         {

iris/src/process_eye.py

@@ -1,18 +1,24 @@
+import copy
 import pathlib
 
 import cv2
+import matplotlib as mpl
 import matplotlib.pyplot as plt
+import numpy as np
+from pipeline import pipeline_conf
 
 import iris
 
-from pipeline import pipeline_conf
-
 
 def process(path: str) -> None:
     path = pathlib.Path(path)
 
-    iris_pipeline = iris.IRISPipeline(env=iris.IRISPipeline.DEBUGGING_ENVIRONMENT, config=pipeline_conf)
+    print("Loading pipeline...")
+    iris_pipeline = iris.IRISPipeline(
+        env=iris.IRISPipeline.DEBUGGING_ENVIRONMENT, config=pipeline_conf
+    )
     iris_visualizer = iris.visualisation.IRISVisualizer()
+    print("Pipeline loaded")
 
     outputs = []
     for image_path in (path / "input").iterdir():
@@ -20,38 +26,61 @@ def process(path: str) -> None:
             continue
 
         print(f"Processing {image_path}...")
+        try:
+            eye_side = "right" if "left" not in image_path.stem else "left"
 
-        eye_side = "right" if "left" not in image_path.stem else "left"
+            image_pixels = cv2.imread(str(image_path), cv2.IMREAD_GRAYSCALE)
+            img = iris.IRImage(
+                img_data=image_pixels,
+                eye_side=eye_side,
+            )
+            output = iris_pipeline(
+                iris.IRImage(
+                    img_data=image_pixels, image_id="image_id", eye_side="right"
+                )
+            )
 
-        image_pixels = cv2.imread(str(image_path), cv2.IMREAD_GRAYSCALE)
-        img = iris.IRImage(
-            img_data=image_pixels,
-            eye_side=eye_side,
-        )
-        output = iris_pipeline(iris.IRImage(img_data=image_pixels, image_id="image_id", eye_side="right"))
+            iris_visualizer.plot_segmentation_map(
+                ir_image=img,
+                segmap=iris_pipeline.call_trace["segmentation"],
+            )
+            (path / "segmentation").mkdir(parents=True, exist_ok=True)
+            plt.savefig(path / "segmentation" / f"{image_path.stem}.png")
+            plt.close()
 
-        iris_visualizer.plot_segmentation_map(
-            ir_image=img,
-            segmap=iris_pipeline.call_trace["segmentation"],
-        )
-        (path / "segmentation").mkdir(parents=True, exist_ok=True)
-        plt.savefig(path / "segmentation" / f"{image_path.stem}.png")
+            iris_visualizer.plot_all_geometry(
+                ir_image=img,
+                geometry_polygons=iris_pipeline.call_trace["geometry_estimation"],
+                eye_orientation=iris_pipeline.call_trace["eye_orientation"],
+                eye_center=iris_pipeline.call_trace["eye_center_estimation"],
+            )
+            (path / "geometry").mkdir(parents=True, exist_ok=True)
+            plt.savefig(path / "geometry" / f"{image_path.stem}.png")
+            plt.close()
 
-        iris_visualizer.plot_all_geometry(
-            ir_image=img,
-            geometry_polygons=iris_pipeline.call_trace["geometry_estimation"],
-            eye_orientation=iris_pipeline.call_trace["eye_orientation"],
-            eye_center=iris_pipeline.call_trace["eye_center_estimation"],
-        )
-        (path / "geometry").mkdir(parents=True, exist_ok=True)
-        plt.savefig(path / "geometry" / f"{image_path.stem}.png")
+            (path / "normalized").mkdir(parents=True, exist_ok=True)
+            data = iris_pipeline.call_trace["normalization"].normalized_image
+            plt.imsave(
+                path / "normalized" / f"{image_path.stem}.png", data, cmap="Greys_r"
+            )
 
-        (path / "normalized").mkdir(parents=True, exist_ok=True)
-        data = iris_pipeline.call_trace["normalization"].normalized_image
-        plt.imsave(path / "normalized" / f"{image_path.stem}.png", data, cmap='Greys_r')
+            cm = copy.copy(mpl.colormaps["binary"])
+            cm.set_bad("red")
 
-        (path / "template").mkdir(parents=True, exist_ok=True)
-        data = output["iris_template"].iris_codes[0][:, :, 0]
-        plt.imsave(path / "template" / f"{image_path.stem}.png", data, cmap='binary')
+            data = output["iris_template"].iris_codes[0][:, :, 0]
+            mask = np.invert(output["iris_template"].mask_codes[0][:, :, 0])
+            data = np.ma.masked_array(data, mask=mask)
 
-        outputs.append((image_path, img, output))
+            (path / "template").mkdir(parents=True, exist_ok=True)
+            plt.imsave(path / "template" / f"{image_path.stem}.png", data, cmap=cm)
+
+            outputs.append((image_path, img, output))
+
+            print(f"Finished processing: {image_path}")
+
+        except:
+            print(
+                f"Failed to process: {image_path}. Most likely not centered, out of focus, or eye not opened enough"
+            )
+
+    print("Done")