Skip to content

Commit

Permalink
Adds brightness/contrast controls to surface detector
Browse files Browse the repository at this point in the history
  • Loading branch information
domstoppable committed Sep 30, 2024
1 parent a34a873 commit 6427dc2
Show file tree
Hide file tree
Showing 3 changed files with 40 additions and 1 deletion.
2 changes: 2 additions & 0 deletions pupil_src/shared_modules/surface_tracker/offline_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import logging
import typing as T

import cv2
import pupil_apriltags

from .surface_marker import Surface_Marker, Surface_Marker_Type
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
)
30 changes: 30 additions & 0 deletions pupil_src/shared_modules/surface_tracker/surface_tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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)

Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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",
Expand All @@ -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):
Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit 6427dc2

Please sign in to comment.