euler-loading 2.10.0__tar.gz → 2.11.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_loading-2.10.0 → euler_loading-2.11.0}/PKG-INFO +1 -1
  2. {euler_loading-2.10.0 → euler_loading-2.11.0}/README.md +17 -0
  3. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/cpu/muses.py +79 -0
  4. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/generate/loaders.json +26 -0
  5. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/gpu/muses.py +80 -0
  6. {euler_loading-2.10.0 → euler_loading-2.11.0}/pyproject.toml +1 -1
  7. {euler_loading-2.10.0 → euler_loading-2.11.0}/tests/test_loaders.py +62 -0
  8. {euler_loading-2.10.0 → euler_loading-2.11.0}/.github/workflows/workflow.yml +0 -0
  9. {euler_loading-2.10.0 → euler_loading-2.11.0}/.gitignore +0 -0
  10. {euler_loading-2.10.0 → euler_loading-2.11.0}/docs/loader-attributes.md +0 -0
  11. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/__init__.py +0 -0
  12. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/_dataset_contract.py +0 -0
  13. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/_ds_crawler_utils.py +0 -0
  14. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/_metadata.py +0 -0
  15. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/_resolution.py +0 -0
  16. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/_writing.py +0 -0
  17. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/dataset.py +0 -0
  18. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/indexing.py +0 -0
  19. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/__init__.py +0 -0
  20. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/_annotations.py +0 -0
  21. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/_writer_utils.py +0 -0
  22. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/contracts.py +0 -0
  23. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/cpu/__init__.py +0 -0
  24. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/cpu/generic.py +0 -0
  25. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/cpu/generic_dense_depth.py +0 -0
  26. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/cpu/real_drive_sim.py +0 -0
  27. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/cpu/vkitti2.py +0 -0
  28. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/generate/__init__.py +0 -0
  29. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/generate/__main__.py +0 -0
  30. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/generic.py +0 -0
  31. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/gpu/__init__.py +0 -0
  32. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/gpu/generic.py +0 -0
  33. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/gpu/generic_dense_depth.py +0 -0
  34. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/gpu/real_drive_sim.py +0 -0
  35. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/gpu/vkitti2.py +0 -0
  36. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/muses.py +0 -0
  37. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/real_drive_sim.py +0 -0
  38. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/loaders/vkitti2.py +0 -0
  39. {euler_loading-2.10.0 → euler_loading-2.11.0}/euler_loading/preprocessing.py +0 -0
  40. {euler_loading-2.10.0 → euler_loading-2.11.0}/example.py +0 -0
  41. {euler_loading-2.10.0 → euler_loading-2.11.0}/package-lock.json +0 -0
  42. {euler_loading-2.10.0 → euler_loading-2.11.0}/sample_rds.py +0 -0
  43. {euler_loading-2.10.0 → euler_loading-2.11.0}/tests/__init__.py +0 -0
  44. {euler_loading-2.10.0 → euler_loading-2.11.0}/tests/conftest.py +0 -0
  45. {euler_loading-2.10.0 → euler_loading-2.11.0}/tests/example_rds_calib.json +0 -0
  46. {euler_loading-2.10.0 → euler_loading-2.11.0}/tests/test_dataset.py +0 -0
  47. {euler_loading-2.10.0 → euler_loading-2.11.0}/tests/test_id_schema.py +0 -0
  48. {euler_loading-2.10.0 → euler_loading-2.11.0}/tests/test_indexing.py +0 -0
  49. {euler_loading-2.10.0 → euler_loading-2.11.0}/tests/test_preprocessing.py +0 -0
  50. {euler_loading-2.10.0 → euler_loading-2.11.0}/tests/test_real_dataset.py +0 -0
  51. {euler_loading-2.10.0 → euler_loading-2.11.0}/tests/test_writing.py +0 -0
  52. {euler_loading-2.10.0 → euler_loading-2.11.0}/vkitti_cpu_example_output.json +0 -0
  53. {euler_loading-2.10.0 → euler_loading-2.11.0}/vkitti_gpu_example_output.json +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: euler-loading
