import hashlib from dataclasses import dataclass from typing import Optional import cv2 import numpy as np @dataclass class PreprocessConfig: apply_darken: bool = False apply_brighten: bool = True darken_min: float = 0.3 darken_max: float = 0.8 brighten_method: str = "clahe" # "clahe" or "gamma" brighten_gamma_min: float = 1.2 brighten_gamma_max: float = 1.8 def _hash_seed(text: str) -> int: digest = hashlib.md5(text.encode("utf-8")).hexdigest() return int(digest[:8], 16) def _clamp_uint8(img: np.ndarray) -> np.ndarray: return np.clip(img, 0, 255).astype(np.uint8) def darken_frame(frame: np.ndarray, factor: float) -> np.ndarray: return _clamp_uint8(frame.astype(np.float32) * factor) def brighten_frame(frame: np.ndarray, method: str = "clahe", gamma: float = 1.5, darkness_threshold: int = 80) -> np.ndarray: # the smarter way to brighten is to first check if the frame is actually dark, and only apply brighten if it is. This way we avoid over-brightening already bright frames and introducing noise. gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) average_brightness = np.mean(gray) # If the frame is already bright enough, return it as is to save processing time and avoid over-brightening if average_brightness > darkness_threshold: return frame method = (method or "clahe").lower() if method == "clahe": lab = cv2.cvtColor(frame, cv2.COLOR_BGR2LAB) l, a, b = cv2.split(lab) clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) l2 = clahe.apply(l) merged = cv2.merge((l2, a, b)) return cv2.cvtColor(merged, cv2.COLOR_LAB2BGR) if method == "gamma": inv_gamma = 1.0 / max(gamma, 1e-6) table = np.array( [(i / 255.0) ** inv_gamma * 255 for i in np.arange(256)], dtype=np.uint8 ) return cv2.LUT(frame, table) raise ValueError(f"Unsupported brighten method: {method}") def apply_darken_then_brighten( frame: np.ndarray, *, config: PreprocessConfig, rng: np.random.Generator, darken_factor: Optional[float] = None, ) -> np.ndarray: processed = frame if config.apply_darken: if darken_factor is None: darken_factor = float(rng.uniform(config.darken_min, config.darken_max)) processed = darken_frame(processed, darken_factor) if config.apply_brighten: gamma = float(rng.uniform(config.brighten_gamma_min, config.brighten_gamma_max)) processed = brighten_frame(processed, method=config.brighten_method, gamma=gamma) return processed