euler-preprocess 3.1.0__tar.gz → 3.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (53) hide show
  1. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/PKG-INFO +1 -1
  2. euler_preprocess-3.2.0/euler_preprocess/fog/__init__.py +11 -0
  3. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/atmospheric_light.py +8 -5
  4. euler_preprocess-3.2.0/euler_preprocess/fog/inference.py +379 -0
  5. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/transform.py +151 -25
  6. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess.egg-info/PKG-INFO +1 -1
  7. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess.egg-info/SOURCES.txt +1 -0
  8. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/pyproject.toml +1 -1
  9. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/tests/test_fog_aux_outputs.py +252 -0
  10. euler_preprocess-3.1.0/euler_preprocess/sky_depth/__init__.py +0 -0
  11. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/README.md +0 -0
  12. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/__init__.py +0 -0
  13. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/cli.py +0 -0
  14. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/common/__init__.py +0 -0
  15. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/common/dataset.py +0 -0
  16. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/common/device.py +0 -0
  17. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/common/intrinsics.py +0 -0
  18. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/common/io.py +0 -0
  19. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/common/logging.py +0 -0
  20. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/common/noise.py +0 -0
  21. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/common/normalize.py +0 -0
  22. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/common/output.py +0 -0
  23. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/common/sampling.py +0 -0
  24. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/common/transform.py +0 -0
  25. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/airlight_from_sky.py +0 -0
  26. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/augmentations.py +0 -0
  27. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/capture.py +0 -0
  28. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/dcp_airlight.py +0 -0
  29. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/dcp_airlight_torch.py +0 -0
  30. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/dcp_heuristic_airlight.py +0 -0
  31. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/dcp_heuristic_airlight_torch.py +0 -0
  32. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/foggify.py +0 -0
  33. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/foggify_logging.py +0 -0
  34. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/logging.py +0 -0
  35. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/models.py +0 -0
  36. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/fog/pipeline.py +0 -0
  37. {euler_preprocess-3.1.0/euler_preprocess/fog → euler_preprocess-3.2.0/euler_preprocess/radial}/__init__.py +0 -0
  38. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/radial/transform.py +0 -0
  39. {euler_preprocess-3.1.0/euler_preprocess/radial → euler_preprocess-3.2.0/euler_preprocess/sky_depth}/__init__.py +0 -0
  40. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess/sky_depth/transform.py +0 -0
  41. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess.egg-info/dependency_links.txt +0 -0
  42. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess.egg-info/entry_points.txt +0 -0
  43. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess.egg-info/requires.txt +0 -0
  44. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/euler_preprocess.egg-info/top_level.txt +0 -0
  45. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/setup.cfg +0 -0
  46. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/tests/test_airlight_fallback.py +0 -0
  47. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/tests/test_cli_sample_selection.py +0 -0
  48. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/tests/test_dcp_heuristic_airlight.py +0 -0
  49. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/tests/test_foggify_integration.py +0 -0
  50. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/tests/test_radial.py +0 -0
  51. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/tests/test_sky_depth.py +0 -0
  52. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/tests/test_source_backed_output.py +0 -0
  53. {euler_preprocess-3.1.0 → euler_preprocess-3.2.0}/tests/test_zip_output.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: euler-preprocess
3
- Version: 3.1.0
3
+ Version: 3.2.0
4
4
  Summary: Physics-based preprocessing (fog, etc.) for RGB+depth datasets
5
5
  Requires-Python: >=3.9
6
6
  Description-Content-Type: text/markdown
@@ -0,0 +1,11 @@
1
+ from euler_preprocess.fog.inference import (
2
+ FogInferenceResult,
3
+ render_fog_image,
4
+ render_fog_sample,
5
+ )
6
+
7
+ __all__ = [
8
+ "FogInferenceResult",
9
+ "render_fog_image",
10
+ "render_fog_sample",
11
+ ]
@@ -155,12 +155,13 @@ class AtmosphericLightResolver:
155
155
  items: list[dict],
156
156
  device: Any,
157
157
  contrast_threshold_default: float,
158
+ method: str | None = None,
158
159
  ) -> tuple[list[float], "torch.Tensor"]:
159
160
  """Resolve beta and dampened base L_s for a uniform-model torch batch."""
160
161
  al_spec = items[0]["model_cfg"].get("atmospheric_light", "from_sky")
161
162
  airlight_is_estimated = uses_estimated_airlight(al_spec)
162
163
  if airlight_is_estimated:
163
- ls_base = self._estimate_batch_torch(rgb_batch, items, device)
164
+ ls_base = self._estimate_batch_torch(rgb_batch, items, device, method)
164
165
  ls_base = normalize_atmospheric_light_torch(ls_base)
165
166
  else:
166
167
  ls_values = []
@@ -209,17 +210,19 @@ class AtmosphericLightResolver:
209
210
  rgb_batch: "torch.Tensor",
210
211
  items: list[dict],
211
212
  device: Any,
213
+ method: str | None = None,
212
214
  ) -> "torch.Tensor":
213
- if self.method == "from_sky":
215
+ resolved_method = method or self.method
216
+ if resolved_method == "from_sky":
214
217
  return self._estimate_from_sky_batch_torch(rgb_batch, items, device)
215
- estimator = self._get_estimator_torch(self.method)
218
+ estimator = self._get_estimator_torch(resolved_method)
216
219
  if estimator is None:
217
220
  raise RuntimeError(
218
- f"Torch airlight estimator unavailable for method '{self.method}'."
221
+ f"Torch airlight estimator unavailable for method '{resolved_method}'."
219
222
  )
220
223
  al_list = []
221
224
  for idx, item in enumerate(items):
222
- if self.method == "dcp":
225
+ if resolved_method == "dcp":
223
226
  al_list.append(estimator.compute(rgb_batch[idx]))