3
- Version: 2.10.0
3
+ Version: 2.11.0
4
4
  Summary: Multi-modal PyTorch dataloader using ds-crawler indices
5
5
  Requires-Python: >=3.9
6
6
  Requires-Dist: ds-crawler>=2.3.0
@@ -193,6 +193,17 @@ PyTorch `Dataset`. On construction it:
193
193
  | `hierarchical_modalities` | `dict[str, Modality] \| None` | Optional modalities whose files live at intermediate hierarchy levels (e.g. per-scene intrinsics). These do **not** participate in ID intersection. Each sample will contain a dict `{file_id: loaded_result}` with all files at or above the sample's hierarchy level. Results are cached so shared files are parsed only once. |
194
194
  | `transforms` | `list[Callable[[dict], dict]] \| None` | Applied in order after loading. Each receives and returns the full sample dict. |
195
195
 
196
+ For hierarchical modalities, the ds-crawler file ID is the key in
197
+ `{file_id: loaded_result}`. If files at different ancestor levels use the same
198
+ ID, the deepest matching file wins. This gives calibration files natural
199
+ inheritance semantics: a root-level `intrinsics` file can apply to every sample,
200
+ while a scene- or camera-level `intrinsics` file with the same ID overrides it
201
+ for descendants. Use ds-crawler's `indexing.id.override` when the physical file
202
+ name is incidental, such as `calib.json` or a UUID, and you want a stable
203
+ semantic key like `intrinsics`, `extrinsics`, or `calibration`. Do not use an
204
+ override when multiple distinct files at the same hierarchy level should all be
205
+ returned.
206
+
196
207
  #### Sample dict
197
208
 
198
209
  `dataset[i]` returns:
@@ -361,6 +372,11 @@ Each logical modality must have its own ds-crawler index. Usually that means one
361
372
  Files across regular modalities are matched by IDs from those indexes, so **the indexed hierarchy and naming conventions must be consistent** across modalities up to modality-specific parts captured in the config.
362
373
 
363
374
  Calibration files or other per-scene/per-sequence metadata can be loaded via `hierarchical_modalities`. These files are matched to samples based on their position in the hierarchy — all files at or above a sample's hierarchy level are included and cached for efficiency.
375
+ Root-level calibration files, for example `calib.json` at the top of a shared
376
+ zip, should be indexed by ds-crawler without a hierarchy block. euler-loading
377
+ treats that root file as an ancestor of every sample. If the modality uses
378
+ `collapse_single=True`, the loaded calibration value is returned directly as
379
+ long as exactly one hierarchical file matches the sample.
364
380
 
365
381
  ## Layout-aware augmentation loading
366
382
 
@@ -484,6 +500,7 @@ All built-in loaders accept both filesystem paths (`str`) and in-memory buffers
484
500
  | `reference_rgb` | Clear-weather reference RGB image as float32, normalised to [0, 1] |
485
501
  | `semantic_segmentation` | Single-channel Cityscapes `labelIds` or `labelTrainIds` PNG as class IDs |
486
502
  | `semantic_segmentation_color` | Cityscapes `labelColor` PNG as RGB uint8 labels |
503
+ | `sky_mask` | Binary mask from Cityscapes labels; uses sky `trainId=10` for `labelTrainIds`, `labelId=23` for `labelIds`, or RGB `meta["sky_class"]` such as `[0, 0, 23]` |
487
504
  | `panoptic_segmentation` | COCO-style RGB panoptic PNG decoded to integer segment IDs |
488
505
  | `lidar_point_cloud` / `point_cloud` | Lidar `.bin` file as `(N, 6)` float64: `x, y, z, intensity, ring, timestamp` |
489
506
  | `sparse_depth` | Alias for MUSES lidar points for sparse-depth style workflows |
@@ -12,6 +12,7 @@ Return types
12
12
  containing raw Cityscapes ``labelIds`` or ``labelTrainIds``.
13
13
  - **semantic_segmentation_color** -- ``np.ndarray`` of shape ``(H, W, 3)``
14
14
  uint8 containing Cityscapes RGB colours.
