filtered/preprocess/dehaze.py

#!/usr/bin/env python3
"""
Dark Channel Prior (DCP) 기반 디헤이즈 — He et al., CVPR 2009.

OpenCV에 완제품이 없어 직접 구현한 부분.
ROS에 의존하지 않으므로 주피터/CLI에서 그대로 import해 튜닝할 수 있습니다.

실시간성 최적화 (발표에서의 기여 포인트):
  - 대기광/투과율 추정을 축소 해상도(기본 1/4)에서 수행 → 픽셀 수 1/16
  - guided filter도 축소 해상도에서 수행 후 투과율만 업샘플
  - 복원식만 원본 해상도에서 계산 (벡터 연산 3줄)
  - min filter는 cv2.erode 로 대체 (사각 커널 erode == min filter, C 구현)

화재 도메인 특유의 함정:
  - 표준 DCP의 대기광 A 추정은 "가장 밝은 픽셀"을 고르는데, 우리 장면에서는
    그게 **불씨/화염**일 수 있습니다. A가 과대 추정되면 화면 전체가 어두워지고
    정작 불씨 주변이 뭉개집니다.
  - 대응: 상위 후보의 **평균**을 사용 + a_max로 상한 클리핑 + (옵션) 화면 상단
    영역만 후보로 삼기(sky_ratio). 아래 estimate_atmospheric_light 참고.
"""

from __future__ import annotations

import cv2
import numpy as np


# ---------------------------------------------------------------- 기본 연산


def dark_channel(img: np.ndarray, patch: int) -> np.ndarray:
    """다크 채널: 채널 최솟값 -> patch x patch 최소 필터.

    img: float32 (H, W, 3), 0~1 범위
    """
    min_ch = np.min(img, axis=2)
    if patch <= 1:
        return min_ch
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (patch, patch))
    # erode(사각 커널) == 최소 필터. 직접 슬라이딩 윈도우 짜면 수십 배 느림.
    return cv2.erode(min_ch, kernel)


def estimate_atmospheric_light(
    img: np.ndarray,
    dark: np.ndarray,
    top_ratio: float = 0.001,
    a_max: float = 0.92,
    sky_ratio: float = 1.0,
) -> np.ndarray:
    """대기광 A 추정 (1, 1, 3).

    top_ratio : 다크 채널 상위 몇 %를 후보로 볼지 (원논문 0.1%)
    a_max     : A 상한. 화염 같은 포화 픽셀이 섞였을 때 과대추정 방지
    sky_ratio : 후보를 이미지 상단 몇 비율로 제한할지 (1.0 = 제한 없음).
                지하주차장처럼 하늘이 없으면 1.0 유지, 연기가 위에 깔리면 0.5~0.7.
    """
    h, w = dark.shape
    limit = h if sky_ratio >= 1.0 else max(1, int(h * sky_ratio))

    dark_roi = dark[:limit]
    img_roi = img[:limit]

    n = max(int(dark_roi.size * top_ratio), 1)
    flat_dark = dark_roi.ravel()
    idx = np.argpartition(flat_dark, -n)[-n:]

    candidates = img_roi.reshape(-1, 3)[idx]  # (n, 3)
    # 원논문은 후보 중 "가장 밝은 한 픽셀"을 쓰지만, 스페큘러/화염 한 점에
    # 통째로 끌려갑니다. 평균이 훨씬 안정적입니다.
    a = candidates.mean(axis=0)
    a = np.clip(a, 1e-3, a_max)
    return a.reshape(1, 1, 3).astype(np.float32)


def guided_filter(guide: np.ndarray, src: np.ndarray, radius: int, eps: float) -> np.ndarray:
    """단일 채널 guided filter (He et al.).

    cv2.ximgproc.guidedFilter 는 opencv-contrib-python 이 있어야 쓸 수 있어서
    box filter 5번으로 직접 구현했습니다. 의존성 없이 동작합니다.
    """
    guide = guide.astype(np.float32)
    src = src.astype(np.float32)
    k = (2 * radius + 1, 2 * radius + 1)

    mean_i = cv2.blur(guide, k)
    mean_p = cv2.blur(src, k)
    corr_i = cv2.blur(guide * guide, k)
    corr_ip = cv2.blur(guide * src, k)

    var_i = corr_i - mean_i * mean_i
    cov_ip = corr_ip - mean_i * mean_p

    a = cov_ip / (var_i + eps)
    b = mean_p - a * mean_i

    mean_a = cv2.blur(a, k)
    mean_b = cv2.blur(b, k)
    return mean_a * guide + mean_b


# ---------------------------------------------------------------- 메인 클래스


