From 6427dc2a77c5136fac997faf2562d55e976e5525 Mon Sep 17 00:00:00 2001
From: Dominic Canare <dom@dominiccanare.com>
Date: Mon, 30 Sep 2024 08:48:37 -0400
Subject: [PATCH] Adds brightness/contrast controls to surface detector

---
 .../surface_tracker/offline_utils.py          |  2 ++
 .../surface_marker_detector.py                |  9 +++++-
 .../surface_tracker/surface_tracker.py        | 30 +++++++++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/pupil_src/shared_modules/surface_tracker/offline_utils.py b/pupil_src/shared_modules/surface_tracker/offline_utils.py
index 93529cc1..75e67250 100644
--- a/pupil_src/shared_modules/surface_tracker/offline_utils.py
+++ b/pupil_src/shared_modules/surface_tracker/offline_utils.py
@@ -20,6 +20,8 @@ def from_detector(cls, detector: MarkerDetectorController):
             apriltag_nthreads=detector._apriltag_nthreads,
             apriltag_quad_decimate=detector.apriltag_quad_decimate,
             apriltag_decode_sharpening=detector.apriltag_decode_sharpening,
+            brightness=detector.brightness,
+            contrast=detector.contrast,
         )
 
     def __call__(self, frame):
diff --git a/pupil_src/shared_modules/surface_tracker/surface_marker_detector.py b/pupil_src/shared_modules/surface_tracker/surface_marker_detector.py
index 0c013b1d..27f1c215 100644
--- a/pupil_src/shared_modules/surface_tracker/surface_marker_detector.py
+++ b/pupil_src/shared_modules/surface_tracker/surface_marker_detector.py
@@ -13,6 +13,7 @@
 import logging
 import typing as T
 
+import cv2
 import pupil_apriltags
 
 from .surface_marker import Surface_Marker, Surface_Marker_Type
@@ -181,12 +182,17 @@ def __init__(
         apriltag_nthreads: int = 2,
         apriltag_quad_decimate: float = ...,
         apriltag_decode_sharpening: float = ...,
+        brightness: int = 0,
+        contrast: float = 1.0,
     ):
         self._marker_detector_mode = marker_detector_mode
         self._apriltag_nthreads = apriltag_nthreads
         self._apriltag_quad_decimate = apriltag_quad_decimate
         self._apriltag_decode_sharpening = apriltag_decode_sharpening
 
+        self.brightness = brightness
+        self.contrast = contrast
+
         self.init_detector()
 
     def init_detector(self):
@@ -237,6 +243,7 @@ def apriltag_decode_sharpening(self, value: float):
     def detect_markers_iter(
         self, gray_img, frame_index: int
     ) -> T.Iterable[Surface_Marker]:
+        adjusted_img = cv2.convertScaleAbs(gray_img, alpha=self.contrast, beta=self.brightness)
         yield from self.__detector.detect_markers_iter(
-            gray_img=gray_img, frame_index=frame_index
+            gray_img=adjusted_img, frame_index=frame_index
         )
diff --git a/pupil_src/shared_modules/surface_tracker/surface_tracker.py b/pupil_src/shared_modules/surface_tracker/surface_tracker.py
index 833d6d52..c3b366c1 100644
--- a/pupil_src/shared_modules/surface_tracker/surface_tracker.py
+++ b/pupil_src/shared_modules/surface_tracker/surface_tracker.py
@@ -61,6 +61,8 @@ def __init__(
         use_online_detection: bool = False,
         use_high_res=APRILTAG_HIGH_RES_ON,
         sharpen=APRILTAG_SHARPENING_ON,
+        brightness=0,
+        contrast=1.0,
     ):
         super().__init__(g_pool)
 
@@ -82,6 +84,8 @@ def __init__(
             marker_detector_mode=marker_detector_mode,
             apriltag_quad_decimate=use_high_res,
             apriltag_decode_sharpening=sharpen,
+            brightness=brightness,
+            contrast=contrast,
         )
         self._surface_file_store = Surface_File_Store(parent_dir=self._save_dir)
 
@@ -136,6 +140,7 @@ def init_ui(self):
             hotkey=Hotkey.SURFACE_TRACKER_ADD_SURFACE_CAPTURE_AND_PLAYER_HOTKEY(),
         )
         self.g_pool.quickbar.append(self.add_button)
+
         self._update_ui()
 
     def _update_ui(self):
@@ -191,8 +196,30 @@ def set_marker_detector_mode(value):
                     {"subject": "surface_tracker.marker_detection_params_changed"}
                 )
 
+        def set_brightness(value):
+            if self.marker_detector.brightness != value:
+                self.marker_detector.brightness = value
+                self.notify_all({
+                    "subject": "surface_tracker.marker_detection_params_changed",
+                    "delay": 2.0,
+                })
+
+        def set_contrast(value):
+            if self.marker_detector.contrast != value:
+                self.marker_detector.contrast = value
+                self.notify_all({
+                    "subject": "surface_tracker.marker_detection_params_changed",
+                    "delay": 2.0,
+                })
+
         menu = ui.Growing_Menu("Marker Detection Parameters")
         menu.collapsed = True
+        menu.append(
+            ui.Slider("brightness", self.marker_detector, min=0, max=100, label="Brightness", setter=set_brightness)
+        )
+        menu.append(
+            ui.Slider("contrast", self.marker_detector, min=1.0, max=3.0, label="Contrast", setter=set_contrast)
+        )
         menu.append(
             ui.Selector(
                 "marker_detector_mode",
@@ -205,6 +232,7 @@ def set_marker_detector_mode(value):
         )
 
         menu.append(self._apriltag_marker_param_menu())
+
         return menu
 
     def _apriltag_marker_param_menu(self):
@@ -498,6 +526,8 @@ def get_init_dict(self):
             "marker_detector_mode": marker_detector_mode,
             "use_high_res": self.marker_detector.apriltag_quad_decimate,
             "sharpen": self.marker_detector.apriltag_decode_sharpening,
+            "brightness": self.marker_detector.brightness,
+            "contrast": self.marker_detector.contrast,
         }
 
     def save_surface_definitions_to_file(self):