15
+ - **sky_mask** -- ``np.ndarray`` of shape ``(H, W)`` bool.
15
16
  - **panoptic_segmentation** -- ``np.ndarray`` of shape ``(H, W)`` int64
16
17
  containing COCO-style panoptic segment IDs decoded from RGB.
17
18
  - **lidar_point_cloud** / **sparse_depth** -- ``np.ndarray`` of shape
@@ -35,6 +36,8 @@ from PIL import Image
35
36
  from euler_loading.loaders._annotations import modality_meta
36
37
 
37
38
  _LIDAR_COLUMNS = ["x", "y", "z", "intensity", "ring", "timestamp"]
39
+ _SKY_LABEL_ID = 23
40
+ _SKY_TRAIN_ID = 10
38
41
 
39
42
 
40
43
  def _load_image_rgb(path: Union[str, BinaryIO]) -> np.ndarray:
@@ -56,6 +59,59 @@ def _load_single_channel_labels(path: Union[str, BinaryIO]) -> np.ndarray:
56
59
  return arr
57
60
 
58
61
 
62
+ def _select_sky_class_id(
63
+ path: Union[str, BinaryIO],
64
+ meta: Mapping[str, Any] | None,
65
+ attributes: Mapping[str, Any] | None,
66
+ ) -> int:
67
+ for source in (attributes, meta):
68
+ if not isinstance(source, Mapping):
69
+ continue
70
+ value = source.get("sky_class_id")
71
+ if isinstance(value, (int, np.integer)):
72
+ return int(value)
73
+
74
+ name = str(getattr(path, "name", path)).lower()
75
+ if "labelids" in name:
76
+ return _SKY_LABEL_ID
77
+ return _SKY_TRAIN_ID
78
+
79
+
80
+ def _select_sky_class_color(
81
+ meta: Mapping[str, Any] | None,
82
+ attributes: Mapping[str, Any] | None,
83
+ ) -> np.ndarray | None:
84
+ for source in (attributes, meta):
85
+ if not isinstance(source, Mapping):
86
+ continue
87
+ value = source.get("sky_class")
88
+ if value is None:
89
+ value = source.get("sky_mask")
90
+ if value is None:
91
+ continue
92
+
93
+ color = np.asarray(value, dtype=np.uint8)
94
+ if color.shape != (3,):
95
+ raise ValueError("MUSES sky_class must be a 3-element RGB value")
96
+ return color
97
+ return None
98
+
99
+
100
+ def _load_sky_mask_array(
101
+ path: Union[str, BinaryIO],
102
+ meta: Mapping[str, Any] | None,
103
+ attributes: Mapping[str, Any] | None,
104
+ ) -> np.ndarray:
105
+ color = _select_sky_class_color(meta, attributes)
106
+ if color is not None:
107
+ with Image.open(path) as image:
108
+ arr = np.array(image.convert("RGB"), dtype=np.uint8)
109
+ return np.all(arr == color, axis=-1)
110
+
111
+ labels = _load_single_channel_labels(path)
112
+ return labels == _select_sky_class_id(path, meta, attributes)
113
+
114
+
59
115
  def _load_rgb_uint8(path: Union[str, BinaryIO]) -> np.ndarray:
60
116
  """Load an RGB image without normalisation."""
61
117
  with Image.open(path) as image:
@@ -246,6 +302,29 @@ def semantic_segmentation(path: Union[str, BinaryIO], meta: dict[str, Any] | Non
246
302
  return _load_single_channel_labels(path)
247
303
 
248
304
 
305
+ @modality_meta(
306
+ modality_type="sky_mask",
307
+ dtype="bool",
308
+ shape="HW",
309
+ file_formats=[".png"],
310
+ meta={
311
+ "encoding": "single_channel",
312
+ "sky_class": [0, 0, _SKY_LABEL_ID],
313
+ "sky_label_id": _SKY_LABEL_ID,
314
+ "sky_train_id": _SKY_TRAIN_ID,
315
+ },
316
+ )
317
+ def sky_mask(path: Union[str, BinaryIO], meta: dict[str, Any] | None = None, *, attributes: dict[str, Any] | None = None) -> np.ndarray:
318
+ """Load a MUSES sky mask as an ``(H, W)`` bool array.
319
+
320
+ Reads a single-channel Cityscapes ``labelIds`` / ``labelTrainIds`` PNG, or
321
+ an RGB-encoded class image when ``sky_class`` is provided. Override the
322
+ default ID with ``sky_class_id``, or set ``sky_class`` to an RGB value such
323
+ as ``[0, 0, 23]``.
324
+ """
325
+ return _load_sky_mask_array(path, meta, attributes)
326
+
327
+
249
328
  @modality_meta(
250
329
  modality_type="semantic_segmentation_color",
251
330
  dtype="uint8",
@@ -396,6 +396,32 @@
396
396
  "shape": "CHW"
397
397
  }
398
398
  },