224
227
  else:
225
228
  sky_mask_t = torch.from_numpy(item["sky_mask"]).to(
@@ -0,0 +1,379 @@
1
+ from __future__ import annotations
2
+
3
+ import json
4
+ import tempfile
5
+ from dataclasses import dataclass
6
+ from pathlib import Path
7
+ from typing import Any, Literal, Mapping, Sequence
8
+
9
+ import numpy as np
10
+
11
+ from euler_preprocess.common.intrinsics import (
12
+ planar_to_radial_depth,
13
+ planar_to_radial_depth_torch,
14
+ )
15
+ from euler_preprocess.common.io import load_json
16
+ from euler_preprocess.common.normalize import (
17
+ _is_chw,
18
+ _to_numpy,
19
+ normalize_depth,
20
+ normalize_rgb,
21
+ normalize_rgb_torch,
22
+ normalize_sky_mask,
23
+ )
24
+ from euler_preprocess.common.device import torch_generator_for_index
25
+ from euler_preprocess.fog.transform import FogTransform
26
+
27
+ try:
28
+ import torch
29
+ except ImportError: # pragma: no cover - torch is optional
30
+ torch = None
31
+
32
+
33
+ FogInferenceMode = Literal["cpu", "gpu"]
34
+
35
+
36
+ @dataclass(frozen=True)
37
+ class FogInferenceResult:
38
+ """In-memory fog render output for experimental single-sample inference."""
39
+
40
+ rgb: np.ndarray
41
+ model_name: str
42
+ beta: float
43
+ airlight: np.ndarray
44
+ k_map: np.ndarray
45
+ ls_map: np.ndarray
46
+ scenario_name: str | None = None
47
+
48
+
49
+ def render_fog_image(
50
+ *,
51
+ rgb: Any,
52
+ depth: Any,
53
+ semantic_segmentation: Any,
54
+ intrinsics: Any,
55
+ config_path: str | Path | None = None,
56
+ config: Mapping[str, Any] | None = None,
57
+ scenario_profile_name: str | None = None,
58
+ mode: FogInferenceMode = "cpu",
59
+ sample_id: str | None = None,
60
+ sample_index: int = 0,
61
+ sky_class: Sequence[int | float] | None = (29, 0, 0),
62
+ ) -> np.ndarray:
63
+ """Render and return only the RGB image for one in-memory sample."""
64
+
65
+ return render_fog_sample(
66
+ rgb=rgb,
67
+ depth=depth,
68
+ semantic_segmentation=semantic_segmentation,
69
+ intrinsics=intrinsics,
70
+ config_path=config_path,
71
+ config=config,
72
+ scenario_profile_name=scenario_profile_name,
73
+ mode=mode,
74
+ sample_id=sample_id,
75
+ sample_index=sample_index,
76
+ sky_class=sky_class,
77
+ ).rgb
78
+
79
+
80
+ def render_fog_sample(
81
+ *,
82
+ rgb: Any,
83
+ depth: Any,
84
+ semantic_segmentation: Any,
85
+ intrinsics: Any,
86
+ config_path: str | Path | None = None,
87
+ config: Mapping[str, Any] | None = None,
88
+ scenario_profile_name: str | None = None,
89
+ mode: FogInferenceMode = "cpu",
90
+ sample_id: str | None = None,
91
+ sample_index: int = 0,
92
+ sky_class: Sequence[int | float] | None = (29, 0, 0),
93
+ ) -> FogInferenceResult:
94
+ """Render fog/camera artifacts for a single in-memory sample.
95
+
96
+ The call accepts the four modalities used by :class:`FogTransform` without
97
+ requiring a dataset reader or output backend. ``semantic_segmentation`` may
98
+ be either a boolean sky mask or an RGB semantic label map; RGB maps are
99
+ converted using ``sky_class`` which defaults to ``[29, 0, 0]``.
100
+ """
101
+
102
+ if mode not in ("cpu", "gpu"):
103
+ raise ValueError("mode must be 'cpu' or 'gpu'")
104
+ if config_path is None and config is None:
105
+ raise ValueError("Either config_path or config must be provided")
106
+
107
+ base_config = _load_config(config_path=config_path, config=config)
108
+ runtime_config = _config_for_mode(base_config, mode)
109
+
110
+ with tempfile.TemporaryDirectory(prefix="euler-fog-inference-") as tmpdir:
111
+ tmp_path = Path(tmpdir)
112
+ resolved_config_path = tmp_path / "config.json"
113
+ resolved_config_path.write_text(json.dumps(runtime_config), encoding="utf-8")
114
+ transform = FogTransform(
115
+ config_path=str(resolved_config_path),
116
+ out_path=str(tmp_path / "out"),
117
+ )
118
+ return _render_with_transform(
119
+ transform,
120
+ rgb=rgb,
121
+ depth=depth,
122
+ semantic_segmentation=semantic_segmentation,
123
+ intrinsics=intrinsics,
124
+ scenario_profile_name=scenario_profile_name,
125
+ mode=mode,
126
+ sample_id=sample_id,
127
+ sample_index=sample_index,
128
+ sky_class=sky_class,
129
+ )
130
+
131
+
132
+ def _load_config(
133
+ *,
134
+ config_path: str | Path | None,
135
+ config: Mapping[str, Any] | None,
136
+ ) -> dict[str, Any]:
137
+ if config_path is not None:
138
+ loaded = load_json(Path(config_path))
139
+ if config is None:
140
+ return dict(loaded)
141
+ merged = dict(loaded)
142
+ merged.update(dict(config))
143
+ return merged
144
+ return dict(config or {})
145
+
146
+
147
+ def _config_for_mode(config: dict[str, Any], mode: FogInferenceMode) -> dict[str, Any]:
148
+ resolved = dict(config)
149
+ if mode == "cpu":
150
+ resolved["device"] = "cpu"
151
+ return resolved
152
+
153
+ configured_device = str(resolved.get("device", "")).strip().lower()
154
+ if configured_device in ("", "cpu"):
155
+ resolved["device"] = "gpu"
156
+ return resolved
157
+
158
+
159
+ def _render_with_transform(
160
+ transform: FogTransform,
161
+ *,
162
+ rgb: Any,
163
+ depth: Any,
164
+ semantic_segmentation: Any,
165
+ intrinsics: Any,
166
+ scenario_profile_name: str | None,
167
+ mode: FogInferenceMode,
168
+ sample_id: str | None,
169
+ sample_index: int,
170
+ sky_class: Sequence[int | float] | None,
171
+ ) -> FogInferenceResult:
172
+ sample_id = sample_id or "inference_sample"
173
+ rgb_np = normalize_rgb(rgb)
174
+ depth_np = normalize_depth(
175
+ depth,
176
+ rgb_np.shape[:2],
177
+ transform.resize_depth_flag,
178
+ )
179
+ intrinsics_np = _normalize_intrinsics(intrinsics)
180
+ sky_mask = _normalize_semantic_sky_mask(
181
+ semantic_segmentation,
182
+ target_shape=rgb_np.shape[:2],
183
+ sky_class=sky_class,
184
+ )
185
+
186
+ rng = transform._rng_for(sample_index)
187
+ scenario = _select_scenario_profile(transform, scenario_profile_name)
188
+ plan = transform._resolve_render_plan(
189
+ rng,
190
+ scenario=scenario,
191
+ sample_scenario=scenario is None,
192
+ )
193
+
194
+ if mode == "gpu":
195
+ return _render_gpu(
196
+ transform,
197
+ rgb_np=rgb_np,
198
+ depth_np=depth_np,
199
+ sky_mask=sky_mask,
200
+ intrinsics_np=intrinsics_np,
201
+ plan=plan,
202
+ rng=rng,
203
+ sample_id=sample_id,
204
+ sample_index=sample_index,
205
+ )
206
+ return _render_cpu(
207
+ transform,
208
+ rgb_np=rgb_np,
209
+ depth_np=depth_np,
210
+ sky_mask=sky_mask,
211
+ intrinsics_np=intrinsics_np,
212
+ plan=plan,
213
+ rng=rng,
214
+ sample_id=sample_id,
215
+ )
216
+
217
+
218
+ def _select_scenario_profile(
219
+ transform: FogTransform,
220
+ scenario_profile_name: str | None,
221
+ ) -> dict[str, Any] | None:
222
+ if scenario_profile_name is None:
223
+ return None
224
+ for profile in transform.scenario_profiles:
225
+ if transform._scenario_name(profile) == scenario_profile_name:
226
+ return profile
227
+ known = ", ".join(
228
+ name
229
+ for name in (
230
+ transform._scenario_name(profile)
231
+ for profile in transform.scenario_profiles
232
+ )
233
+ if name is not None
234
+ )
235
+ raise ValueError(
236
+ f"Unknown scenario_profile_name '{scenario_profile_name}'. "
237
+ f"Known: {known or '<none>'}"
238
+ )
239
+
240
+
241
+ def _render_cpu(
242
+ transform: FogTransform,
243
+ *,
244
+ rgb_np: np.ndarray,
245
+ depth_np: np.ndarray,
246
+ sky_mask: np.ndarray,
247
+ intrinsics_np: np.ndarray,
248
+ plan,
249
+ rng: np.random.Generator,
250
+ sample_id: str,
251
+ ) -> FogInferenceResult:
252
+ depth_m = np.maximum(depth_np * transform.depth_scale, 0.0)
253
+ if intrinsics_np is not None:
254
+ depth_m = planar_to_radial_depth(depth_m, intrinsics_np)
255
+
256
+ result = transform.pipeline.process_np(
257
+ rgb=rgb_np,
258
+ depth_m=depth_m,
259
+ sky_mask=sky_mask,
260
+ model_name=plan.model_name,
261
+ model_cfg=plan.model_cfg,
262
+ rng=rng,
263
+ sample_id=sample_id,
264
+ intrinsics=intrinsics_np,
265
+ airlight_method=plan.airlight_method,
266
+ capture_artifacts=plan.capture_artifacts,
267
+ )
268
+ return FogInferenceResult(
269
+ rgb=np.asarray(result.rgb, dtype=np.float32),
270
+ model_name=plan.model_name,
271
+ beta=float(result.beta),
272
+ airlight=np.asarray(result.airlight, dtype=np.float32),
273
+ k_map=np.asarray(result.k_map, dtype=np.float32),
274
+ ls_map=np.asarray(result.ls_map, dtype=np.float32),
275
+ scenario_name=plan.scenario_name,
276
+ )
277
+
278
+
279
+ def _render_gpu(
280
+ transform: FogTransform,
281
+ *,
282
+ rgb_np: np.ndarray,
283
+ depth_np: np.ndarray,
284
+ sky_mask: np.ndarray,
285
+ intrinsics_np: np.ndarray,
286
+ plan,
287
+ rng: np.random.Generator,
288
+ sample_id: str,
289
+ sample_index: int,
290
+ ) -> FogInferenceResult:
291
+ if torch is None or transform.torch_device is None:
292
+ raise RuntimeError("Torch device not configured for GPU inference")
293
+
294
+ device = transform.torch_device
295
+ rgb_t = normalize_rgb_torch(rgb_np, device)
296
+ depth_t = torch.from_numpy(depth_np).to(device=device, dtype=torch.float32)
297
+ depth_t = torch.clamp(depth_t * transform.depth_scale, min=0.0)
298
+ if intrinsics_np is not None:
299
+ K_t = torch.from_numpy(intrinsics_np).to(device=device, dtype=torch.float32)
300
+ depth_t = planar_to_radial_depth_torch(depth_t, K_t)
301
+
302
+ sky_mask_t = torch.from_numpy(sky_mask).to(device=device, dtype=torch.bool)
303
+ estimated_airlight = transform._estimate_airlight_torch(
304
+ rgb_t,
305
+ sky_mask_t,
306
+ sample_id=sample_id,
307
+ method=plan.airlight_method,
308
+ )
309
+ torch_gen = torch_generator_for_index(
310
+ transform.torch_device,
311
+ transform.seed,
312
+ transform.base_rng,
313
+ sample_index,
314
+ )
315
+ foggy_t, beta, airlight_t, k_map_t, ls_map_t = transform._apply_model_torch(
316
+ rgb_t,
317
+ depth_t,
318
+ plan.model_name,
319
+ plan.model_cfg,
320
+ rng,
321
+ estimated_airlight,
322
+ torch_gen,
323
+ sample_id=sample_id,
324
+ intrinsics=intrinsics_np,
325
+ depth_m=depth_t,
326
+ capture_artifacts=plan.capture_artifacts,
327
+ )
328
+ return FogInferenceResult(
329
+ rgb=torch.clamp(foggy_t, 0.0, 1.0).detach().cpu().numpy().astype(np.float32),
330
+ model_name=plan.model_name,
331
+ beta=float(beta),
332
+ airlight=airlight_t.detach().cpu().numpy().astype(np.float32),
333
+ k_map=k_map_t.detach().cpu().numpy().astype(np.float32),
334
+ ls_map=ls_map_t.detach().cpu().numpy().astype(np.float32),
335
+ scenario_name=plan.scenario_name,
336
+ )
337
+
338
+
339
+ def _normalize_intrinsics(intrinsics: Any) -> np.ndarray:
340
+ if isinstance(intrinsics, Mapping):
341
+ if "intrinsics" in intrinsics:
342
+ intrinsics = intrinsics["intrinsics"]
343
+ elif all(key in intrinsics for key in ("fx", "fy", "cx", "cy")):
344
+ return np.array(
345
+ [
346
+ [float(intrinsics["fx"]), 0.0, float(intrinsics["cx"])],
347
+ [0.0, float(intrinsics["fy"]), float(intrinsics["cy"])],
348
+ [0.0, 0.0, 1.0],
349
+ ],
350
+ dtype=np.float32,
351
+ )
352
+ arr = _to_numpy(intrinsics).astype(np.float32)
353
+ if arr.shape != (3, 3):
354
+ raise ValueError(f"intrinsics must have shape (3, 3), got {arr.shape}")
355
+ return arr
356
+
357
+
358
+ def _normalize_semantic_sky_mask(
359
+ semantic_segmentation: Any,
360
+ *,
361
+ target_shape: tuple[int, int],
362
+ sky_class: Sequence[int | float] | None,
363
+ ) -> np.ndarray:
364
+ semantic = _to_numpy(semantic_segmentation)
365
+ if _is_chw(semantic):
366
+ semantic = np.transpose(semantic, (1, 2, 0))
367
+
368
+ if semantic.ndim == 3 and semantic.shape[-1] >= 3 and sky_class is not None:
369
+ sky_value = np.asarray(sky_class, dtype=semantic.dtype).reshape(1, 1, -1)
370
+ mask = np.all(semantic[..., : sky_value.shape[-1]] == sky_value, axis=-1)
371
+ else:
372
+ mask = normalize_sky_mask(semantic)
373
+
374
+ if mask.shape != target_shape:
375
+ raise ValueError(
376
+ f"semantic_segmentation shape {mask.shape} does not match "
377
+ f"image shape {target_shape}"
378
+ )
379
+ return mask.astype(bool, copy=False)
@@ -116,6 +116,9 @@ _SCENARIO_CONTROL_KEYS = {
116
116
  }
117
117
 
118
118
 
119
+ _GPU_BATCH_SCOPE_VALUES = {"sample", "batch"}
120
+
121
+
119
122
  @dataclass(frozen=True)
120
123
  class RenderPlan:
121
124
  model_name: str
@@ -124,6 +127,23 @@ class RenderPlan:
124
127
  capture_artifacts: CaptureArtifactPipeline | None = None
125
128
  scenario_name: str | None = None
126
129
 
130
+
131
+ def _freeze_distribution_specs(value: Any, rng: np.random.Generator) -> Any:
132
+ """Resolve distribution specs while preserving ordinary config structure."""
133
+ if isinstance(value, dict):
134
+ if "dist" in value:
135
+ return sample_value(value, rng)
136
+ return {
137
+ key: _freeze_distribution_specs(child, rng)
138
+ for key, child in value.items()
139
+ }
140
+ if isinstance(value, list):
141
+ return [_freeze_distribution_specs(child, rng) for child in value]
142
+ if isinstance(value, tuple):
143
+ return tuple(_freeze_distribution_specs(child, rng) for child in value)
144
+ return value
145
+
146
+
127
147
  # Use the canonical euler-loading ``generic.map_2d`` / ``map_3d`` modality
128
148
  # annotations. Auxiliary outputs are written as ``.npy`` files in the
129
149
  # layout the matching loader expects (``map_2d`` → ``(H, W)``,
@@ -244,6 +264,10 @@ class FogTransform(Transform):
244
264
  )