class DarkChannelDehazer:
    """DCP 디헤이저.

    process(bgr_uint8) -> bgr_uint8  형태로 씁니다.
    파라미터는 ROS 파라미터에서 그대로 갈아끼울 수 있게 attribute로 노출.
    """

    def __init__(
        self,
        omega: float = 0.95,
        t0: float = 0.1,
        patch: int = 15,
        scale: float = 0.25,
        use_guided: bool = True,
        guided_radius: int = 8,
        guided_eps: float = 1e-3,
        a_top_ratio: float = 0.001,
        a_max: float = 0.92,
        sky_ratio: float = 1.0,
        a_smoothing: float = 0.0,
    ):
        self.omega = omega          # 1.0으로 하면 원근감이 사라져 부자연스러움 (원논문 0.95)
        self.t0 = t0                # 투과율 하한. 낮을수록 진한 연기까지 복원하지만 노이즈 폭발
        self.patch = patch          # 원본 해상도 기준 패치 크기
        self.scale = scale          # 추정 단계 축소 배율 (★ 실시간성의 핵심)
        self.use_guided = use_guided
        self.guided_radius = guided_radius  # 축소 해상도 기준 반경
        self.guided_eps = guided_eps
        self.a_top_ratio = a_top_ratio
        self.a_max = a_max
        self.sky_ratio = sky_ratio

        # 대기광 A의 프레임 간 EMA 계수 (0 = 끔, 0.9 = 강한 평활).
        # A는 매 프레임 장면 내용에서 추정되므로, 로봇이 움직여 화면이 바뀌면
        # **같은 불씨가 프레임마다 다른 밝기로 복원**됩니다. 영상에서는 깜빡임으로
        # 보이고, YOLO 학습 데이터로 쓰면 같은 물체의 외형 분산이 커집니다.
        # 시계열(동영상·rosbag·실주행)에는 켜고, 정지영상 비교 실험에는 끕니다.
        # ★ 켜면 출력이 이전 프레임에 의존하므로 더 이상 결정론적이지 않습니다.
        self.a_smoothing = a_smoothing
        self._a_ema: np.ndarray | None = None

        self.last_transmission: np.ndarray | None = None  # 디버깅/발표 그림용
        self.last_a: np.ndarray | None = None

    def reset_state(self) -> None:
        """프레임 간 누적 상태 초기화. 새 bag/영상을 시작할 때 호출."""
        self._a_ema = None

    # ------------------------------------------------------------------

    def _patch_for_scale(self) -> int:
        """축소본에서 같은 '실제 공간 범위'를 덮도록 패치도 함께 줄입니다.

        이걸 안 줄이면 축소본에서 패치가 상대적으로 너무 커져 투과율이
        과도하게 뭉개집니다(halo).
        """
        p = int(round(self.patch * self.scale))
        p = max(3, p)
        return p if p % 2 == 1 else p + 1

    def estimate_transmission(self, img_f: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """축소본에서 A와 투과율 t를 추정해 (t_full, A) 반환. img_f: float32 0~1 원본 해상도."""
        h, w = img_f.shape[:2]

        if self.scale < 1.0:
            small = cv2.resize(img_f, None, fx=self.scale, fy=self.scale,
                               interpolation=cv2.INTER_AREA)
        else:
            small = img_f

        patch_s = self._patch_for_scale()

        dark = dark_channel(small, patch_s)
        a = estimate_atmospheric_light(small, dark, self.a_top_ratio,
                                       self.a_max, self.sky_ratio)

        # 프레임 간 평활 — t 계산에 들어가기 **전에** 적용해야
        # 투과율까지 함께 안정됩니다.
        if self.a_smoothing > 0.0:
            if self._a_ema is None:
                self._a_ema = a.copy()
            else:
                k = float(np.clip(self.a_smoothing, 0.0, 0.99))
                self._a_ema = k * self._a_ema + (1.0 - k) * a
            a = self._a_ema

        # t = 1 - omega * darkchannel(I / A)
        t_small = 1.0 - self.omega * dark_channel(small / a, patch_s)

        if self.use_guided:
            gray = cv2.cvtColor((small * 255).astype(np.uint8), cv2.COLOR_BGR2GRAY)
            gray = gray.astype(np.float32) / 255.0
            t_small = guided_filter(gray, t_small, self.guided_radius, self.guided_eps)

        if self.scale < 1.0:
            # 투과율은 저주파 신호라 업샘플해도 손실이 거의 없습니다. 이게 성립하는
            # 덕분에 축소 추정이 정당화됩니다.
            t = cv2.resize(t_small, (w, h), interpolation=cv2.INTER_LINEAR)
        else:
            t = t_small

        return t, a

    # ------------------------------------------------------------------

    def process(self, bgr: np.ndarray) -> np.ndarray:
        """bgr uint8 -> 디헤이즈된 bgr uint8."""
        img_f = bgr.astype(np.float32) / 255.0

        t, a = self.estimate_transmission(img_f)
        t = np.clip(t, self.t0, 1.0)

        self.last_transmission = t
        self.last_a = a

        # 복원식 J = (I - A) / t + A
        out = (img_f - a) / t[..., None] + a
        return np.clip(out * 255.0, 0, 255).astype(np.uint8)

    def process_lowlight(
        self,
        bgr: np.ndarray,
        omega: float = 0.8,
        t0: float = 0.25,
    ) -> np.ndarray:
        """저조도 보정 (옵션).

        "저조도 영상을 반전하면 헤이즈 영상과 통계적으로 닮는다"는 관찰
        (Dong et al., 2011)을 이용해 **같은 DCP 코드를 재사용**합니다.
            반전 -> 디헤이즈 -> 반전
        CLAHE가 국소 대비를 올리는 것과 달리 이쪽은 전역 밝기를 끌어올립니다.

        ★ 반전 영역은 파라미터가 달라야 합니다 (로컬 테스트로 발견한 함정).
          - `a_max`: 연기 영역에서는 화염에 끌리는 걸 막으려 0.92로 조였지만,
            반전 영상은 원래 전체가 밝아 A가 정말 1.0 근처입니다. 여기에
            0.92를 그대로 쓰면 t가 하한까지 눌려 결과가 **오히려 어두워집니다.**
          - `omega`/`t0`: 반전 영상은 '헤이즈'가 화면 전체에 깔린 상태라
            강도를 낮추고(0.8) 하한을 올려야(0.25) 노이즈가 안 터집니다.

        조명이 완전히 꺼진 구간에서만 켜세요. 노이즈가 같이 증폭됩니다.
        """
        sub = DarkChannelDehazer(
            omega=omega,
            t0=t0,
            patch=self.patch,
            scale=self.scale,
            use_guided=self.use_guided,
            guided_radius=self.guided_radius,
            guided_eps=self.guided_eps,
            a_top_ratio=self.a_top_ratio,
            a_max=1.0,          # ★ 반전 영역에서는 조이지 않음
            sky_ratio=1.0,      # 반전되면 위/아래 의미가 뒤집혀 제한이 무의미
        )
        out = 255 - sub.process(255 - bgr)
        self.last_transmission = sub.last_transmission
        self.last_a = sub.last_a
        return out


# ---------------------------------------------------------------- CLAHE


class ClaheEnhancer:
    """LAB의 L 채널에만 CLAHE.

    BGR 각 채널에 따로 걸면 채널별 히스토그램이 제각각 늘어나 **색이 틀어집니다**.
    LAB는 밝기(L)와 색(a,b)이 분리돼 있어 L만 건드리면 색상은 보존됩니다.
    """

    def __init__(self, clip_limit: float = 2.0, tile_grid: tuple[int, int] = (8, 8)):
        self.clip_limit = clip_limit
        self.tile_grid = tuple(tile_grid)
        self._clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=self.tile_grid)

    def update(self, clip_limit: float, tile_grid) -> None:
        self.clip_limit = clip_limit
        self.tile_grid = tuple(tile_grid)
        self._clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=self.tile_grid)

    def process(self, bgr: np.ndarray) -> np.ndarray:
        lab = cv2.cvtColor(bgr, cv2.COLOR_BGR2LAB)
        lab[:, :, 0] = self._clahe.apply(lab[:, :, 0])
        return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)


def apply_gamma(bgr: np.ndarray, gamma: float) -> np.ndarray:
    """감마 보정: out = (in/255)^gamma * 255.

    **gamma < 1 이면 밝아지고, gamma > 1 이면 어두워집니다.**

    감마는 문헌마다 지수를 뒤집어 쓰는 경우가 있어(어떤 코드는 1/gamma를 씀)
    방향이 반대가 되기 쉽습니다. 여기서는 위 식으로 고정했고
    tests/test_dehaze.py::TestGamma 가 방향을 잠가둡니다.

    LUT 256개 조회라 해상도와 무관하게 비용이 거의 0입니다.
    """
    if abs(gamma - 1.0) < 1e-3:
        return bgr
    g = max(gamma, 1e-3)
    lut = np.clip((np.arange(256) / 255.0) ** g * 255.0, 0, 255).astype(np.uint8)
    return cv2.LUT(bgr, lut)