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.
Files changed (36) hide show
  1. {nvidia_codec-0.0.21/src/nvidia_codec.egg-info → nvidia_codec-0.0.22}/PKG-INFO +7 -2
  2. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/README.md +6 -1
  3. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/setup.cfg +1 -1
  4. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/__init__.py +1 -1
  5. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/cuda.py +1 -1
  6. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/decode.py +4 -2
  7. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/nvcuvid.py +0 -17
  8. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/libavformat.py +6 -4
  9. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/color.py +2 -3
  10. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/decode.py +36 -32
  11. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/player.py +2 -0
  12. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22/src/nvidia_codec.egg-info}/PKG-INFO +7 -2
  13. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec.egg-info/SOURCES.txt +2 -1
  14. nvidia_codec-0.0.22/tests/test_smoke.py +131 -0
  15. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/LICENSE +0 -0
  16. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/pyproject.toml +0 -0
  17. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/common.py +0 -0
  18. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/__init__.py +0 -0
  19. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/common.py +0 -0
  20. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/core/cuviddec.py +0 -0
  21. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/__init__.py +0 -0
  22. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/adapt.py +0 -0
  23. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/common.py +0 -0
  24. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/__init__.py +0 -0
  25. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/libavcodec.py +0 -0
  26. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/libavformat.py +0 -0
  27. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/libavutil.py +0 -0
  28. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/libavcodec.py +0 -0
  29. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/libavutil.py +0 -0
  30. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/test.py +0 -0
  31. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/__init__.py +0 -0
  32. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/compat.py +0 -0
  33. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/demux.py +0 -0
  34. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec.egg-info/dependency_links.txt +0 -0
  35. {nvidia_codec-0.0.21 → nvidia_codec-0.0.22}/src/nvidia_codec.egg-info/requires.txt +0 -0
  36. {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.21
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 8.x (for demuxing)
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 8.x (for demuxing)
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
@@ -1,6 +1,6 @@
1
1
  [metadata]
2
2
  name = nvidia-codec
3
- version = 0.0.21
3
+ version = 0.0.22
4
4
  author = Zesen Qian
5
5
  author_email = nvidia-codec@riaqn.org
6
6
  description = Pythonic interface to Nvidia Codec
@@ -19,7 +19,7 @@ Quick Start:
19
19
  Requirements:
20
20
  - NVIDIA GPU with NVDEC support
21
21
  - NVIDIA driver with libnvcuvid.so
22
- - FFmpeg 8.x (for demuxing)
22
+ - FFmpeg 9.x (for demuxing)
23
23
  - PyTorch with CUDA support
24
24
 
25
25
  Supported Codecs:
@@ -8,7 +8,7 @@ CUcontext = c_void_p
8
8
  CUresult = c_int
9
9
  cudaError = c_int
10
10
 
11
- CUdeviceptr = c_uint
11
+ CUdeviceptr = c_ulonglong # CUdeviceptr_v2: unsigned long long
12
12
 
13
13
  lib = cdll.LoadLibrary('libcuda.so')
14
14
  librt = cdll.LoadLibrary('libcudart.so')
@@ -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(CUVIDEOFORMAT))
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
- "stream": self.params.output_stream,
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
- lib = cdll.LoadLibrary("libavformat.so")
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 < 62:
27
+ if _major != 63:
26
28
  raise ImportError(
27
- f"nvidia-codec requires FFmpeg 8.x (libavformat 62+), "
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)) / (225 << bit8)
138
- v = (V.type(torch.float) - (128 << bit8)) / (225 << 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
- """Wraps a Surface with semaphore ownership. Releases on free."""
34
+ """Maps a Picture under the decoder's output-surface semaphore.
37
35
 
38
- def __init__(self, surface):
39
- surface.decoder._surface_sem.acquire()
40
- self._surface = surface
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.map(self._pic, stream))
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 * 2 (double the minimum).
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 carries its own slot+surface ownership (→
275
- # OwnedSurface); a bare Picture maps via the decoder (→ plain
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.map(pic, stream)
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.21
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 8.x (for demuxing)
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
@@ -31,4 +31,5 @@ src/nvidia_codec/utils/color.py
31
31
  src/nvidia_codec/utils/compat.py
32
32
  src/nvidia_codec/utils/decode.py
33
33
  src/nvidia_codec/utils/demux.py
34
- src/nvidia_codec/utils/player.py
34
+ src/nvidia_codec/utils/player.py
35
+ tests/test_smoke.py
@@ -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