245
265
  self.augmentation_specs = list(self.augmentation_config.specs)
246
266
  self.scenario_profiles = self._parse_scenario_profiles(self.config)
267
+ (
268
+ self.gpu_scenario_scope,
269
+ self.gpu_condition_parameter_scope,
270
+ ) = self._parse_gpu_batching_config(self.config)
247
271
  self._configure_output_layout_metadata()
248
272
  self._written_configs: set[str] = set()
249
273
  self.torch_device = None
@@ -317,6 +341,15 @@ class FogTransform(Transform):
317
341
  return np.random.default_rng(np.random.SeedSequence(seed_parts))
318
342
  return self.base_rng
319
343
 
344
+ def _rng_for_batch(self, batch_index: int):
345
+ if self.seed is not None:
346
+ return np.random.default_rng(
347
+ np.random.SeedSequence(
348
+ [int(self.seed), int(batch_index), 1_000_003]
349
+ )
350
+ )
351
+ return self.base_rng
352
+
320
353
  def _get_airlight_estimator(self, method: str):
321
354
  return self.atmospheric_light.get_estimator(method)
322
355
 
@@ -380,6 +413,45 @@ class FogTransform(Transform):
380
413
  return tuple(profiles)
381
414
  return ()
382
415
 
416
+ def _parse_gpu_batching_config(
417
+ self,
418
+ config: dict[str, Any],
419
+ ) -> tuple[str, str]:
420
+ raw = config.get("gpu_batching", {})
421
+ if raw is None or raw is False:
422
+ raw = {}
423
+ if raw is True:
424
+ raw = {"scenario_scope": "batch"}
425
+ if not isinstance(raw, dict):
426
+ raise ValueError("gpu_batching must be a boolean or object")
427
+
428
+ scenario_scope = str(
429
+ raw.get("scenario_scope", config.get("gpu_scenario_scope", "sample"))
430
+ ).lower()
431
+ if scenario_scope not in _GPU_BATCH_SCOPE_VALUES:
432
+ raise ValueError(
433
+ "gpu_batching.scenario_scope must be 'sample' or 'batch'"
434
+ )
435
+
436
+ if "condition_parameter_scope" in raw:
437
+ condition_scope = str(raw["condition_parameter_scope"]).lower()
438
+ elif "sample_condition_once_per_batch" in raw:
439
+ condition_scope = (
440
+ "batch" if bool(raw["sample_condition_once_per_batch"]) else "sample"
441
+ )
442
+ else:
443
+ condition_scope = "batch" if scenario_scope == "batch" else "sample"
444
+ if condition_scope not in _GPU_BATCH_SCOPE_VALUES:
445
+ raise ValueError(
446
+ "gpu_batching.condition_parameter_scope must be 'sample' or 'batch'"
447
+ )
448
+ if condition_scope == "batch" and scenario_scope != "batch":
449
+ raise ValueError(
450
+ "gpu_batching.condition_parameter_scope='batch' requires "
451
+ "scenario_scope='batch'"
452
+ )
453
+ return scenario_scope, condition_scope
454
+
383
455
  def _sample_scenario_profile(
384
456
  self,
385
457
  rng: np.random.Generator,
@@ -466,25 +538,60 @@ class FogTransform(Transform):
466
538
  self,
467
539
  rng: np.random.Generator,
468
540
  augmentation: FogAugmentationSpec | None = None,
541
+ *,
542
+ scenario: dict[str, Any] | None = None,
543
+ sample_scenario: bool = True,
544
+ freeze_sampled_parameters: bool = False,
469
545
  ) -> RenderPlan:
470
- scenario = self._sample_scenario_profile(rng)
546
+ if sample_scenario:
547
+ scenario = self._sample_scenario_profile(rng)
471
548
  if scenario is None:
549
+ effective_config = (
550
+ _freeze_distribution_specs(self.config, rng)
551
+ if freeze_sampled_parameters
552
+ else self.config
553
+ )
554
+ models_cfg = (
555
+ effective_config.get("models")
556
+ or effective_config.get("fog_models")
557
+ or {}
558
+ )
472
559
  if augmentation is not None:
473
- model_name, model_cfg = self._resolve_augmented_model(augmentation)
560
+ base_cfg = resolve_model_config(augmentation.model_name, models_cfg)
561
+ model_cfg = deep_merge(base_cfg, augmentation.model_overrides)
562
+ if freeze_sampled_parameters:
563
+ model_cfg = _freeze_distribution_specs(model_cfg, rng)
474
564
  return RenderPlan(
475
- model_name=model_name,
565
+ model_name=augmentation.model_name,
476
566
  model_cfg=model_cfg,
477
567
  airlight_method=augmentation.airlight_method,
478
- )
479
- model_name = select_model(self.config, rng)
480
- model_cfg = resolve_model_config(model_name, self.models_cfg)
481
- return RenderPlan(model_name=model_name, model_cfg=model_cfg)
568
+ capture_artifacts=(
569
+ CaptureArtifactPipeline.from_config(effective_config)
570
+ if freeze_sampled_parameters
571
+ else None
572
+ ),
573
+ )
574
+ model_name = select_model(effective_config, rng)
575
+ model_cfg = resolve_model_config(model_name, models_cfg)
576
+ if freeze_sampled_parameters:
577
+ model_cfg = _freeze_distribution_specs(model_cfg, rng)
578
+ return RenderPlan(
579
+ model_name=model_name,
580
+ model_cfg=model_cfg,
581
+ capture_artifacts=(
582
+ CaptureArtifactPipeline.from_config(effective_config)
583
+ if freeze_sampled_parameters
584
+ else None
585
+ ),
586
+ )
482
587
 
