添加光照一致化ffc程序

src/orisys/frame_processors.py

@@ -1,5 +1,8 @@
+from __future__ import annotations
+
 import cv2
 import numpy as np
+from pathlib import Path
 
 from .unroll import UVRemap
 
@@ -229,3 +232,180 @@ def create_frame_processor(sensor_type, stitcher=None, uv_remap=None, **uv_kwarg
             raise ValueError("UV unroll sensor requires a UVRemap instance.")
         return UVUnrollFrameProcessor(uv_remap, **uv_kwargs)
     raise ValueError(f"Unsupported sensor type: {sensor_type}")
+
+
+class FFCFrameProcessor:
+    """Flat-field correction with offline weight-map calibration and fast online correction."""
+
+    def __init__(self, blur_kernel=151, blur_sigma=0.0, epsilon=1e-6, masked_outside=True):
+        self.blur_kernel = ensure_odd(blur_kernel)
+        self.blur_sigma = float(blur_sigma)
+        self.epsilon = float(epsilon)
+        self.masked_outside = bool(masked_outside)
+        self.weight_map = None
+        self.weight_map_3ch = None
+        self.roi_mask = None
+        self._multiplied_buffer = None
+        self._output_buffer = None
+        self._buffer_shape = None
+        self._buffer_dtype = None
+        self._roi_bbox = None
+
+    def calibrate(self, reference_image, roi_mask=None):
+        gray = to_gray_float32(reference_image)
+        roi_mask = normalize_mask(roi_mask, gray.shape)
+        active = roi_mask > 0
+        if not np.any(active):
+            raise ValueError("ROI mask is empty; cannot calibrate FFC.")
+
+        background = cv2.GaussianBlur(
+            gray,
+            (self.blur_kernel, self.blur_kernel),
+            sigmaX=self.blur_sigma,
+            sigmaY=self.blur_sigma,
+            borderType=cv2.BORDER_REPLICATE,
+        )
+        mean_value = float(background[active].mean())
+        weight = mean_value / (background + self.epsilon)
+
+        outside_value = 0.0 if self.masked_outside else 1.0
+        self.weight_map = np.where(active, weight, outside_value).astype(np.float32)
+        self.weight_map_3ch = None
+        self.roi_mask = roi_mask
+        self._roi_bbox = self._compute_roi_bbox(roi_mask)
+        self._clear_runtime_buffers()
+        return self.weight_map
+
+    def correct(self, image):
+        if self.weight_map is None:
+            raise RuntimeError("FFC is not calibrated. Call calibrate() or load() first.")
+        if image.shape[:2] != self.weight_map.shape:
+            raise ValueError(
+                f"image shape {image.shape[:2]} does not match weight map {self.weight_map.shape}"
+            )
+        if image.ndim not in (2, 3):
+            raise ValueError("image must have shape HxW or HxWxC")
+
+        self._ensure_runtime_buffers(image)
+        if self.masked_outside:
+            self._output_buffer.fill(0)
+
+        if self._roi_bbox is None:
+            src = image
+            weight = self.weight_map if image.ndim == 2 else self.weight_map_3ch
+            dst_view = self._output_buffer
+        else:
+            x1, y1, x2, y2 = self._roi_bbox
+            src = image[y1:y2, x1:x2]
+            if image.ndim == 2:
+                weight = self.weight_map[y1:y2, x1:x2]
+            else:
+                weight = self.weight_map_3ch[y1:y2, x1:x2]
+            dst_view = self._output_buffer[y1:y2, x1:x2]
+
+        cv2.multiply(src, weight, dst=self._multiplied_buffer, dtype=cv2.CV_32F)
+        if np.issubdtype(image.dtype, np.uint8):
+            cv2.convertScaleAbs(self._multiplied_buffer, dst=dst_view, alpha=1.0, beta=0.0)
+            return self._output_buffer
+
+        corrected = saturate_to_dtype(self._multiplied_buffer, image.dtype)
+        dst_view[...] = corrected
+        return self._output_buffer
+
+    def save(self, path):
+        if self.weight_map is None:
+            raise RuntimeError("No weight map to save.")
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        np.savez_compressed(
+            path,
+            weight_map=self.weight_map,
+            roi_mask=self.roi_mask,
+            blur_kernel=np.array(self.blur_kernel, dtype=np.int32),
+            blur_sigma=np.array(self.blur_sigma, dtype=np.float32),
+            epsilon=np.array(self.epsilon, dtype=np.float32),
+            masked_outside=np.array(self.masked_outside, dtype=np.bool_),
+        )
+
+    def load(self, path):
+        data = np.load(path)
+        self.weight_map = data["weight_map"].astype(np.float32, copy=False)
+        self.weight_map_3ch = None
+        self.roi_mask = data["roi_mask"].astype(np.uint8, copy=False)
+        self._roi_bbox = self._compute_roi_bbox(self.roi_mask)
+        self.blur_kernel = int(data.get("blur_kernel", self.blur_kernel))
+        self.blur_sigma = float(data.get("blur_sigma", self.blur_sigma))
+        self.epsilon = float(data.get("epsilon", self.epsilon))
+        self.masked_outside = bool(data.get("masked_outside", self.masked_outside))
+        self._clear_runtime_buffers()
+        return self
+
+    def _clear_runtime_buffers(self):
+        self._multiplied_buffer = None
+        self._output_buffer = None
+        self._buffer_shape = None
+        self._buffer_dtype = None
+
+    @staticmethod
+    def _compute_roi_bbox(roi_mask):
+        ys, xs = np.nonzero(roi_mask > 0)
+        if xs.size == 0:
+            return None
+        return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1
+
+    def _ensure_runtime_buffers(self, image):
+        shape = image.shape
+        dtype = image.dtype
+        if self._buffer_shape == shape and self._buffer_dtype == dtype:
+            return
+
+        self._buffer_shape = shape
+        self._buffer_dtype = dtype
+        self._output_buffer = np.empty(shape, dtype=dtype)
+
+        if image.ndim == 3:
+            channels = image.shape[2]
+            self.weight_map_3ch = np.repeat(self.weight_map[:, :, None], channels, axis=2)
+            if self._roi_bbox is None:
+                self._multiplied_buffer = np.empty(shape, dtype=np.float32)
+            else:
+                x1, y1, x2, y2 = self._roi_bbox
+                self._multiplied_buffer = np.empty((y2 - y1, x2 - x1, channels), dtype=np.float32)
+        else:
+            self.weight_map_3ch = None
+            if self._roi_bbox is None:
+                self._multiplied_buffer = np.empty(shape, dtype=np.float32)
+            else:
+                x1, y1, x2, y2 = self._roi_bbox
+                self._multiplied_buffer = np.empty((y2 - y1, x2 - x1), dtype=np.float32)
+
+
+def ensure_odd(value, minimum=3):
+    value = max(int(value), int(minimum))
+    return value if value % 2 == 1 else value + 1
+
+
+def to_gray_float32(image):
+    if image.ndim == 2:
+        gray = image
+    elif image.ndim == 3:
+        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    else:
+        raise ValueError("image must have shape HxW or HxWxC")
+    return gray.astype(np.float32, copy=False)
+
+
+def normalize_mask(mask, shape):
+    if mask is None:
+        return np.ones(shape, dtype=np.uint8) * 255
+    mask = np.asarray(mask)
+    if mask.shape[:2] != shape:
+        mask = cv2.resize(mask, (shape[1], shape[0]), interpolation=cv2.INTER_NEAREST)
+    return np.where(mask > 0, 255, 0).astype(np.uint8)
+
+
+def saturate_to_dtype(values, dtype):
+    if np.issubdtype(dtype, np.integer):
+        info = np.iinfo(dtype)
+        return np.clip(values, info.min, info.max).astype(dtype)
+    return values.astype(dtype, copy=False)