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.
- {mvdata-0.9.0 → mvdata-0.9.2}/PKG-INFO +1 -1
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/__init__.py +0 -12
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/__init__.py +6 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/decode.py +39 -5
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/frames.py +104 -24
- mvdata-0.9.2/mvdata/codec/native_yuv.py +169 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/probe.py +25 -1
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/video_stream_reader.py +26 -2
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata.egg-info/PKG-INFO +1 -1
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata.egg-info/SOURCES.txt +1 -2
- {mvdata-0.9.0 → mvdata-0.9.2}/pyproject.toml +1 -1
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_video_stream_reader.py +205 -3
- mvdata-0.9.0/mvdata/multivideo_decode_benchmark.py +0 -216
- mvdata-0.9.0/tests/test_multivideo_decode_benchmark.py +0 -127
- {mvdata-0.9.0 → mvdata-0.9.2}/README.md +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/cloud_storage.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/_imports.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/encode.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/codec/select.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/dataset_base.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/downloader.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/gpu_policy.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/gpu_support.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/image_metrics.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/legacy_writer.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/multivideo.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/multivideo_slicer.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/multivideo_writer.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/nvdec_parallel.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/nvenc_codec.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/per_frame.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/ranged.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/ranged_writer.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/stash_utils.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/utils.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/write_progress.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata/writer_base.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata.egg-info/dependency_links.txt +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata.egg-info/requires.txt +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/mvdata.egg-info/top_level.txt +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/setup.cfg +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_dataset_base_defaults.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_gpu_policy.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_gpu_support.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_image_metrics.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_multivideo_bit_depth.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_multivideo_slicer.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_nvdec_parallel.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_per_camera.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_ranged_mixed_streams.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_ranged_nvenc_roundtrip.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_ranged_resume.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_ranged_stream_discovery.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_roundtrip.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_s3_downloader.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_stash_bit_depth.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_stash_comprehensive.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_stash_policy.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_stash_regenerate.py +0 -0
- {mvdata-0.9.0 → mvdata-0.9.2}/tests/test_write_progress.py +0 -0
|
@@ -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(
|
|
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
|
|
79
|
-
|
|
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
|
-
|
|
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
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
max_out =
|
|
191
|
-
target_dtype = np.uint8
|
|
192
|
-
|
|
193
|
-
|
|
194
|
-
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
|
|
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
|
-
|
|
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.
|
|
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,
|
|
@@ -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.
|
|
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"
|