483
588
  payload = self._scenario_payload(scenario)
484
589
  effective_config = deep_merge(
485
590
  self.config,
486
591
  self._scenario_config_override(payload),
487
592
  )
593
+ if freeze_sampled_parameters:
594
+ effective_config = _freeze_distribution_specs(effective_config, rng)
488
595
  models_cfg = (
489
596
  effective_config.get("models")
490
597
  or effective_config.get("fog_models")
@@ -505,6 +612,8 @@ class FogTransform(Transform):
505
612
  model_cfg = resolve_model_config(model_name, models_cfg)
506
613
  model_cfg = deep_merge(model_cfg, scenario_overrides)
507
614
  airlight_method = self._scenario_airlight_method(payload)
615
+ if freeze_sampled_parameters:
616
+ model_cfg = _freeze_distribution_specs(model_cfg, rng)
508
617
 
509
618
  return RenderPlan(
510
619
  model_name=model_name,
@@ -514,6 +623,20 @@ class FogTransform(Transform):
514
623
  scenario_name=self._scenario_name(scenario),
515
624
  )
516
625
 
626
+ def _resolve_gpu_batch_render_plan(self, batch_index: int) -> RenderPlan | None:
627
+ if self.gpu_scenario_scope != "batch":
628
+ return None
629
+ rng = self._rng_for_batch(batch_index)
630
+ scenario = self._sample_scenario_profile(rng)
631
+ return self._resolve_render_plan(
632
+ rng,
633
+ scenario=scenario,
634
+ sample_scenario=False,
635
+ freeze_sampled_parameters=(
636
+ self.gpu_condition_parameter_scope == "batch"
637
+ ),
638
+ )
639
+
517
640
  def _source_extension(self, sample: dict, backend: Any | None = None) -> str:
518
641
  meta = sample.get("meta")
519
642
  source_modality = (
@@ -977,7 +1100,10 @@ class FogTransform(Transform):
977
1100
  saved_paths: list[Path] = []
978
1101
 
979
1102
  with progress_bar(total, "GPU", self.logger) as bar:
980
- for batch in iter_batches(enumerate(samples), self.gpu_batch_size):
1103
+ for batch_index, batch in enumerate(
1104
+ iter_batches(enumerate(samples), self.gpu_batch_size)
1105
+ ):
1106
+ batch_plan = self._resolve_gpu_batch_render_plan(batch_index)
981
1107
  items: list[dict] = []
982
1108
  for global_index, sample in batch:
983
1109
  rgb = _to_numpy(sample["rgb"])
@@ -993,7 +1119,7 @@ class FogTransform(Transform):
993
1119
  )
994
1120
  else:
995
1121
  rng = self.base_rng
996
- plan = self._resolve_render_plan(rng)
1122
+ plan = batch_plan or self._resolve_render_plan(rng)
997
1123
  items.append(
998
1124
  {
999
1125
  "sample_id": sample["id"],
@@ -1022,22 +1148,19 @@ class FogTransform(Transform):
1022
1148
  grouped.setdefault(shape, []).append(item)
1023
1149
 
1024
1150
  for group_items in grouped.values():
1025
- if self.scenario_profiles:
1026
- uniform_items = []
1027
- other_items = list(group_items)
1028
- else:
1029
- uniform_items = [
1030
- item
1031
- for item in group_items
1032
- if item["model_name"] == "uniform"
1033
- ]
1034
- other_items = [
1035
- item
1036
- for item in group_items
1037
- if item["model_name"] != "uniform"
1038
- ]
1039
-
1040
- if uniform_items:
1151
+ uniform_groups: dict[int, list[dict]] = {}
1152
+ other_items = []
1153
+ for item in group_items:
1154
+ if item["model_name"] != "uniform":
1155
+ other_items.append(item)
1156
+ continue
1157
+ capture_artifacts = item.get("capture_artifacts")
1158
+ capture_key = (
1159
+ 0 if capture_artifacts is None else id(capture_artifacts)
1160
+ )
1161
+ uniform_groups.setdefault(capture_key, []).append(item)
1162
+
1163
+ for uniform_items in uniform_groups.values():
1041
1164
  rgb_batch = torch.stack(
1042
1165
  [
1043
1166
  normalize_rgb_torch(item["rgb"], device)
@@ -1073,6 +1196,7 @@ class FogTransform(Transform):
1073
1196
  contrast_threshold_default=(
1074
1197
  self.contrast_threshold_default
1075
1198
  ),
1199
+ method=uniform_items[0].get("airlight_method"),
1076
1200
  )
1077
1201
  )
1078
1202
  k_tensor = torch.tensor(
@@ -1082,6 +1206,7 @@ class FogTransform(Transform):
1082
1206
  foggy = rgb_batch * t[..., None] + ls_base[
1083
1207
  :, None, None, :
1084
1208
  ] * (1.0 - t[..., None])
1209
+ capture_artifacts = uniform_items[0].get("capture_artifacts")
1085
1210
  foggy = self.pipeline.apply_capture_torch_batch(
1086
1211
  foggy,
1087
1212
  contexts=tuple(
@@ -1100,6 +1225,7 @@ class FogTransform(Transform):
1100
1225
  )
1101
1226
  for idx, item in enumerate(uniform_items)
1102
1227
  ),
1228
+ capture_artifacts=capture_artifacts,
1103
1229
  )
1104
1230
 
1105
1231
  height = int(rgb_batch.shape[1])
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: euler-preprocess
3
- Version: 3.1.0
3
+ Version: 3.2.0
4
4
  Summary: Physics-based preprocessing (fog, etc.) for RGB+depth datasets
5
5
  Requires-Python: >=3.9
6
6
  Description-Content-Type: text/markdown
@@ -30,6 +30,7 @@ euler_preprocess/fog/dcp_heuristic_airlight.py
30
30
  euler_preprocess/fog/dcp_heuristic_airlight_torch.py
31
31
  euler_preprocess/fog/foggify.py
32
32
  euler_preprocess/fog/foggify_logging.py
33
+ euler_preprocess/fog/inference.py
33
34
  euler_preprocess/fog/logging.py
34
35
  euler_preprocess/fog/models.py
35
36
  euler_preprocess/fog/pipeline.py
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "euler-preprocess"
3
- version = "3.1.0"
3
+ version = "3.2.0"
4
4
  description = "Physics-based preprocessing (fog, etc.) for RGB+depth datasets"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.9"
@@ -27,6 +27,7 @@ from euler_loading.loaders.cpu.generic import (
27
27
  from euler_preprocess.common.output import prepare_output_backends
28
28
  import euler_preprocess.fog.capture as capture_module
29
29
  from euler_preprocess.fog.capture import CaptureArtifactPipeline, CaptureContext
30
+ from euler_preprocess.fog.inference import render_fog_image, render_fog_sample
30
31
  from euler_preprocess.fog.models import visibility_to_k
31
32
  from euler_preprocess.fog.transform import (
32
33
  ATMOSPHERIC_LIGHT_SLOT,
@@ -1234,6 +1235,174 @@ def test_scenario_profile_correlates_model_and_capture_overrides(
1234
1235
  np.testing.assert_allclose(result.rgb, 0.2)
1235
1236
 
1236
1237
 
1238
+ def test_gpu_batching_freezes_batch_condition_parameters(
1239
+ tmp_path: Path,
1240
+ ) -> None:
1241
+ cfg = {
1242
+ "airlight": "from_sky",
1243
+ "device": "cpu",
1244
+ "seed": 11,
1245
+ "gpu_batch_size": 2,
1246
+ "gpu_batching": {"scenario_scope": "batch"},
1247
+ "capture": {
1248
+ "stages": [
1249
+ {
1250
+ "type": "sensor",
1251
+ "enabled": True,
1252
+ "iso": {"dist": "uniform", "min": 800.0, "max": 1600.0},
1253
+ "shot_noise_scale": {
1254
+ "dist": "uniform",
1255
+ "min": 0.1,
1256
+ "max": 0.2,
1257
+ },
1258
+ }
1259
+ ]
1260
+ },
1261
+ "scenario_profiles": [
1262
+ {
1263
+ "name": "batch_dense",
1264
+ "weight": 1.0,
1265
+ "model": "uniform",
1266
+ "models": {
1267
+ "uniform": {
1268
+ "visibility_m": {
1269
+ "dist": "uniform",
1270
+ "min": 20.0,
1271
+ "max": 40.0,
1272
+ },
1273
+ "atmospheric_light": [0.2, 0.2, 0.2],
1274
+ }
1275
+ },
1276
+ "capture_overrides": {
1277
+ "sensor": {
1278
+ "read_noise_electrons": {
1279
+ "dist": "uniform",
1280
+ "min": 2.0,
1281
+ "max": 4.0,
1282
+ }
1283
+ }
1284
+ },
1285
+ }
1286
+ ],
1287
+ }
1288
+ config_path = tmp_path / "gpu_batching_config.json"
1289
+ config_path.write_text(json.dumps(cfg))
1290
+ transform = FogTransform(
1291
+ config_path=str(config_path),
1292
+ out_path=str(tmp_path / "out"),
1293
+ )
1294
+
1295
+ plan = transform._resolve_gpu_batch_render_plan(0)
1296
+
1297
+ assert plan is not None
1298
+ assert transform.gpu_scenario_scope == "batch"
1299
+ assert transform.gpu_condition_parameter_scope == "batch"
1300
+ assert plan.scenario_name == "batch_dense"
1301
+ assert isinstance(plan.model_cfg["visibility_m"], float)
1302
+ assert 20.0 <= plan.model_cfg["visibility_m"] <= 40.0
1303
+ assert plan.capture_artifacts is not None
1304
+ sensor_cfg = plan.capture_artifacts.stages[0].config
1305
+ assert sensor_cfg["enabled"] is True
1306
+ assert isinstance(sensor_cfg["iso"], float)
1307
+ assert isinstance(sensor_cfg["shot_noise_scale"], float)
1308
+ assert isinstance(sensor_cfg["read_noise_electrons"], float)
1309
+
1310
+
1311
+ def test_gpu_batch_scenario_keeps_uniform_batch_path(
1312
+ tmp_path: Path,
1313
+ monkeypatch: pytest.MonkeyPatch,
1314
+ ) -> None:
1315
+ torch = pytest.importorskip("torch")
1316
+ cfg = {
1317
+ "airlight": "from_sky",
1318
+ "device": "cpu",
1319
+ "seed": 13,
1320
+ "gpu_batch_size": 2,
1321
+ "gpu_batching": {
1322
+ "scenario_scope": "batch",
1323
+ "condition_parameter_scope": "batch",
1324
+ },
1325
+ "scenario_profiles": [
1326
+ {
1327
+ "name": "batch_uniform",
1328
+ "weight": 1.0,
1329
+ "model": "uniform",
1330
+ "airlight_method": "dcp_heuristic",
1331
+ "models": {
1332
+ "uniform": {
1333
+ "visibility_m": 80.0,
1334
+ "atmospheric_light": [0.35, 0.35, 0.35],
1335
+ }
1336
+ },
1337
+ "capture": {"stages": []},
1338
+ }
1339
+ ],
1340
+ }
1341
+ config_path = tmp_path / "gpu_batch_uniform_config.json"
1342
+ config_path.write_text(json.dumps(cfg))
1343
+ transform = FogTransform(
1344
+ config_path=str(config_path),
1345
+ out_path=str(tmp_path / "out"),
1346
+ )
1347
+ transform.torch_device = torch.device("cpu")
1348
+
1349
+ calls = []
1350
+ original = transform.pipeline.apply_capture_torch_batch
1351
+ resolve_calls = []
1352
+ original_resolve = transform.atmospheric_light.resolve_uniform_batch_torch
1353
+
1354
+ def spy_apply_capture_torch_batch(
1355
+ rgb_batch,
1356
+ *,
1357
+ contexts,
1358
+ capture_artifacts=None,
1359
+ ):
1360
+ calls.append((len(contexts), capture_artifacts))
1361
+ return original(
1362
+ rgb_batch,
1363
+ contexts=contexts,
1364
+ capture_artifacts=capture_artifacts,
1365
+ )
1366
+
1367
+ def spy_resolve_uniform_batch_torch(**kwargs):
1368
+ resolve_calls.append(kwargs.get("method"))
1369
+ return original_resolve(**kwargs)
1370
+
1371
+ monkeypatch.setattr(
1372
+ transform.pipeline,
1373
+ "apply_capture_torch_batch",
1374
+ spy_apply_capture_torch_batch,
1375
+ )
1376
+ monkeypatch.setattr(
1377
+ transform.atmospheric_light,
1378
+ "resolve_uniform_batch_torch",
1379
+ spy_resolve_uniform_batch_torch,
1380
+ )
1381
+
1382
+ samples = [
1383
+ {
1384
+ "id": "sample_a",
1385
+ "rgb": np.full((4, 5, 3), 0.6, dtype=np.float32),
1386
+ "depth": np.full((4, 5), 10.0, dtype=np.float32),
1387
+ "semantic_segmentation": np.ones((4, 5), dtype=bool),
1388
+ },
1389
+ {
1390
+ "id": "sample_b",
1391
+ "rgb": np.full((4, 5, 3), 0.4, dtype=np.float32),
1392
+ "depth": np.full((4, 5), 12.0, dtype=np.float32),
1393
+ "semantic_segmentation": np.ones((4, 5), dtype=bool),
1394
+ },
1395
+ ]
1396
+
1397
+ saved = transform._generate_fog_gpu(samples)
1398
+
1399
+ assert len(saved) == 2
1400
+ assert len(calls) == 1
1401
+ assert calls[0][0] == 2
1402
+ assert calls[0][1] is not None
1403
+ assert resolve_calls == ["dcp_heuristic"]
1404
+
1405
+
1237
1406
  def test_capture_camera_profile_supplies_stage_defaults() -> None:
1238
1407
  pipeline = CaptureArtifactPipeline.from_config(
1239
1408
  {
@@ -1257,6 +1426,89 @@ def test_capture_camera_profile_supplies_stage_defaults() -> None:
1257
1426
  assert pipeline.stages[0].config["read_noise_electrons"] == 3.0
1258
1427
 
1259
1428
 
1429
+ def test_render_fog_sample_accepts_modalities_and_named_scenario() -> None:
1430
+ config = {
1431
+ "airlight": "from_sky",
1432
+ "device": "cpu",
1433
+ "seed": 17,
1434
+ "capture": {
1435
+ "stages": [
1436
+ {
1437
+ "type": "exposure",
1438
+ "gain": 1.0,
1439
+ "condition_profiles": [
1440
+ {"name": "clean", "weight": 1.0, "gain": 1.0},
1441
+ {"name": "dark", "weight": 0.0, "gain": 0.25},
1442
+ ],
1443
+ }
1444
+ ]
1445
+ },
1446
+ "scenario_profiles": [
1447
+ {
1448
+ "name": "clean_reference",
1449
+ "weight": 1.0,
1450
+ "model": "uniform",
1451
+ "models": {
1452
+ "uniform": {
1453
+ "visibility_m": 30.0,
1454
+ "atmospheric_light": [0.2, 0.2, 0.2],
1455
+ }
1456
+ },
1457
+ "capture_overrides": {"exposure": {"condition_profile": "clean"}},
1458
+ },
1459
+ {
1460
+ "name": "dark_reference",
1461
+ "weight": 1.0,
1462
+ "model": "uniform",
1463
+ "models": {
1464
+ "uniform": {
1465
+ "visibility_m": 30.0,
1466
+ "atmospheric_light": [0.2, 0.2, 0.2],
1467
+ }
1468
+ },
1469
+ "capture_overrides": {"exposure": {"condition_profile": "dark"}},
1470
+ },
1471
+ ],
1472
+ }
1473
+ rgb = np.full((4, 5, 3), 0.8, dtype=np.float32)
1474
+ depth = np.zeros((4, 5), dtype=np.float32)
1475
+ semantic = np.zeros((4, 5, 3), dtype=np.uint8)
1476
+ semantic[:2, :, :] = np.array([29, 0, 0], dtype=np.uint8)
1477
+ intrinsics = np.array(
1478
+ [[100.0, 0.0, 2.0], [0.0, 100.0, 1.5], [0.0, 0.0, 1.0]],
1479
+ dtype=np.float32,
1480
+ )
1481
+
1482
+ result = render_fog_sample(
1483
+ rgb=rgb,
1484
+ depth=depth,
1485
+ semantic_segmentation=semantic,
1486
+ intrinsics=intrinsics,
1487
+ config=config,
1488
+ scenario_profile_name="dark_reference",
1489
+ mode="cpu",
1490
+ )
1491
+ image = render_fog_image(
1492
+ rgb=rgb,
1493
+ depth=depth,
1494
+ semantic_segmentation=semantic,
1495
+ intrinsics={
1496
+ "fx": 100.0,
1497
+ "fy": 100.0,
1498
+ "cx": 2.0,
1499
+ "cy": 1.5,
1500
+ },
1501
+ config=config,
1502
+ scenario_profile_name="dark_reference",
1503
+ mode="cpu",
1504
+ )
1505
+
1506
+ assert result.scenario_name == "dark_reference"
1507
+ assert result.model_name == "uniform"
1508
+ np.testing.assert_allclose(result.rgb, 0.2)
1509
+ np.testing.assert_allclose(image, result.rgb)
1510
+
1511
+
1260
1512
  def test_optics_stage_uses_intrinsics_principal_point() -> None:
1261
1513
  pipeline = CaptureArtifactPipeline.from_config(
1262
1514
  {