nvidia-codec 0.0.21__tar.gz → 0.0.22__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.
- {nvidia_codec-0.0.21/src/nvidia_codec.egg-info → nvidia_codec-0.0.22}/PKG-INFO +7 -2
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/README.md +6 -1
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/setup.cfg +1 -1
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/__init__.py +1 -1
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/cuda.py +1 -1
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/decode.py +4 -2
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/nvcuvid.py +0 -17
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/libavformat.py +6 -4
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/color.py +2 -3
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/decode.py +36 -32
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/player.py +2 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22/src/nvidia_codec.egg-info}/PKG-INFO +7 -2
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec.egg-info/SOURCES.txt +2 -1
- nvidia_codec-0.0.22/tests/test_smoke.py +131 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/LICENSE +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/pyproject.toml +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/common.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/__init__.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/common.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/cuviddec.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/__init__.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/adapt.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/common.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/__init__.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/libavcodec.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/libavformat.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/libavutil.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/libavcodec.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/libavutil.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/test.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/__init__.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/compat.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/demux.py +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec.egg-info/dependency_links.txt +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec.egg-info/requires.txt +0 -0
- {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec.egg-info/top_level.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: nvidia-codec
|
|
3
|
-
Version: 0.0.
|
|
3
|
+
Version: 0.0.22
|
|
4
4
|
Summary: Pythonic interface to Nvidia Codec
|
|
5
5
|
Home-page: https://github.com/riaqn/python-nvidia-codec
|
|
6
6
|
Author: Zesen Qian
|
|
@@ -26,7 +26,7 @@ A Pythonic library for NVIDIA Video Codec (NVDEC).
|
|
|
26
26
|
- Python 3.10+
|
|
27
27
|
- NVIDIA GPU with NVDEC support
|
|
28
28
|
- NVIDIA driver with `libnvcuvid.so`
|
|
29
|
-
- FFmpeg
|
|
29
|
+
- FFmpeg 9.x (for demuxing)
|
|
30
30
|
|
|
31
31
|
## Installation
|
|
32
32
|
|
|
@@ -57,6 +57,11 @@ for timestamp, frame in player.frames(torch.uint8):
|
|
|
57
57
|
print(f'{timestamp}: {frame.shape}')
|
|
58
58
|
```
|
|
59
59
|
|
|
60
|
+
`frames()` without `start` continues from the demuxer's current position; it
|
|
61
|
+
only begins at the first frame on a freshly opened player. After a
|
|
62
|
+
`screenshot()` or a previous `frames()` pass, pass `start=timedelta(0)` to
|
|
63
|
+
decode from the beginning again.
|
|
64
|
+
|
|
60
65
|
### Screenshot (single frame extraction)
|
|
61
66
|
|
|
62
67
|
```python
|
|
@@ -9,7 +9,7 @@ A Pythonic library for NVIDIA Video Codec (NVDEC).
|
|
|
9
9
|
- Python 3.10+
|
|
10
10
|
- NVIDIA GPU with NVDEC support
|
|
11
11
|
- NVIDIA driver with `libnvcuvid.so`
|
|
12
|
-
- FFmpeg
|
|
12
|
+
- FFmpeg 9.x (for demuxing)
|
|
13
13
|
|
|
14
14
|
## Installation
|
|
15
15
|
|
|
@@ -40,6 +40,11 @@ for timestamp, frame in player.frames(torch.uint8):
|
|
|
40
40
|
print(f'{timestamp}: {frame.shape}')
|
|
41
41
|
```
|
|
42
42
|
|
|
43
|
+
`frames()` without `start` continues from the demuxer's current position; it
|
|
44
|
+
only begins at the first frame on a freshly opened player. After a
|
|
45
|
+
`screenshot()` or a previous `frames()` pass, pass `start=timedelta(0)` to
|
|
46
|
+
decode from the beginning again.
|
|
47
|
+
|
|
43
48
|
### Screenshot (single frame extraction)
|
|
44
49
|
|
|
45
50
|
```python
|
|
@@ -37,7 +37,7 @@ class Surface:
|
|
|
37
37
|
self.pts = pts
|
|
38
38
|
|
|
39
39
|
self.params = CUVIDPROCPARAMS()
|
|
40
|
-
memmove(byref(self.params), byref(params), sizeof(
|
|
40
|
+
memmove(byref(self.params), byref(params), sizeof(CUVIDPROCPARAMS))
|
|
41
41
|
|
|
42
42
|
self.params.output_stream = stream
|
|
43
43
|
|
|
@@ -117,7 +117,9 @@ class Surface:
|
|
|
117
117
|
"strides": strides,
|
|
118
118
|
"version": 3,
|
|
119
119
|
"data": (self.c_devptr.value, False),
|
|
120
|
-
|
|
120
|
+
# CAI v3 forbids stream 0; CUDA's stream 0 is the legacy
|
|
121
|
+
# default stream, which the protocol spells as 1.
|
|
122
|
+
"stream": self.params.output_stream or 1,
|
|
121
123
|
}
|
|
122
124
|
|
|
123
125
|
|
|
@@ -120,23 +120,6 @@ class CUVIDPARSERPARAMS(Structure):
|
|
|
120
120
|
]
|
|
121
121
|
CUstream = c_void_p # according to /opt/cuda/include/cuda.h
|
|
122
122
|
CUvideoparser = c_void_p # according to nvcuvid.h
|
|
123
|
-
class CUVIDPROCPARAMS(Structure):
|
|
124
|
-
_fields_ = [('progressive_frame', c_int),
|
|
125
|
-
('second_field', c_int),
|
|
126
|
-
('top_field_first', c_int),
|
|
127
|
-
('unpaired_field', c_int),
|
|
128
|
-
('reserved_flags', c_uint),
|
|
129
|
-
('reserved_zero', c_uint),
|
|
130
|
-
('raw_input_dptr', c_ulonglong),
|
|
131
|
-
('raw_input_pitch', c_uint),
|
|
132
|
-
('raw_input_format', c_uint),
|
|
133
|
-
('raw_output_dptr', c_ulonglong),
|
|
134
|
-
('raw_output_pitch', c_uint),
|
|
135
|
-
('Reserved1', c_uint),
|
|
136
|
-
('output_stream', CUstream),
|
|
137
|
-
('Reserved', c_uint * 46),
|
|
138
|
-
('histogram_dptr', POINTER(c_ulonglong)),
|
|
139
|
-
('Reserved2', c_void_p*1) ]
|
|
140
123
|
class CUvideopacketflags(Flag):
|
|
141
124
|
ENDOFSTREAM = 1
|
|
142
125
|
TIMESTAMP = auto()
|
|
@@ -17,14 +17,16 @@ import logging
|
|
|
17
17
|
|
|
18
18
|
log = logging.getLogger(__name__)
|
|
19
19
|
|
|
20
|
-
|
|
20
|
+
# Pin the ABI major to ffmpeg 9 (libavformat 63), matching the libavcodec.so.63
|
|
21
|
+
# pin in ffmpeg/libavcodec.py: struct layouts + enum values in include/ track
|
|
22
|
+
# ffmpeg 9, so a system upgrade fails loudly at load instead of mis-reading.
|
|
23
|
+
lib = cdll.LoadLibrary("libavformat.so.63")
|
|
21
24
|
|
|
22
|
-
# Check FFmpeg version compatibility (requires FFmpeg 8.x / libavformat 62+)
|
|
23
25
|
_version = lib.avformat_version()
|
|
24
26
|
_major = _version >> 16
|
|
25
|
-
if _major
|
|
27
|
+
if _major != 63:
|
|
26
28
|
raise ImportError(
|
|
27
|
-
f"nvidia-codec requires FFmpeg
|
|
29
|
+
f"nvidia-codec requires FFmpeg 9.x (libavformat 63), "
|
|
28
30
|
f"but found libavformat {_major}.{(_version >> 8) & 0xff}.{_version & 0xff}"
|
|
29
31
|
)
|
|
30
32
|
|
|
@@ -103,7 +103,6 @@ def convert(surface : Surface,
|
|
|
103
103
|
assert height % 2 == 0
|
|
104
104
|
assert width % 2 == 0
|
|
105
105
|
assert height // 2 * 3 == h
|
|
106
|
-
assert YUV.dtype == torch.uint8
|
|
107
106
|
|
|
108
107
|
Y = YUV[:height,:]
|
|
109
108
|
UV = YUV[height:,:]
|
|
@@ -134,8 +133,8 @@ def convert(surface : Surface,
|
|
|
134
133
|
#normalize
|
|
135
134
|
if source_range == AVColorRange.MPEG:
|
|
136
135
|
y = (Y.type(torch.float) - (16 << bit8)) / (220 << bit8)
|
|
137
|
-
u = (U.type(torch.float) - (128 << bit8)) / (
|
|
138
|
-
v = (V.type(torch.float) - (128 << bit8)) / (
|
|
136
|
+
u = (U.type(torch.float) - (128 << bit8)) / (224 << bit8)
|
|
137
|
+
v = (V.type(torch.float) - (128 << bit8)) / (224 << bit8)
|
|
139
138
|
elif source_range == AVColorRange.JPEG:
|
|
140
139
|
y = Y.type(torch.float) / (256 << bit8)
|
|
141
140
|
u = (U.type(torch.float) - (128 << bit8)) / (256 << bit8)
|
|
@@ -13,11 +13,9 @@ from ..ffmpeg.include.libavutil import (
|
|
|
13
13
|
AV_NOPTS_VALUE,
|
|
14
14
|
AVColorRange,
|
|
15
15
|
AVColorSpace,
|
|
16
|
-
AVDISCARD_NONE,
|
|
17
|
-
AVDISCARD_NONKEY,
|
|
18
16
|
)
|
|
19
17
|
from .compat import av2cuda, extract_stream_ptr
|
|
20
|
-
from ..core.decode import BaseDecoder, Picture, Surface
|
|
18
|
+
from ..core.decode import BaseDecoder, Picture, Surface, decide_surface_format
|
|
21
19
|
from ..ffmpeg.common import AVException
|
|
22
20
|
from .color import convert, to_yuv420p
|
|
23
21
|
from .demux import VideoTrack
|
|
@@ -33,11 +31,21 @@ log = logging.getLogger(__name__)
|
|
|
33
31
|
|
|
34
32
|
|
|
35
33
|
class OwnedSurface:
|
|
36
|
-
"""
|
|
34
|
+
"""Maps a Picture under the decoder's output-surface semaphore.
|
|
37
35
|
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
36
|
+
The semaphore (sized to ulNumOutputSurfaces) must be held BEFORE
|
|
37
|
+
cuvidMapVideoFrame64 runs — two concurrent maps with one output
|
|
38
|
+
surface would overwrite each other — so the map happens here, not
|
|
39
|
+
in the caller. free() unmaps and releases.
|
|
40
|
+
"""
|
|
41
|
+
|
|
42
|
+
def __init__(self, decoder, pic, stream):
|
|
43
|
+
decoder._surface_sem.acquire()
|
|
44
|
+
try:
|
|
45
|
+
self._surface = decoder.map(pic, stream)
|
|
46
|
+
except BaseException:
|
|
47
|
+
decoder._surface_sem.release()
|
|
48
|
+
raise
|
|
41
49
|
|
|
42
50
|
def __getattr__(self, name):
|
|
43
51
|
return getattr(self._surface, name)
|
|
@@ -83,7 +91,7 @@ class OwnedPicture:
|
|
|
83
91
|
|
|
84
92
|
def map(self, stream=0):
|
|
85
93
|
assert self._pic is not None, "Picture already freed"
|
|
86
|
-
return OwnedSurface(self._decoder
|
|
94
|
+
return OwnedSurface(self._decoder, self._pic, stream)
|
|
87
95
|
|
|
88
96
|
|
|
89
97
|
class Decoder(BaseDecoder):
|
|
@@ -100,7 +108,10 @@ class Decoder(BaseDecoder):
|
|
|
100
108
|
num_pictures: Callable(min_pics) -> int, returns the total number of
|
|
101
109
|
picture slots (must be >= min_pics and <= 32). Higher values allow
|
|
102
110
|
better pipelining between decode and consumer threads.
|
|
103
|
-
Default: lambda min_pics: min_pics
|
|
111
|
+
Default: lambda min_pics: min_pics + 4.
|
|
112
|
+
allow_high: Keep >8-bit sources at 16 bits (P016 / YUV444_16Bit
|
|
113
|
+
surfaces) instead of letting NVDEC round them to 8-bit NV12 /
|
|
114
|
+
YUV444. convert() handles both; to_yuv420p() is 8-bit only.
|
|
104
115
|
"""
|
|
105
116
|
|
|
106
117
|
def __init__(
|
|
@@ -111,6 +122,7 @@ class Decoder(BaseDecoder):
|
|
|
111
122
|
cropping=None,
|
|
112
123
|
target_rect=None,
|
|
113
124
|
device=None,
|
|
125
|
+
allow_high=False,
|
|
114
126
|
):
|
|
115
127
|
self.track = track
|
|
116
128
|
self._slot_locks = None # initialized in decide(), after we know num_pictures
|
|
@@ -147,6 +159,10 @@ class Decoder(BaseDecoder):
|
|
|
147
159
|
"num_pictures": num_pics,
|
|
148
160
|
"num_surfaces": 1,
|
|
149
161
|
}
|
|
162
|
+
if allow_high:
|
|
163
|
+
d["surface_format"] = decide_surface_format(
|
|
164
|
+
p["chroma_format"], p["bit_depth"],
|
|
165
|
+
p["supported_surface_formats"], allow_high=True)
|
|
150
166
|
if target_size is not None:
|
|
151
167
|
d["target_size"] = target_size
|
|
152
168
|
if cropping is not None:
|
|
@@ -242,25 +258,6 @@ class Decoder(BaseDecoder):
|
|
|
242
258
|
self.with_callbacks(lambda: self.send(None), post_decode=lambda pic: None)
|
|
243
259
|
self._recent.clear()
|
|
244
260
|
|
|
245
|
-
def _decode_packets(self, keyframes_only=False):
|
|
246
|
-
"""Read packets from current position, filter, decode.
|
|
247
|
-
|
|
248
|
-
Calls self.post_decode(pic) for each frame.
|
|
249
|
-
"""
|
|
250
|
-
self.track.stream.discard = AVDISCARD_NONKEY if keyframes_only else AVDISCARD_NONE
|
|
251
|
-
packets = self.track.fc.read_packets(self.track.stream)
|
|
252
|
-
it = self.bsf.filter(packets, flush=True, reuse=True)
|
|
253
|
-
self.with_callbacks(lambda: self.send(None), post_decode=lambda pic: None)
|
|
254
|
-
|
|
255
|
-
for pkt in it:
|
|
256
|
-
pts = pkt.av.pts if pkt.av.pts != AV_NOPTS_VALUE else pkt.av.dts
|
|
257
|
-
arr = np.ctypeslib.as_array(pkt.av.data, (pkt.av.size,))
|
|
258
|
-
self.send((pts, arr))
|
|
259
|
-
if keyframes_only:
|
|
260
|
-
self.bsf.flush()
|
|
261
|
-
self.send(None)
|
|
262
|
-
self.send(None)
|
|
263
|
-
|
|
264
261
|
def _map_and_convert(self, pic, dtype, pixel_format='rgb'):
|
|
265
262
|
"""Map a Picture (or OwnedPicture) to a tensor. Frees pic and surface.
|
|
266
263
|
Returns (timedelta, frame) on success, (timedelta, exception) on error.
|
|
@@ -268,16 +265,16 @@ class Decoder(BaseDecoder):
|
|
|
268
265
|
pixel_format='rgb' (default) → (3,H,W) RGB via convert() (color matrix).
|
|
269
266
|
pixel_format='yuv420p' → 1D planar I420 uint8 via to_yuv420p() (no color
|
|
270
267
|
math, half the bytes) — for piping straight to an encoder."""
|
|
268
|
+
t = None
|
|
271
269
|
try:
|
|
272
270
|
t = self.track.pts2time(pic.pts)
|
|
273
271
|
stream = extract_stream_ptr(torch.cuda.current_stream(self.device))
|
|
274
|
-
# OwnedPicture
|
|
275
|
-
#
|
|
276
|
-
# Surface). Both expose .free().
|
|
272
|
+
# Both OwnedPicture (owns its slot) and bare Picture map under
|
|
273
|
+
# the output-surface semaphore; .free() on either is safe.
|
|
277
274
|
if isinstance(pic, OwnedPicture):
|
|
278
275
|
surface = pic.map(stream)
|
|
279
276
|
else:
|
|
280
|
-
surface = self
|
|
277
|
+
surface = OwnedSurface(self, pic, stream)
|
|
281
278
|
except Exception as e:
|
|
282
279
|
return (t, e)
|
|
283
280
|
finally:
|
|
@@ -683,6 +680,13 @@ class Decoder(BaseDecoder):
|
|
|
683
680
|
"""Yield (int_pts, frame_or_exception) for every decoded frame in
|
|
684
681
|
display order. Decodes ALL frames (B-frames included).
|
|
685
682
|
|
|
683
|
+
POSITION: without `start` / `start_pts` this does NOT rewind. It
|
|
684
|
+
decodes from wherever the demuxer currently is — the beginning on
|
|
685
|
+
a fresh Decoder, but after a `screenshot()`, `screenshots()` or an
|
|
686
|
+
exhausted `frames()` it resumes from there (and yields nothing if
|
|
687
|
+
that was EOF). Pass `start=timedelta(0)` to decode from the start
|
|
688
|
+
of the file regardless of what ran before.
|
|
689
|
+
|
|
686
690
|
The yielded `int_pts` is the raw track pts (same units as
|
|
687
691
|
`pic.pts`, including the demuxer's `start_time` offset). Pass
|
|
688
692
|
through `track.pts2time(pts)` if you need a `timedelta`. Yielding
|
|
@@ -20,6 +20,7 @@ class Player(Decoder):
|
|
|
20
20
|
target_rect=None,
|
|
21
21
|
device=None,
|
|
22
22
|
track_idx=None,
|
|
23
|
+
**decoder_kwargs,
|
|
23
24
|
):
|
|
24
25
|
tracks = [t for t in parse(url) if isinstance(t, VideoTrack)]
|
|
25
26
|
if not tracks:
|
|
@@ -38,5 +39,6 @@ class Player(Decoder):
|
|
|
38
39
|
cropping=cropping,
|
|
39
40
|
target_rect=target_rect,
|
|
40
41
|
device=device,
|
|
42
|
+
**decoder_kwargs,
|
|
41
43
|
)
|
|
42
44
|
self.url = url
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: nvidia-codec
|
|
3
|
-
Version: 0.0.
|
|
3
|
+
Version: 0.0.22
|
|
4
4
|
Summary: Pythonic interface to Nvidia Codec
|
|
5
5
|
Home-page: https://github.com/riaqn/python-nvidia-codec
|
|
6
6
|
Author: Zesen Qian
|
|
@@ -26,7 +26,7 @@ A Pythonic library for NVIDIA Video Codec (NVDEC).
|
|
|
26
26
|
- Python 3.10+
|
|
27
27
|
- NVIDIA GPU with NVDEC support
|
|
28
28
|
- NVIDIA driver with `libnvcuvid.so`
|
|
29
|
-
- FFmpeg
|
|
29
|
+
- FFmpeg 9.x (for demuxing)
|
|
30
30
|
|
|
31
31
|
## Installation
|
|
32
32
|
|
|
@@ -57,6 +57,11 @@ for timestamp, frame in player.frames(torch.uint8):
|
|
|
57
57
|
print(f'{timestamp}: {frame.shape}')
|
|
58
58
|
```
|
|
59
59
|
|
|
60
|
+
`frames()` without `start` continues from the demuxer's current position; it
|
|
61
|
+
only begins at the first frame on a freshly opened player. After a
|
|
62
|
+
`screenshot()` or a previous `frames()` pass, pass `start=timedelta(0)` to
|
|
63
|
+
decode from the beginning again.
|
|
64
|
+
|
|
60
65
|
### Screenshot (single frame extraction)
|
|
61
66
|
|
|
62
67
|
```python
|
|
@@ -0,0 +1,131 @@
|
|
|
1
|
+
"""Smoke test: synthesize a short H.264 clip with ffmpeg, decode it with NVDEC.
|
|
2
|
+
|
|
3
|
+
Requires an NVIDIA GPU, libnvcuvid, FFmpeg 9 libraries and the `ffmpeg` CLI.
|
|
4
|
+
Run with: python -m pytest tests/ (or simply: python tests/test_smoke.py)
|
|
5
|
+
"""
|
|
6
|
+
from datetime import timedelta
|
|
7
|
+
import os
|
|
8
|
+
import shutil
|
|
9
|
+
import subprocess
|
|
10
|
+
import tempfile
|
|
11
|
+
|
|
12
|
+
import pytest
|
|
13
|
+
import torch
|
|
14
|
+
|
|
15
|
+
from nvidia_codec.core.cuviddec import cudaVideoSurfaceFormat
|
|
16
|
+
from nvidia_codec.utils import Player, parse, VideoTrack
|
|
17
|
+
|
|
18
|
+
W, H, FPS, SECONDS = 320, 240, 25, 2
|
|
19
|
+
|
|
20
|
+
pytestmark = pytest.mark.skipif(
|
|
21
|
+
not (shutil.which("ffmpeg") and torch.cuda.is_available()),
|
|
22
|
+
reason="needs ffmpeg CLI and a CUDA GPU",
|
|
23
|
+
)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
@pytest.fixture(scope="module")
|
|
27
|
+
def clip():
|
|
28
|
+
d = tempfile.mkdtemp(prefix="nvc-smoke-")
|
|
29
|
+
path = os.path.join(d, "clip.mp4")
|
|
30
|
+
subprocess.run(
|
|
31
|
+
["ffmpeg", "-v", "error", "-y", "-f", "lavfi",
|
|
32
|
+
f"-i", f"testsrc=size={W}x{H}:rate={FPS}:duration={SECONDS}",
|
|
33
|
+
"-pix_fmt", "yuv420p", "-c:v", "libx264", "-g", "12", path],
|
|
34
|
+
check=True, timeout=60,
|
|
35
|
+
)
|
|
36
|
+
yield path
|
|
37
|
+
shutil.rmtree(d, ignore_errors=True)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def test_parse(clip):
|
|
41
|
+
tracks = [t for t in parse(clip) if isinstance(t, VideoTrack)]
|
|
42
|
+
assert len(tracks) == 1
|
|
43
|
+
t = tracks[0]
|
|
44
|
+
assert (t.width, t.height) == (W, H)
|
|
45
|
+
assert t.mime_codec.startswith("avc1.")
|
|
46
|
+
assert abs(t.duration.total_seconds() - SECONDS) < 0.2
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def test_frames_rgb(clip):
|
|
50
|
+
with Player(clip) as p:
|
|
51
|
+
n = 0
|
|
52
|
+
last = -1
|
|
53
|
+
for pts, frame in p.frames(torch.uint8):
|
|
54
|
+
assert not isinstance(frame, Exception), frame
|
|
55
|
+
assert frame.shape == (3, H, W) and frame.dtype == torch.uint8
|
|
56
|
+
assert frame.is_cuda
|
|
57
|
+
assert pts > last
|
|
58
|
+
last = pts
|
|
59
|
+
n += 1
|
|
60
|
+
assert n == FPS * SECONDS
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
def test_frames_yuv420p_window(clip):
|
|
64
|
+
with Player(clip) as p:
|
|
65
|
+
out = list(p.frames(torch.uint8, start=timedelta(seconds=0.5),
|
|
66
|
+
end=timedelta(seconds=1.0), pixel_format="yuv420p"))
|
|
67
|
+
assert 0 < len(out) <= FPS // 2 + 2
|
|
68
|
+
for _, frame in out:
|
|
69
|
+
assert frame.shape == (H * W * 3 // 2,)
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def test_screenshot(clip):
|
|
73
|
+
with Player(clip, target_size=lambda h, w: (h // 2, w // 2)) as p:
|
|
74
|
+
t, frame = p.screenshot(timedelta(seconds=1), torch.float32, accurate=True)
|
|
75
|
+
assert frame.shape == (3, H // 2, W // 2)
|
|
76
|
+
assert abs(t.total_seconds() - 1.0) < 1 / FPS + 1e-3
|
|
77
|
+
assert 0.0 <= frame.min() and frame.max() <= 1.0
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
def test_screenshots(clip):
|
|
81
|
+
with Player(clip) as p:
|
|
82
|
+
shots = [s for s in p.screenshots(torch.uint8, timedelta(seconds=0.5))]
|
|
83
|
+
assert len(shots) >= 3
|
|
84
|
+
for _, frame, _ in shots:
|
|
85
|
+
assert not isinstance(frame, Exception), frame
|
|
86
|
+
assert frame.shape == (3, H, W)
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
def _has_encoder(name):
|
|
90
|
+
out = subprocess.run(["ffmpeg", "-hide_banner", "-encoders"],
|
|
91
|
+
capture_output=True, text=True, timeout=30).stdout
|
|
92
|
+
return f" {name} " in out
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
@pytest.fixture(scope="module")
|
|
96
|
+
def clip10():
|
|
97
|
+
if not _has_encoder("libx265"):
|
|
98
|
+
pytest.skip("ffmpeg lacks libx265")
|
|
99
|
+
d = tempfile.mkdtemp(prefix="nvc-smoke10-")
|
|
100
|
+
path = os.path.join(d, "clip10.mp4")
|
|
101
|
+
subprocess.run(
|
|
102
|
+
["ffmpeg", "-v", "error", "-y", "-f", "lavfi",
|
|
103
|
+
"-i", f"testsrc=size={W}x{H}:rate={FPS}:duration=1",
|
|
104
|
+
"-pix_fmt", "yuv420p10le", "-c:v", "libx265",
|
|
105
|
+
"-x265-params", "log-level=none", "-g", "12", path],
|
|
106
|
+
check=True, timeout=120,
|
|
107
|
+
)
|
|
108
|
+
yield path
|
|
109
|
+
shutil.rmtree(d, ignore_errors=True)
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
def test_10bit_allow_high_matches_8bit(clip10):
|
|
113
|
+
"""P016 (allow_high) and NV12 paths must agree on the same 10-bit source."""
|
|
114
|
+
with Player(clip10) as p:
|
|
115
|
+
_, lo = p.screenshot(timedelta(seconds=0.5), torch.float32, accurate=True)
|
|
116
|
+
assert p.surface_format == cudaVideoSurfaceFormat.NV12
|
|
117
|
+
with Player(clip10, allow_high=True) as p:
|
|
118
|
+
_, hi = p.screenshot(timedelta(seconds=0.5), torch.float32, accurate=True)
|
|
119
|
+
assert p.surface_format == cudaVideoSurfaceFormat.P016
|
|
120
|
+
assert hi.shape == (3, H, W)
|
|
121
|
+
assert 0.0 <= hi.min() and hi.max() <= 1.0
|
|
122
|
+
# yuv420p export is 8-bit only: the error is reported in-band
|
|
123
|
+
_, r = next(iter(p.frames(torch.uint8, start=timedelta(0),
|
|
124
|
+
pixel_format="yuv420p")))
|
|
125
|
+
assert isinstance(r, ValueError)
|
|
126
|
+
# 10-bit rounds to 8-bit at most 1/255 off, plus matrix rounding
|
|
127
|
+
assert (hi - lo).abs().max() < 0.02
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
if __name__ == "__main__":
|
|
131
|
+
raise SystemExit(pytest.main([__file__, "-v"]))
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|