mvdata 0.9.0__tar.gz → 0.9.2__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 (60) hide show
  1. {mvdata-0.9.0 → mvdata-0.9.2}/PKG-INFO +1 -1
  2. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/__init__.py +0 -12
  3. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/__init__.py +6 -0
  4. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/decode.py +39 -5
  5. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/frames.py +104 -24
  6. mvdata-0.9.2/mvdata/codec/native_yuv.py +169 -0
  7. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/probe.py +25 -1
  8. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/video_stream_reader.py +26 -2
  9. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata.egg-info/PKG-INFO +1 -1
  10. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata.egg-info/SOURCES.txt +1 -2
  11. {mvdata-0.9.0 → mvdata-0.9.2}/pyproject.toml +1 -1
  12. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_video_stream_reader.py +205 -3
  13. mvdata-0.9.0/mvdata/multivideo_decode_benchmark.py +0 -216
  14. mvdata-0.9.0/tests/test_multivideo_decode_benchmark.py +0 -127
  15. {mvdata-0.9.0 → mvdata-0.9.2}/README.md +0 -0
  16. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/cloud_storage.py +0 -0
  17. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/_imports.py +0 -0
  18. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/encode.py +0 -0
  19. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/select.py +0 -0
  20. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/dataset_base.py +0 -0
  21. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/downloader.py +0 -0
  22. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/gpu_policy.py +0 -0
  23. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/gpu_support.py +0 -0
  24. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/image_metrics.py +0 -0
  25. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/legacy_writer.py +0 -0
  26. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/multivideo.py +0 -0
  27. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/multivideo_slicer.py +0 -0
  28. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/multivideo_writer.py +0 -0
  29. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/nvdec_parallel.py +0 -0
  30. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/nvenc_codec.py +0 -0
  31. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/per_frame.py +0 -0
  32. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/ranged.py +0 -0
  33. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/ranged_writer.py +0 -0
  34. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/stash_utils.py +0 -0
  35. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/utils.py +0 -0
  36. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/write_progress.py +0 -0
  37. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/writer_base.py +0 -0
  38. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata.egg-info/dependency_links.txt +0 -0
  39. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata.egg-info/requires.txt +0 -0
  40. {mvdata-0.9.0 → mvdata-0.9.2}/mvdata.egg-info/top_level.txt +0 -0
  41. {mvdata-0.9.0 → mvdata-0.9.2}/setup.cfg +0 -0
  42. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_dataset_base_defaults.py +0 -0
  43. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_gpu_policy.py +0 -0
  44. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_gpu_support.py +0 -0
  45. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_image_metrics.py +0 -0
  46. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_multivideo_bit_depth.py +0 -0
  47. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_multivideo_slicer.py +0 -0
  48. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_nvdec_parallel.py +0 -0
  49. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_per_camera.py +0 -0
  50. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_ranged_mixed_streams.py +0 -0
  51. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_ranged_nvenc_roundtrip.py +0 -0
  52. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_ranged_resume.py +0 -0
  53. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_ranged_stream_discovery.py +0 -0
  54. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_roundtrip.py +0 -0
  55. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_s3_downloader.py +0 -0
  56. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_stash_bit_depth.py +0 -0
  57. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_stash_comprehensive.py +0 -0
  58. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_stash_policy.py +0 -0
  59. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_stash_regenerate.py +0 -0
  60. {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_write_progress.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: mvdata
3
- Version: 0.9.0
3
+ Version: 0.9.2
4
4
  Summary: Gracia Dataset Convention - Python library for working with multi-view video datasets
5
5
  Author: Gracia Team
6
6
  License: MIT
@@ -37,13 +37,6 @@ from .gpu_support import (
37
37
  supports_gpu_decode,
38
38
  )
39
39
  from .image_metrics import calculate_psnr
40
- from .multivideo_decode_benchmark import (
41
- MultiVideoDecodeBenchmarkResult,
42
- NvdecRuntimeInfo,
43
- benchmark_multivideo_decode,
44
- inspect_nvdec_runtime,
45
- summarize_multivideo_dataset,
46
- )
47
40
 
48
41
 
49
42
  def _build_missing_class(class_name: str, install_hint: str, import_error: Exception):
@@ -138,11 +131,6 @@ __all__ = [
138
131
  'MultiVideoStreamEligibility',
139
132
  'MultiVideoStreamSliceInfo',
140
133
  'calculate_psnr',
141
- 'NvdecRuntimeInfo',
142
- 'MultiVideoDecodeBenchmarkResult',
143
- 'inspect_nvdec_runtime',
144
- 'summarize_multivideo_dataset',
145
- 'benchmark_multivideo_decode',
146
134
  'GpuDecodeMediaSupport',
147
135
  'GpuDecodeSupportReport',
148
136
  'gpu_decode_support',
@@ -5,6 +5,7 @@ This subpackage is split into small modules:
5
5
  - :mod:`mvdata.codec._imports` – lazy optional-dep imports and constants
6
6
  - :mod:`mvdata.codec.probe` – bit-depth probing and GPU capability checks
7
7
  - :mod:`mvdata.codec.frames` – decoded-frame → HWC RGB (numpy and cupy)
8
+ - :mod:`mvdata.codec.native_yuv` – native high-bit-depth NVDEC YUV → RGB
8
9
  - :mod:`mvdata.codec.encode` – RGB → NVENC elementary stream → MP4
9
10
  - :mod:`mvdata.codec.decode` – NVDEC / PyAV decode of MP4
10
11
  - :mod:`mvdata.codec.select` – round-trip PSNR and auto codec selection
@@ -57,6 +58,7 @@ from .frames import (
57
58
  _pyav_frame_to_rgb,
58
59
  numpy_to_cupy_rgb,
59
60
  )
61
+ from .native_yuv import native_nvdec_to_rgb_cupy, native_nvdec_to_rgb_numpy
60
62
  from .probe import (
61
63
  infer_video_bit_depth_from_frame,
62
64
  infer_video_bit_depth_from_pixel_format_name,
@@ -66,6 +68,7 @@ from .probe import (
66
68
  nvdec_decode_compatibility_issue_for_path,
67
69
  nvenc_encode_compatibility_issue,
68
70
  nvenc_encode_compatibility_issue_for_rgb,
71
+ probe_video_color_metadata_pyav,
69
72
  probe_video_bit_depth,
70
73
  probe_video_bit_depth_pyav,
71
74
  probe_video_stream_metadata,
@@ -99,6 +102,8 @@ __all__ = [
99
102
  "infer_video_bit_depth_from_video_format",
100
103
  "mux_elementary_to_mp4",
101
104
  "normalize_codec_name",
105
+ "native_nvdec_to_rgb_cupy",
106
+ "native_nvdec_to_rgb_numpy",
102
107
  "numpy_to_cupy_rgb",
103
108
  "nvdec_decode_compatibility_issue",
104
109
  "nvdec_decode_compatibility_issue_for_path",
@@ -106,6 +111,7 @@ __all__ = [
106
111
  "nvenc_encode_compatibility_issue_for_rgb",
107
112
  "pad_rgb_for_nvenc",
108
113
  "preferred_decoder_for_mp4",
114
+ "probe_video_color_metadata_pyav",
109
115
  "probe_video_bit_depth",
110
116
  "probe_video_bit_depth_pyav",
111
117
  "probe_video_stream_metadata",
@@ -10,20 +10,30 @@ import numpy as np
10
10
  from ..gpu_policy import nvdec_decode_allowed
11
11
  from ._imports import try_import_av, try_import_cupy, try_import_pynvvideocodec, try_import_torch
12
12
  from .frames import _decoded_frame_to_rgb_numpy, _pyav_frame_to_rgb
13
+ from .native_yuv import native_nvdec_to_rgb_numpy
13
14
  from .probe import (
14
15
  infer_video_bit_depth_from_stream,
15
16
  nvdec_decode_compatibility_issue,
16
17
  nvdec_decode_compatibility_issue_for_path,
18
+ probe_video_color_metadata_pyav,
19
+ probe_video_bit_depth,
17
20
  probe_video_stream_metadata,
18
21
  )
19
22
 
20
23
 
21
- def _try_open_nvdec(nvc, path_str: str, gpu_id: int, use_device_memory: bool):
24
+ def _try_open_nvdec(
25
+ nvc,
26
+ path_str: str,
27
+ gpu_id: int,
28
+ use_device_memory: bool,
29
+ *,
30
+ output_color_type=None,
31
+ ):
22
32
  """Open a SimpleDecoder, tolerating API-version differences in the kwarg surface."""
23
33
  base = dict(
24
34
  gpu_id=gpu_id,
25
35
  use_device_memory=use_device_memory,
26
- output_color_type=nvc.OutputColorType.RGB,
36
+ output_color_type=output_color_type or nvc.OutputColorType.RGB,
27
37
  )
28
38
  for extra in ({}, {"need_scanned_stream_metadata": True}):
29
39
  try:
@@ -69,14 +79,27 @@ def decode_mp4_to_rgb_nvdec(nvc, mp4_path: Path, gpu_id: int, expect_count: int)
69
79
  )
70
80
  if issue is not None:
71
81
  raise RuntimeError(f"NVDEC decode unsupported for {mp4_path}: {issue}")
82
+ source_bit_depth = probe_video_bit_depth(mp4_path, nvc=nvc)
83
+ color_metadata = probe_video_color_metadata_pyav(mp4_path)
72
84
 
73
85
  torch_mod = try_import_torch()
74
86
  has_dlpack = torch_mod is not None or try_import_cupy() is not None
75
87
  path_str = str(mp4_path)
88
+ output_color_type = nvc.OutputColorType.NATIVE if source_bit_depth > 8 else nvc.OutputColorType.RGB
89
+ use_dev_options = [True] if source_bit_depth > 8 else ([True, False] if has_dlpack else [False])
76
90
 
77
91
  last_err: Exception | None = None
78
- for use_dev in ([True, False] if has_dlpack else [False]):
79
- dec = _try_open_nvdec(nvc, path_str, gpu_id, use_dev)
92
+ for use_dev in use_dev_options:
93
+ if source_bit_depth > 8:
94
+ dec = _try_open_nvdec(
95
+ nvc,
96
+ path_str,
97
+ gpu_id,
98
+ use_dev,
99
+ output_color_type=output_color_type,
100
+ )
101
+ else:
102
+ dec = _try_open_nvdec(nvc, path_str, gpu_id, use_dev)
80
103
  if dec is None:
81
104
  continue
82
105
  try:
@@ -84,7 +107,18 @@ def decode_mp4_to_rgb_nvdec(nvc, mp4_path: Path, gpu_id: int, expect_count: int)
84
107
  frames: List[np.ndarray] = []
85
108
  for i in range(n):
86
109
  raw = dec.get_batch_frames_by_index([i])[0]
87
- frames.append(_decoded_frame_to_rgb_numpy(raw, torch_mod, bit_depth=bit_depth))
110
+ if source_bit_depth > 8:
111
+ frames.append(
112
+ native_nvdec_to_rgb_numpy(
113
+ raw,
114
+ bit_depth=source_bit_depth,
115
+ **color_metadata,
116
+ )
117
+ )
118
+ else:
119
+ frames.append(
120
+ _decoded_frame_to_rgb_numpy(raw, torch_mod, bit_depth=source_bit_depth)
121
+ )
88
122
  return frames
89
123
  except Exception as e:
90
124
  last_err = e
@@ -52,6 +52,10 @@ def _device_and_stream_ctx(
52
52
 
53
53
  def _target_dtype_for(arr_dtype, *, bit_depth: int | None):
54
54
  if arr_dtype == np.uint8:
55
+ if (bit_depth or 8) > 8:
56
+ raise ValueError(
57
+ f"High-bit-depth decode returned uint8 output for {bit_depth}-bit source"
58
+ )
55
59
  return np.uint8
56
60
  if arr_dtype == np.uint16:
57
61
  return np.uint16
@@ -83,6 +87,9 @@ def _normalize_rgb_numpy_output(
83
87
  arr = arr.astype(target, copy=False)
84
88
  copy_result = True
85
89
 
90
+ arr, scaled = _normalize_uint16_source_depth_scale(arr, bit_depth=bit_depth)
91
+ copy_result = copy_result or scaled
92
+
86
93
  if copy_result:
87
94
  return np.array(arr, copy=True, order="C")
88
95
  if not arr.flags.c_contiguous:
@@ -95,15 +102,17 @@ def _decoded_frame_to_rgb_numpy(frame, torch_mod, *, bit_depth: int | None = Non
95
102
  if hasattr(frame, "__dlpack__"):
96
103
  try:
97
104
  arr = np.from_dlpack(frame)
98
- return _normalize_rgb_numpy_output(arr, bit_depth=bit_depth, copy_result=True)
99
105
  except Exception:
100
106
  pass
107
+ else:
108
+ return _normalize_rgb_numpy_output(arr, bit_depth=bit_depth, copy_result=True)
101
109
  if torch_mod is not None:
102
110
  try:
103
111
  arr = torch_mod.from_dlpack(frame).detach().cpu().numpy()
104
- return _normalize_rgb_numpy_output(arr, bit_depth=bit_depth)
105
112
  except Exception:
106
113
  pass
114
+ else:
115
+ return _normalize_rgb_numpy_output(arr, bit_depth=bit_depth)
107
116
  cp = try_import_cupy()
108
117
  if cp is not None:
109
118
  arr = cp.asnumpy(cp.from_dlpack(frame))
@@ -161,11 +170,72 @@ def numpy_to_cupy_rgb(
161
170
  return out
162
171
 
163
172
 
173
+ def _source_depth_rgb_format(depth: int) -> str:
174
+ return "rgb24" if depth <= 8 else f"gbrp{depth}le"
175
+
176
+
177
+ def _pyav_source_color_range(frame):
178
+ color_range = getattr(frame, "color_range", None)
179
+ text = str(color_range)
180
+ try:
181
+ from av.video.reformatter import ColorRange
182
+ except Exception:
183
+ return color_range
184
+ if text in {"2", "JPEG", "jpeg", "pc"}:
185
+ return ColorRange.JPEG
186
+ if text in {"1", "MPEG", "mpeg", "tv"}:
187
+ return ColorRange.MPEG
188
+ return color_range
189
+
190
+
191
+ def _pyav_to_source_depth_rgb(frame, *, bit_depth: int):
192
+ fmt = _source_depth_rgb_format(bit_depth)
193
+ reformat = getattr(frame, "reformat", None)
194
+ if callable(reformat):
195
+ try:
196
+ return reformat(
197
+ format=fmt,
198
+ src_color_range=_pyav_source_color_range(frame),
199
+ ).to_ndarray()
200
+ except Exception:
201
+ pass
202
+ return frame.to_ndarray(format=fmt)
203
+
204
+
205
+ def _normalize_uint16_source_depth_scale(
206
+ arr: np.ndarray, *, bit_depth: int | None
207
+ ) -> tuple[np.ndarray, bool]:
208
+ if bit_depth is None or bit_depth <= 8 or bit_depth >= 16 or arr.dtype != np.uint16:
209
+ return arr, False
210
+
211
+ source_max = (1 << bit_depth) - 1
212
+ if arr.size == 0 or int(arr.max()) <= source_max:
213
+ return arr, False
214
+
215
+ shift = 16 - bit_depth
216
+ out = np.right_shift(arr, shift)
217
+ return np.clip(out, 0, source_max).astype(np.uint16, copy=False), True
218
+
219
+
220
+ def _rgb48_to_source_depth(arr: np.ndarray, *, bit_depth: int) -> np.ndarray:
221
+ arr = np.asarray(arr)
222
+ arr, _ = _ensure_hwc(arr)
223
+ if bit_depth >= 16:
224
+ return _normalize_rgb_numpy_output(arr, bit_depth=bit_depth, copy_result=True)
225
+
226
+ shift = 16 - bit_depth
227
+ source_max = (1 << bit_depth) - 1
228
+ out = np.right_shift(arr, shift)
229
+ return np.clip(out, 0, source_max).astype(np.uint16, copy=False)
230
+
231
+
164
232
  def _pyav_frame_to_rgb(frame, *, bit_depth: int | None = None) -> np.ndarray:
165
233
  """Convert a PyAV VideoFrame to an HWC RGB numpy array.
166
234
 
167
235
  Full-range yuv444p streams are converted manually to dodge swscale's
168
- limited-range default when VUI is unspecified.
236
+ limited-range default when VUI is unspecified. High-bit-depth outputs use
237
+ source-depth scale, matching pyavif: 10-bit values are 0..1023, 12-bit
238
+ values are 0..4095, stored in uint16 arrays.
169
239
  """
170
240
  pf = str(frame.format.name)
171
241
  depth = int(bit_depth) if bit_depth is not None else infer_video_bit_depth_from_frame(frame)
@@ -176,24 +246,34 @@ def _pyav_frame_to_rgb(frame, *, bit_depth: int | None = None) -> np.ndarray:
176
246
  if pf.startswith("yuv444p") and (
177
247
  is_full or cr is None or cr_text in ("0", "UNSPECIFIED", "unspecified")
178
248
  ):
179
- yuv = frame.to_ndarray(format=pf)
180
- y = yuv[0].astype(np.float32)
181
- max_value = float((1 << depth) - 1)
182
- center = (max_value + 1.0) / 2.0
183
- cb = yuv[1].astype(np.float32) - center
184
- crv = yuv[2].astype(np.float32) - center
185
-
186
- if depth > 8:
187
- max_out = 65535.0
188
- target_dtype = np.uint16
189
- else:
190
- max_out = 255.0
191
- target_dtype = np.uint8
192
- scale = max_out / max_value
193
-
194
- r = np.clip((y + 1.402 * crv) * scale + 0.5, 0, max_out).astype(target_dtype)
195
- g = np.clip((y - 0.344136 * cb - 0.714136 * crv) * scale + 0.5, 0, max_out).astype(target_dtype)
196
- b = np.clip((y + 1.772 * cb) * scale + 0.5, 0, max_out).astype(target_dtype)
197
- return np.stack([r, g, b], axis=-1)
198
-
199
- return frame.to_ndarray(format="rgb48le" if depth > 8 else "rgb24")
249
+ try:
250
+ yuv = frame.to_ndarray(format=pf)
251
+ except ValueError:
252
+ yuv = None
253
+ if yuv is not None:
254
+ y = yuv[0].astype(np.float32)
255
+ max_value = float((1 << depth) - 1)
256
+ center = (max_value + 1.0) / 2.0
257
+ cb = yuv[1].astype(np.float32) - center
258
+ crv = yuv[2].astype(np.float32) - center
259
+
260
+ max_out = max_value
261
+ target_dtype = np.uint16 if depth > 8 else np.uint8
262
+ scale = max_out / max_value
263
+
264
+ r = np.clip((y + 1.402 * crv) * scale + 0.5, 0, max_out).astype(target_dtype)
265
+ g = np.clip((y - 0.344136 * cb - 0.714136 * crv) * scale + 0.5, 0, max_out).astype(target_dtype)
266
+ b = np.clip((y + 1.772 * cb) * scale + 0.5, 0, max_out).astype(target_dtype)
267
+ return np.stack([r, g, b], axis=-1)
268
+
269
+ if depth <= 8:
270
+ return frame.to_ndarray(format="rgb24")
271
+
272
+ try:
273
+ return _normalize_rgb_numpy_output(
274
+ _pyav_to_source_depth_rgb(frame, bit_depth=depth),
275
+ bit_depth=depth,
276
+ copy_result=True,
277
+ )
278
+ except ValueError:
279
+ return _rgb48_to_source_depth(frame.to_ndarray(format="rgb48le"), bit_depth=depth)
@@ -0,0 +1,169 @@
1
+ """Native NVDEC YUV surface conversion helpers."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Any
6
+
7
+ import numpy as np
8
+
9
+ from ._imports import try_import_cupy
10
+
11
+
12
+ class _CaiView:
13
+ def __init__(self, view: Any, *, data_offset_bytes: int = 0):
14
+ self._view = view
15
+ self._data_offset_bytes = int(data_offset_bytes)
16
+
17
+ @property
18
+ def __cuda_array_interface__(self):
19
+ cai = dict(self._view.__cuda_array_interface__)
20
+ cai["shape"] = tuple(cai["shape"])
21
+
22
+ typestr = str(cai.get("typestr", ""))
23
+ itemsize = np.dtype(typestr).itemsize if typestr else 1
24
+ # PyNvVideoCodec reports these strides in samples for native planes.
25
+ cai["strides"] = tuple(int(stride) * itemsize for stride in cai["strides"])
26
+ if self._data_offset_bytes:
27
+ ptr, readonly = cai["data"]
28
+ cai["data"] = (int(ptr) + self._data_offset_bytes, readonly)
29
+ return cai
30
+
31
+
32
+ def _cupy_or_raise():
33
+ cp = try_import_cupy()
34
+ if cp is None:
35
+ raise RuntimeError(
36
+ "Native high-bit-depth NVDEC conversion requires CuPy. "
37
+ "Install mvdata[cuda] or use the PyAV fallback."
38
+ )
39
+ return cp
40
+
41
+
42
+ def _format_name(frame: Any) -> str:
43
+ return str(getattr(frame, "format", "")).split(".")[-1]
44
+
45
+
46
+ def _coefficients(color_space: Any) -> tuple[float, float]:
47
+ text = str(color_space).lower()
48
+ if "2020" in text or text == "9":
49
+ return 0.2627, 0.0593
50
+ if "709" in text or text == "1":
51
+ return 0.2126, 0.0722
52
+ return 0.299, 0.114
53
+
54
+
55
+ def _is_limited_range(color_range: Any) -> bool:
56
+ text = str(color_range).lower()
57
+ return "mpeg" in text or text == "1"
58
+
59
+
60
+ def _native_planes(frame: Any):
61
+ return [_CaiView(view) for view in frame.cuda()]
62
+
63
+
64
+ def _native_yuv444_planes(frame: Any):
65
+ views = frame.cuda()
66
+ if len(views) != 3:
67
+ raise ValueError(f"Expected 3 native YUV444 planes, got {len(views)}")
68
+ cai = dict(views[0].__cuda_array_interface__)
69
+ shape = tuple(cai["shape"])
70
+ strides = tuple(cai["strides"])
71
+ itemsize = np.dtype(cai.get("typestr", "|u2")).itemsize
72
+ plane_bytes = int(shape[0]) * int(strides[0]) * itemsize
73
+ return [_CaiView(views[0], data_offset_bytes=i * plane_bytes) for i in range(3)]
74
+
75
+
76
+ def _plane_to_source_float(cp, plane, *, bit_depth: int):
77
+ shift = 16 - int(bit_depth)
78
+ values = cp.asarray(plane)
79
+ if shift > 0:
80
+ values = cp.right_shift(values, shift)
81
+ return values.astype(cp.float32, copy=False)
82
+
83
+
84
+ def _single_channel(values):
85
+ return values[..., 0] if values.ndim == 3 and values.shape[-1] == 1 else values
86
+
87
+
88
+ def _yuv_to_rgb(cp, y, cb, cr, *, bit_depth: int, color_range: Any, color_space: Any):
89
+ source_max = float((1 << int(bit_depth)) - 1)
90
+ if _is_limited_range(color_range):
91
+ scale = 1 << (int(bit_depth) - 8)
92
+ y = (y - 16.0 * scale) * (source_max / (219.0 * scale))
93
+ center = 128.0 * scale
94
+ chroma_scale = ((source_max + 1.0) / 2.0) / (112.0 * scale)
95
+ cb = (cb - center) * chroma_scale
96
+ cr = (cr - center) * chroma_scale
97
+ else:
98
+ center = (source_max + 1.0) / 2.0
99
+ cb = cb - center
100
+ cr = cr - center
101
+
102
+ kr, kb = _coefficients(color_space)
103
+ kg = 1.0 - kr - kb
104
+ r_cr = 2.0 * (1.0 - kr)
105
+ b_cb = 2.0 * (1.0 - kb)
106
+ g_cb = 2.0 * kb * (1.0 - kb) / kg
107
+ g_cr = 2.0 * kr * (1.0 - kr) / kg
108
+
109
+ r = y + r_cr * cr
110
+ g = y - g_cb * cb - g_cr * cr
111
+ b = y + b_cb * cb
112
+ rgb = cp.stack([r, g, b], axis=-1)
113
+ return cp.clip(rgb + 0.5, 0, source_max).astype(cp.uint16)
114
+
115
+
116
+ def native_nvdec_to_rgb_cupy(
117
+ frame: Any,
118
+ *,
119
+ bit_depth: int,
120
+ color_range: Any = None,
121
+ color_space: Any = None,
122
+ cuda_stream: int | None = None,
123
+ ):
124
+ """Convert a native high-bit-depth NVDEC frame to HWC RGB on the GPU."""
125
+ if int(bit_depth) <= 8:
126
+ raise ValueError("Native high-bit-depth conversion requires bit_depth > 8")
127
+
128
+ cp = _cupy_or_raise()
129
+ stream_ctx = (
130
+ cp.cuda.ExternalStream(int(cuda_stream))
131
+ if cuda_stream is not None and int(cuda_stream) != 0
132
+ else None
133
+ )
134
+
135
+ def _convert():
136
+ fmt = _format_name(frame)
137
+ planes = _native_planes(frame)
138
+ if fmt in {"P016", "P010"}:
139
+ y_view, uv_view = planes
140
+ y = _single_channel(_plane_to_source_float(cp, y_view, bit_depth=bit_depth))
141
+ uv = _plane_to_source_float(cp, uv_view, bit_depth=bit_depth)
142
+ cb = cp.repeat(cp.repeat(uv[..., 0], 2, axis=0), 2, axis=1)[: y.shape[0], : y.shape[1]]
143
+ cr = cp.repeat(cp.repeat(uv[..., 1], 2, axis=0), 2, axis=1)[: y.shape[0], : y.shape[1]]
144
+ elif fmt == "YUV444_16Bit":
145
+ y_view, cb_view, cr_view = _native_yuv444_planes(frame)
146
+ y = _single_channel(_plane_to_source_float(cp, y_view, bit_depth=bit_depth))
147
+ cb = _single_channel(_plane_to_source_float(cp, cb_view, bit_depth=bit_depth))
148
+ cr = _single_channel(_plane_to_source_float(cp, cr_view, bit_depth=bit_depth))
149
+ else:
150
+ raise ValueError(f"Unsupported native NVDEC format for high-bit RGB: {fmt}")
151
+
152
+ return _yuv_to_rgb(
153
+ cp,
154
+ y,
155
+ cb,
156
+ cr,
157
+ bit_depth=bit_depth,
158
+ color_range=color_range,
159
+ color_space=color_space,
160
+ )
161
+
162
+ if stream_ctx is None:
163
+ return _convert()
164
+ with stream_ctx:
165
+ return _convert()
166
+
167
+
168
+ def native_nvdec_to_rgb_numpy(frame: Any, **kwargs) -> np.ndarray:
169
+ return native_nvdec_to_rgb_cupy(frame, **kwargs).get()
@@ -80,7 +80,30 @@ def probe_video_bit_depth_pyav(path: Path) -> int:
80
80
  return infer_video_bit_depth_from_stream(container.streams.video[0])
81
81
 
82
82
 
83
+ def probe_video_color_metadata_pyav(path: Path) -> dict[str, Any]:
84
+ av = try_import_av()
85
+ if av is None:
86
+ return {"color_range": None, "color_space": None}
87
+ try:
88
+ with av.open(str(path)) as container:
89
+ if not container.streams.video:
90
+ return {"color_range": None, "color_space": None}
91
+ stream = container.streams.video[0]
92
+ return {
93
+ "color_range": getattr(stream, "color_range", None),
94
+ "color_space": getattr(stream, "colorspace", None),
95
+ }
96
+ except Exception:
97
+ return {"color_range": None, "color_space": None}
98
+
99
+
83
100
  def probe_video_bit_depth(path: Path, nvc: Any | None = None) -> int:
101
+ pyav_error: Exception | None = None
102
+ try:
103
+ return probe_video_bit_depth_pyav(path)
104
+ except Exception as exc:
105
+ pyav_error = exc
106
+
84
107
  if nvc is None:
85
108
  try:
86
109
  nvc = try_import_pynvvideocodec()
@@ -91,7 +114,8 @@ def probe_video_bit_depth(path: Path, nvc: Any | None = None) -> int:
91
114
  return int(probe_video_stream_metadata(nvc, path)["bitdepth"])
92
115
  except Exception:
93
116
  pass
94
- return probe_video_bit_depth_pyav(path)
117
+
118
+ raise RuntimeError(f"Could not probe video bit depth for {path}") from pyav_error
95
119
 
96
120
 
97
121
  def probe_video_stream_metadata(nvc, path: Path | str) -> dict[str, Any]:
@@ -53,8 +53,14 @@ class _VideoNvdecFrameSource:
53
53
  self._bit_depth = bit_depth
54
54
  self._nvenc_codec = nvenc_codec
55
55
  self._torch_mod = nvenc_codec.try_import_torch()
56
+ self._native_high_bit = self._bit_depth > 8
57
+ self._color_metadata = (
58
+ nvenc_codec.probe_video_color_metadata_pyav(path) if self._native_high_bit else {}
59
+ )
60
+ if self._native_high_bit and nvenc_codec.try_import_cupy() is None:
61
+ raise RuntimeError("Native high-bit-depth NVDEC conversion requires CuPy")
56
62
  has_dlpack = self._torch_mod is not None or nvenc_codec.try_import_cupy() is not None
57
- use_dev_options = [True, False] if has_dlpack else [False]
63
+ use_dev_options = [True] if self._native_high_bit else ([True, False] if has_dlpack else [False])
58
64
 
59
65
  nvc = nvenc_codec.try_import_pynvvideocodec()
60
66
  issue = nvenc_codec.nvdec_decode_compatibility_issue_for_path(nvc, path, gpu_id)
@@ -63,7 +69,16 @@ class _VideoNvdecFrameSource:
63
69
 
64
70
  self._decoder = None
65
71
  for use_dev in use_dev_options:
66
- self._decoder = nvenc_codec._try_open_nvdec(nvc, str(path), gpu_id, use_dev)
72
+ if self._native_high_bit:
73
+ self._decoder = nvenc_codec._try_open_nvdec(
74
+ nvc,
75
+ str(path),
76
+ gpu_id,
77
+ use_dev,
78
+ output_color_type=nvc.OutputColorType.NATIVE,
79
+ )
80
+ else:
81
+ self._decoder = nvenc_codec._try_open_nvdec(nvc, str(path), gpu_id, use_dev)
67
82
  if self._decoder is not None:
68
83
  break
69
84
  if self._decoder is None:
@@ -103,6 +118,15 @@ class _VideoNvdecFrameSource:
103
118
  return self._decode_sequential(index)
104
119
 
105
120
  def _materialize(self, raw, *, device: Device, cuda_stream: int | None):
121
+ if self._native_high_bit:
122
+ kwargs = {
123
+ "bit_depth": self._bit_depth,
124
+ "cuda_stream": cuda_stream,
125
+ **self._color_metadata,
126
+ }
127
+ if device == "cuda":
128
+ return self._nvenc_codec.native_nvdec_to_rgb_cupy(raw, **kwargs)
129
+ return self._nvenc_codec.native_nvdec_to_rgb_numpy(raw, **kwargs)
106
130
  if device == "cuda":
107
131
  return self._nvenc_codec._decoded_frame_to_rgb_cupy(
108
132
  raw,
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: mvdata
3
- Version: 0.9.0
3
+ Version: 0.9.2
4
4
  Summary: Gracia Dataset Convention - Python library for working with multi-view video datasets
5
5
  Author: Gracia Team
6
6
  License: MIT
@@ -9,7 +9,6 @@ mvdata/gpu_support.py
9
9
  mvdata/image_metrics.py
10
10
  mvdata/legacy_writer.py
11
11
  mvdata/multivideo.py
12
- mvdata/multivideo_decode_benchmark.py
13
12
  mvdata/multivideo_slicer.py
14
13
  mvdata/multivideo_writer.py
15
14
  mvdata/nvdec_parallel.py
@@ -32,6 +31,7 @@ mvdata/codec/_imports.py
32
31
  mvdata/codec/decode.py
33
32
  mvdata/codec/encode.py
34
33
  mvdata/codec/frames.py
34
+ mvdata/codec/native_yuv.py
35
35
  mvdata/codec/probe.py
36
36
  mvdata/codec/select.py
37
37
  tests/test_dataset_base_defaults.py
@@ -39,7 +39,6 @@ tests/test_gpu_policy.py
39
39
  tests/test_gpu_support.py
40
40
  tests/test_image_metrics.py
41
41
  tests/test_multivideo_bit_depth.py
42
- tests/test_multivideo_decode_benchmark.py
43
42
  tests/test_multivideo_slicer.py
44
43
  tests/test_nvdec_parallel.py
45
44
  tests/test_per_camera.py
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "mvdata"
7
- version = "0.9.0"
7
+ version = "0.9.2"
8
8
  description = "Gracia Dataset Convention - Python library for working with multi-view video datasets"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.12,<3.13"