399
+ {
400
+ "type": "sky_mask",
401
+ "function": "sky_mask",
402
+ "hierarchical": false,
403
+ "dtype": "bool",
404
+ "access": "key",
405
+ "file_formats": [
406
+ ".png"
407
+ ],
408
+ "meta": {
409
+ "encoding": "single_channel",
410
+ "sky_class": [
411
+ 0,
412
+ 0,
413
+ 23
414
+ ],
415
+ "sky_label_id": 23,
416
+ "sky_train_id": 10
417
+ },
418
+ "cpu": {
419
+ "shape": "HW"
420
+ },
421
+ "gpu": {
422
+ "shape": "1HW"
423
+ }
424
+ },
399
425
  {
400
426
  "type": "sparse_depth",
401
427
  "function": "sparse_depth",
@@ -12,6 +12,7 @@ Return types
12
12
  containing raw Cityscapes ``labelIds`` or ``labelTrainIds``.
13
13
  - **semantic_segmentation_color** -- ``torch.ByteTensor`` of shape
14
14
  ``(3, H, W)`` containing Cityscapes RGB colours.
15
+ - **sky_mask** -- ``torch.BoolTensor`` of shape ``(1, H, W)``.
15
16
  - **panoptic_segmentation** -- ``torch.LongTensor`` of shape ``(1, H, W)``
16
17
  containing COCO-style panoptic segment IDs decoded from RGB.
17
18
  - **lidar_point_cloud** / **sparse_depth** -- ``torch.DoubleTensor`` of shape
@@ -36,6 +37,8 @@ from PIL import Image
36
37
  from euler_loading.loaders._annotations import modality_meta
37
38
 
38
39
  _LIDAR_COLUMNS = ["x", "y", "z", "intensity", "ring", "timestamp"]
40
+ _SKY_LABEL_ID = 23
41
+ _SKY_TRAIN_ID = 10
39
42
 
40
43
 
41
44
  def _load_image_rgb(path: Union[str, BinaryIO]) -> np.ndarray:
@@ -57,6 +60,59 @@ def _load_single_channel_labels(path: Union[str, BinaryIO]) -> np.ndarray:
57
60
  return arr
58
61
 
59
62
 
63
+ def _select_sky_class_id(
64
+ path: Union[str, BinaryIO],
65
+ meta: Mapping[str, Any] | None,
66
+ attributes: Mapping[str, Any] | None,
67
+ ) -> int:
68
+ for source in (attributes, meta):
69
+ if not isinstance(source, Mapping):
70
+ continue
71
+ value = source.get("sky_class_id")
72
+ if isinstance(value, (int, np.integer)):
73
+ return int(value)
74
+
75
+ name = str(getattr(path, "name", path)).lower()
76
+ if "labelids" in name:
77
+ return _SKY_LABEL_ID
78
+ return _SKY_TRAIN_ID
79
+
80
+
81
+ def _select_sky_class_color(
82
+ meta: Mapping[str, Any] | None,
83
+ attributes: Mapping[str, Any] | None,
84
+ ) -> np.ndarray | None:
85
+ for source in (attributes, meta):
86
+ if not isinstance(source, Mapping):
87
+ continue
88
+ value = source.get("sky_class")
89
+ if value is None:
90
+ value = source.get("sky_mask")
91
+ if value is None:
92
+ continue
93
+
94
+ color = np.asarray(value, dtype=np.uint8)
95
+ if color.shape != (3,):
96
+ raise ValueError("MUSES sky_class must be a 3-element RGB value")
97
+ return color
98
+ return None
99
+
100
+
101
+ def _load_sky_mask_array(
102
+ path: Union[str, BinaryIO],
103
+ meta: Mapping[str, Any] | None,
104
+ attributes: Mapping[str, Any] | None,
105
+ ) -> np.ndarray:
106
+ color = _select_sky_class_color(meta, attributes)
107
+ if color is not None:
108
+ with Image.open(path) as image:
109
+ arr = np.array(image.convert("RGB"), dtype=np.uint8)
110
+ return np.all(arr == color, axis=-1)
111
+
112
+ labels = _load_single_channel_labels(path)
113
+ return labels == _select_sky_class_id(path, meta, attributes)
114
+
115
+
60
116
  def _load_rgb_uint8(path: Union[str, BinaryIO]) -> np.ndarray:
61
117
  """Load an RGB image without normalisation."""
62
118
  with Image.open(path) as image:
@@ -250,6 +306,30 @@ def semantic_segmentation(path: Union[str, BinaryIO], meta: dict[str, Any] | Non
250
306
  return torch.from_numpy(arr).unsqueeze(0).contiguous()
251
307
 
252
308
 
309
+ @modality_meta(
310
+ modality_type="sky_mask",
311
+ dtype="bool",
312
+ shape="1HW",
313
+ file_formats=[".png"],
314
+ meta={
315
+ "encoding": "single_channel",
316
+ "sky_class": [0, 0, _SKY_LABEL_ID],
317
+ "sky_label_id": _SKY_LABEL_ID,
318
+ "sky_train_id": _SKY_TRAIN_ID,
319
+ },
320
+ )
321
+ def sky_mask(path: Union[str, BinaryIO], meta: dict[str, Any] | None = None, *, attributes: dict[str, Any] | None = None) -> torch.Tensor:
322
+ """Load a MUSES sky mask as a ``(1, H, W)`` bool tensor.
323
+
324
+ Reads a single-channel Cityscapes ``labelIds`` / ``labelTrainIds`` PNG, or
325
+ an RGB-encoded class image when ``sky_class`` is provided. Override the
326
+ default ID with ``sky_class_id``, or set ``sky_class`` to an RGB value such
327
+ as ``[0, 0, 23]``.
328
+ """
329
+ mask = _load_sky_mask_array(path, meta, attributes)
330
+ return torch.from_numpy(mask).unsqueeze(0).contiguous()
331
+
332
+
253
333
  @modality_meta(
254
334
  modality_type="semantic_segmentation_color",
255
335
  dtype="uint8",
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "euler-loading"
3
- version = "2.10.0"
3
+ version = "2.11.0"
4
4
  description = "Multi-modal PyTorch dataloader using ds-crawler indices"
5
5
  requires-python = ">=3.9"
6
6
  dependencies = [
@@ -158,6 +158,7 @@ MUSES_LOADER_NAMES = [
158
158
  "reference_rgb",
159
159
  "semantic_segmentation",
160
160
  "semantic_segmentation_color",
161
+ "sky_mask",
161
162
  "panoptic_segmentation",
162
163
  "lidar_point_cloud",
163
164
  "point_cloud",
@@ -189,6 +190,33 @@ def muses_semantic_path(tmp_path):
189
190
  return str(path), arr
190
191
 
191
192
 
193
+ @pytest.fixture()
194
+ def muses_sky_mask_path(tmp_path):
195
+ arr = np.array([[0, 10], [10, 255]], dtype=np.uint8)
196
+ path = tmp_path / "REC0001_frame_000001_gt_labelTrainIds.png"
197
+ Image.fromarray(arr, mode="L").save(path)
198
+ return str(path), arr == 10
199
+
200
+
201
+ @pytest.fixture()
202
+ def muses_sky_label_ids_path(tmp_path):
203
+ arr = np.array([[0, 23], [23, 255]], dtype=np.uint8)
204
+ path = tmp_path / "REC0001_frame_000001_gt_labelIds.png"
205
+ Image.fromarray(arr, mode="L").save(path)
206
+ return str(path), arr == 23
207
+
208
+
209
+ @pytest.fixture()
210
+ def muses_sky_rgb_class_path(tmp_path):
211
+ arr = np.zeros((2, 2, 3), dtype=np.uint8)
212
+ arr[0, 1] = [0, 0, 23]
213
+ arr[1, 0] = [0, 0, 23]
214
+ arr[1, 1] = [0, 0, 7]
215
+ path = tmp_path / "REC0001_frame_000001_gt_labelIds_rgb.png"
216
+ Image.fromarray(arr, mode="RGB").save(path)
217
+ return str(path), np.all(arr == np.array([0, 0, 23], dtype=np.uint8), axis=-1)
218
+
219
+
192
220
  @pytest.fixture()
193
221
  def muses_semantic_color_path(tmp_path):
194
222
  arr = np.array(
@@ -320,6 +348,23 @@ class TestMUSESGPULoaders:
320
348
  assert result.shape == (1, 2, 2)
321
349
  assert torch.equal(result, torch.from_numpy(expected.astype(np.int64)).unsqueeze(0))
322
350
 
351
+ def test_sky_mask_uses_train_ids(self, muses_sky_mask_path):
352
+ path, expected = muses_sky_mask_path
353
+ result = gpu_muses.sky_mask(path)
354
+ assert result.dtype == torch.bool
355
+ assert result.shape == (1, 2, 2)
356
+ assert torch.equal(result, torch.from_numpy(expected).unsqueeze(0))
357
+
358
+ def test_sky_mask_uses_label_ids_from_filename(self, muses_sky_label_ids_path):
359
+ path, expected = muses_sky_label_ids_path
360
+ result = gpu_muses.sky_mask(path)
361
+ assert torch.equal(result, torch.from_numpy(expected).unsqueeze(0))
362
+
363
+ def test_sky_mask_uses_rgb_sky_class_meta(self, muses_sky_rgb_class_path):
364
+ path, expected = muses_sky_rgb_class_path
365
+ result = gpu_muses.sky_mask(path, meta={"sky_class": [0, 0, 23]})
366
+ assert torch.equal(result, torch.from_numpy(expected).unsqueeze(0))
367
+
323
368
  def test_semantic_color_is_chw_uint8(self, muses_semantic_color_path):
324
369
  path, expected = muses_semantic_color_path
325
370
  result = gpu_muses.semantic_segmentation_color(path)
@@ -393,6 +438,23 @@ class TestMUSESCPULoaders:
393
438
  assert result.shape == (2, 2)
394
439
  assert np.array_equal(result, expected.astype(np.int64))
395
440
 
441
+ def test_sky_mask_uses_train_ids(self, muses_sky_mask_path):
442
+ path, expected = muses_sky_mask_path
443
+ result = cpu_muses.sky_mask(path)
444
+ assert result.dtype == np.bool_
445
+ assert result.shape == (2, 2)
446
+ assert np.array_equal(result, expected)
447
+
448
+ def test_sky_mask_uses_label_ids_from_filename(self, muses_sky_label_ids_path):
449
+ path, expected = muses_sky_label_ids_path
450
+ result = cpu_muses.sky_mask(path)
451
+ assert np.array_equal(result, expected)
452
+
453
+ def test_sky_mask_uses_rgb_sky_class_meta(self, muses_sky_rgb_class_path):
454
+ path, expected = muses_sky_rgb_class_path
455
+ result = cpu_muses.sky_mask(path, meta={"sky_class": [0, 0, 23]})
456
+ assert np.array_equal(result, expected)
457
+
396
458
  def test_semantic_color_is_hwc_uint8(self, muses_semantic_color_path):
397
459
  path, expected = muses_semantic_color_path
398
460
  result = cpu_muses.semantic_segmentation_color(path)