euler-preprocess 3.6.0__tar.gz → 3.9.0__tar.gz

euler_preprocess-3.9.0/MANIFEST.in

@@ -0,0 +1 @@
+recursive-exclude tests/fixtures *.zip

{euler_preprocess-3.6.0 → euler_preprocess-3.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: euler-preprocess
-Version: 3.6.0
+Version: 3.9.0
 Summary: Physics-based preprocessing (fog, etc.) for RGB+depth datasets
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown

{euler_preprocess-3.6.0 → euler_preprocess-3.9.0}/euler_preprocess/fog/inference.py

@@ -264,6 +264,7 @@ def _render_cpu(
         intrinsics=intrinsics_np,
         airlight_method=plan.airlight_method,
         capture_artifacts=plan.capture_artifacts,
+        clear_weather=plan.clear_weather,
     )
     return FogInferenceResult(
         rgb=np.asarray(result.rgb, dtype=np.float32),
@@ -300,30 +301,25 @@ def _render_gpu(
         depth_t = planar_to_radial_depth_torch(depth_t, K_t)
 
     sky_mask_t = torch.from_numpy(sky_mask).to(device=device, dtype=torch.bool)
-    estimated_airlight = transform._estimate_airlight_torch(
-        rgb_t,
-        sky_mask_t,
-        sample_id=sample_id,
-        method=plan.airlight_method,
-    )
     torch_gen = torch_generator_for_index(
         transform.torch_device,
         transform.seed,
         transform.base_rng,
         sample_index,
     )
-    foggy_t, beta, airlight_t, k_map_t, ls_map_t = transform._apply_model_torch(
+    foggy_t, beta, airlight_t, k_map_t, ls_map_t = transform._process_torch_pipeline(
         rgb_t,
         depth_t,
+        sky_mask_t,
         plan.model_name,
         plan.model_cfg,
         rng,
-        estimated_airlight,
         torch_gen,
         sample_id=sample_id,
         intrinsics=intrinsics_np,
-        depth_m=depth_t,
+        airlight_method=plan.airlight_method,
         capture_artifacts=plan.capture_artifacts,
+        clear_weather=plan.clear_weather,
     )
     return FogInferenceResult(
         rgb=torch.clamp(foggy_t, 0.0, 1.0).detach().cpu().numpy().astype(np.float32),

{euler_preprocess-3.6.0 → euler_preprocess-3.9.0}/euler_preprocess/fog/models.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import logging
 import math
+from typing import Any
 
 import numpy as np
 
@@ -1088,6 +1089,115 @@ def apply_scene_illumination_np(
     return (radiance * scale[..., None]).astype(np.float32, copy=False), ev_map
 
 
+def apply_scene_illumination_torch(
+    rgb_lin: "torch.Tensor",
+    depth_m: "torch.Tensor",
+    k_field,
+    model_cfg: dict,
+    rng: np.random.Generator,
+    sky_mask: "torch.Tensor | None" = None,
+) -> tuple["torch.Tensor", "torch.Tensor"]:
+    """Torch equivalent of :func:`apply_scene_illumination_np`."""
+    height, width = int(depth_m.shape[0]), int(depth_m.shape[1])
+    cfg = _resolve_scene_illumination_config(model_cfg)
+    enabled = cfg.get("enabled", False)
+    if not isinstance(enabled, bool):
+        raise ValueError("scene_illumination.enabled must be a boolean")
+    if not enabled:
+        return rgb_lin.to(dtype=torch.float32), torch.zeros(
+            (height, width),
+            device=rgb_lin.device,
+            dtype=torch.float32,
+        )
+
+    global_ev = _sample_float(
+        cfg.get("global_ev", 0.0),
+        rng,
+        "scene_illumination.global_ev",
+    )
+    near_ev = _sample_float(
+        cfg.get("near_ev", 0.0),
+        rng,
+        "scene_illumination.near_ev",
+    )
+    decay_depth = max(
+        _sample_float(
+            cfg.get("near_decay_depth_m", 15.0),
+            rng,
+            "scene_illumination.near_decay_depth_m",
+        ),
+        1e-6,
+    )
+    fog_ev = _sample_float(
+        cfg.get("fog_coupled_ev", 0.0),
+        rng,
+        "scene_illumination.fog_coupled_ev",
+    )
+
+    depth = _sanitize_depth_torch(depth_m).to(
+        device=rgb_lin.device,
+        dtype=rgb_lin.dtype,
+    )
+    k_map = torch.clamp(
+        broadcast_k_field_torch(
+            k_field,
+            height,
+            width,
+            device=rgb_lin.device,
+            dtype=rgb_lin.dtype,
+        ),
+        min=0.0,
+    )
+    near_weight = torch.exp(-depth / decay_depth)
+    fog_opacity = 1.0 - torch.exp(-k_map * depth)
+    ev_map = (
+        global_ev + near_ev * near_weight + fog_ev * torch.clamp(fog_opacity, 0.0, 1.0)
+    )
+
+    if sky_mask is not None:
+        sky = torch.as_tensor(
+            sky_mask,
+            device=rgb_lin.device,
+            dtype=rgb_lin.dtype,
+        )
+        if tuple(sky.shape) != (height, width):
+            raise ValueError(
+                f"sky_mask must have shape ({height}, {width}); got {tuple(sky.shape)}"
+            )
+        sky_weight = float(
+            np.clip(
+                _sample_float(
+                    cfg.get("sky_weight", 0.0),
+                    rng,
+                    "scene_illumination.sky_weight",
+                ),
+                0.0,
+                1.0,
+            )
+        )
+        ev_map = ev_map * (1.0 - torch.clamp(sky, 0.0, 1.0) * (1.0 - sky_weight))
+
+    min_scale = max(
+        _sample_float(
+            cfg.get("min_radiance_scale", 0.08),
+            rng,
+            "scene_illumination.min_radiance_scale",
+        ),
+        0.0,
+    )
+    scale = torch.maximum(
+        torch.exp2(-ev_map),
+        torch.full_like(ev_map, min_scale),
+    )
+    radiance = torch.nan_to_num(
+        rgb_lin.to(dtype=torch.float32),
+        nan=0.0,
+        posinf=1.0,
+        neginf=0.0,
+    )
+    return radiance * scale[..., None], ev_map
+
+
 def _apply_scene_airlight_dampening_np(
     ls_field: np.ndarray,
     ev_map: np.ndarray,
@@ -1109,6 +1219,33 @@ def _apply_scene_airlight_dampening_np(
     return np.clip(ls_map * scale[..., None], 0.0, 1.0).astype(np.float32, copy=False)
 
 
+def _apply_scene_airlight_dampening_torch(
+    ls_field: "torch.Tensor",
+    ev_map: "torch.Tensor",
+    model_cfg: dict,
+    rng: np.random.Generator,
+) -> "torch.Tensor":
+    cfg = _resolve_scene_illumination_config(model_cfg)
+    if not bool(cfg.get("enabled", False)):
+        return ls_field.to(dtype=torch.float32)
+    ratio = _sample_float(
+        cfg.get("airlight_ev_ratio", 0.0),
+        rng,
+        "scene_illumination.airlight_ev_ratio",
+    )
+    if ratio <= 0.0:
+        return ls_field.to(dtype=torch.float32)
+    scale = torch.exp2(-torch.clamp(ev_map, min=0.0) * ratio)
+    ls_map = broadcast_ls_field_torch(
+        ls_field,
+        int(ev_map.shape[0]),
+        int(ev_map.shape[1]),
+        device=ev_map.device,
+        dtype=ev_map.dtype,
+    )
+    return torch.clamp(ls_map * scale[..., None], 0.0, 1.0)
+
+
 def apply_fog(
     rgb: np.ndarray, depth_m: np.ndarray, k_field: np.ndarray, ls_field: np.ndarray
 ) -> np.ndarray:
@@ -1140,6 +1277,127 @@ def apply_fog_torch(
     return rgb * t + ls_field * (1.0 - t)
 
 
+def broadcast_k_field_torch(
+    k_field,
+    height: int,
+    width: int,
+    *,
+    device,
+    dtype,
+) -> "torch.Tensor":
+    """Return ``k_field`` as a Torch ``(H, W)`` map."""
+    field = torch.as_tensor(k_field, device=device, dtype=dtype)
+    if field.ndim == 0:
+        return field.expand(height, width)
+    if tuple(field.shape) == (height, width):
+        return field
+    raise ValueError(
+        f"k_field must be scalar or shape ({height}, {width}); got {tuple(field.shape)}"
+    )
+
+
+def broadcast_ls_field_torch(
+    ls_field,
+    height: int,
+    width: int,
+    *,
+    device,
+    dtype,
+) -> "torch.Tensor":
+    """Return ``ls_field`` as a Torch ``(H, W, 3)`` map."""
+    field = torch.as_tensor(ls_field, device=device, dtype=dtype)
+    if tuple(field.shape) == (3,):
+        return field.view(1, 1, 3).expand(height, width, 3)
+    if tuple(field.shape) == (1, 1, 3):
+        return field.expand(height, width, 3)
+    if tuple(field.shape) == (height, width, 3):
+        return field
+    raise ValueError(
+        f"ls_field must have shape (3,), (1, 1, 3), or "
+        f"({height}, {width}, 3); got {tuple(field.shape)}"
+    )
+
+
+def render_fog_fields_np(
+    rgb: np.ndarray,
+    depth_m: np.ndarray,
+    k_field,
+    ls_field: np.ndarray,
+    model_cfg: dict,
+    rng: np.random.Generator,
+    *,
+    sky_mask: np.ndarray | None = None,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Render resolved fog fields with the shared CPU scene semantics."""
+    height, width = depth_m.shape
+    scene_rgb, ev_map = apply_scene_illumination_np(
+        rgb,
+        depth_m,
+        k_field,
+        model_cfg,
+        rng,
+        sky_mask=sky_mask,
+    )
+    rendered_ls_field = _apply_scene_airlight_dampening_np(
+        ls_field,
+        ev_map,
+        model_cfg,
+        rng,
+    )
+    foggy = apply_fog(scene_rgb, depth_m, k_field, rendered_ls_field)
+    return (
+        foggy,
+        broadcast_k_field(k_field, height, width),
+        broadcast_ls_field(rendered_ls_field, height, width),
+    )
+
+
+def render_fog_fields_torch(
+    rgb: "torch.Tensor",
+    depth_m: "torch.Tensor",
+    k_field,
+    ls_field: "torch.Tensor",
+    model_cfg: dict,
+    rng: np.random.Generator,
+    *,
+    sky_mask: "torch.Tensor | None" = None,
+) -> tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]:
+    """Render resolved fog fields with CPU-equivalent Torch semantics."""
+    height, width = int(depth_m.shape[0]), int(depth_m.shape[1])
+    scene_rgb, ev_map = apply_scene_illumination_torch(
+        rgb,
+        depth_m,
+        k_field,
+        model_cfg,
+        rng,
+        sky_mask=sky_mask,
+    )
+    rendered_ls_field = _apply_scene_airlight_dampening_torch(
+        ls_field,
+        ev_map,
+        model_cfg,
+        rng,
+    )
+    foggy = apply_fog_torch(scene_rgb, depth_m, k_field, rendered_ls_field)
+    return (
+        foggy,
+        broadcast_k_field_torch(
+            k_field,
+            height,
+            width,
+            device=rgb.device,
+            dtype=rgb.dtype,
+        ),
+        broadcast_ls_field_torch(
+            rendered_ls_field,
+            height,
+            width,
+            device=rgb.device,
+            dtype=rgb.dtype,
+        ),
+    )
+
+
 def select_model(config: dict, rng: np.random.Generator) -> str:
     selection = config.get("selection")
     if selection is None:
@@ -1288,16 +1546,13 @@ def apply_model(
     else:
         ls_field = ls_base.reshape(1, 1, 3)
 
-    scene_rgb, ev_map = apply_scene_illumination_np(
+    foggy, k_map, ls_map = render_fog_fields_np(
         rgb,
         depth_m,
         k_field,
+        ls_field,
         model_cfg,
         rng,
         sky_mask=sky_mask,
     )
-    ls_field = _apply_scene_airlight_dampening_np(ls_field, ev_map, model_cfg, rng)
-    foggy = apply_fog(scene_rgb, depth_m, k_field, ls_field)
-    k_map = broadcast_k_field(k_field, height, width)
-    ls_map = broadcast_ls_field(ls_field, height, width)
     return foggy, k_mean, ls_base, k_map, ls_map

{euler_preprocess-3.6.0 → euler_preprocess-3.9.0}/euler_preprocess/fog/pipeline.py

@@ -5,7 +5,12 @@ from typing import Any
 
 import numpy as np
 
-from euler_preprocess.common.color import linear_to_srgb, srgb_to_linear
+from euler_preprocess.common.color import (
+    linear_to_srgb,
+    linear_to_srgb_torch,
+    srgb_to_linear,
+    srgb_to_linear_torch,
+)
 from euler_preprocess.fog.atmospheric_light import AtmosphericLightResolver
 from euler_preprocess.fog.capture import CaptureArtifactPipeline, CaptureContext
 from euler_preprocess.fog.models import apply_model
@@ -71,9 +76,7 @@ class FogProcessingPipeline:
         sample_id: str | None,
         airlight_method: str | None = None,
     ) -> FogPipelineResult:
-        render_input_space = _normalize_render_input_space(
-            model_cfg.get("render_input_space", self.render_input_space)
-        )
+        render_input_space = self.render_input_space_for(model_cfg)
         rgb_for_render = srgb_to_linear(rgb) if render_input_space == "srgb" else rgb
         estimated_airlight = self.atmospheric_light.estimate_np(
             rgb_for_render,
@@ -101,6 +104,61 @@ class FogProcessingPipeline:
             ls_map=ls_map,
         )
 
+    def render_input_space_for(self, model_cfg: dict) -> str:
+        """Resolve the input encoding used by both scene-rendering backends."""
+        return _normalize_render_input_space(
+            model_cfg.get("render_input_space", self.render_input_space)
+        )
+
+    def prepare_render_rgb_torch(self, rgb, model_cfg: dict):
+        """Convert display RGB to the scene-linear domain when configured."""
+        if self.render_input_space_for(model_cfg) == "srgb":
+            return srgb_to_linear_torch(rgb)
+        return rgb
+
+    def restore_render_rgb_torch(
+        self,
+        rgb,
+        model_cfg: dict,
+        *,
+        clear_weather: bool = False,
+    ):
+        """Return a rendered Torch image to the CPU pathway's output domain."""
+        if not clear_weather and self.render_input_space_for(model_cfg) == "srgb":
+            return linear_to_srgb_torch(rgb)
+        return rgb
+
+    def capture_context(
+        self,
+        *,
+        sample_id: str | None,
+        rng: np.random.Generator,
+        model_cfg: dict,
+        sky_mask,
+        airlight,
+        intrinsics=None,
+        depth_m=None,
+        k_map=None,
+        device=None,
+    ) -> CaptureContext:
+        """Build the complete capture context shared by CPU and Torch routes."""
+        return CaptureContext(
+            sample_id=sample_id,
+            rng=rng,
+            device=device,
+            intrinsics=intrinsics,
+            depth_m=depth_m,
+            k_map=k_map,
+            attributes={
+                "sky_mask": sky_mask,
+                "airlight": airlight,
+                "render_input_space": model_cfg.get(
+                    "render_input_space",
+                    self.render_input_space,
+                ),
+            },
+        )
+
     def process_np(
         self,
         *,
@@ -114,33 +172,39 @@ class FogProcessingPipeline:
         intrinsics: np.ndarray | None = None,
         airlight_method: str | None = None,
         capture_artifacts: CaptureArtifactPipeline | None = None,
+        clear_weather: bool = False,
     ) -> FogPipelineResult:
-        result = self.render_scene_np(
-            rgb=rgb,
-            depth_m=depth_m,
-            sky_mask=sky_mask,
-            model_name=model_name,
-            model_cfg=model_cfg,
-            rng=rng,
-            sample_id=sample_id,
-            airlight_method=airlight_method,
-        )
+        if clear_weather:
+            height, width = depth_m.shape
+            result = FogPipelineResult(
+                rgb=rgb,
+                beta=0.0,
+                airlight=np.zeros(3, dtype=np.float32),
+                k_map=np.zeros((height, width), dtype=np.float32),
+                ls_map=np.zeros((height, width, 3), dtype=np.float32),
+            )
+        else:
+            result = self.render_scene_np(
+                rgb=rgb,
+                depth_m=depth_m,
+                sky_mask=sky_mask,
+                model_name=model_name,
+                model_cfg=model_cfg,
+                rng=rng,
+                sample_id=sample_id,
+                airlight_method=airlight_method,
+            )
         return self.apply_capture_np(
             result,
-            context=CaptureContext(
+            context=self.capture_context(
                 sample_id=sample_id,
                 rng=rng,
+                model_cfg=model_cfg,
+                sky_mask=sky_mask,
+                airlight=result.airlight,
                 intrinsics=intrinsics,
                 depth_m=depth_m,
                 k_map=result.k_map,
-                attributes={
-                    "sky_mask": sky_mask,
-                    "airlight": result.airlight,
-                    "render_input_space": model_cfg.get(
-                        "render_input_space",
-                        self.render_input_space,
-                    ),
-                },
             ),
             capture_artifacts=capture_artifacts,
         )