nvidia-codec 0.0.21__tar.gz → 0.0.23__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.23}/PKG-INFO +25 -3
  2. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/README.md +24 -2
  3. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/setup.cfg +1 -1
  4. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/__init__.py +1 -1
  5. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/core/cuda.py +1 -1
  6. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/core/decode.py +4 -2
  7. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/core/nvcuvid.py +0 -17
  8. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/ffmpeg/include/libavformat.py +1 -0
  9. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/ffmpeg/include/libavutil.py +7 -5
  10. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/ffmpeg/libavformat.py +6 -14
  11. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/utils/color.py +2 -3
  12. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/utils/decode.py +76 -32
  13. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/utils/demux.py +23 -8
  14. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/utils/player.py +6 -1
  15. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23/src/nvidia_codec.egg-info}/PKG-INFO +25 -3
  16. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec.egg-info/SOURCES.txt +2 -1
  17. nvidia_codec-0.0.23/tests/test_smoke.py +188 -0
  18. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/LICENSE +0 -0
  19. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/pyproject.toml +0 -0
  20. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/common.py +0 -0
  21. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/core/__init__.py +0 -0
  22. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/core/common.py +0 -0
  23. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/core/cuviddec.py +0 -0
  24. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/ffmpeg/__init__.py +0 -0
  25. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/ffmpeg/adapt.py +0 -0
  26. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/ffmpeg/common.py +0 -0
  27. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/ffmpeg/include/__init__.py +0 -0
  28. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/ffmpeg/include/libavcodec.py +0 -0
  29. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/ffmpeg/libavcodec.py +0 -0
  30. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/ffmpeg/libavutil.py +0 -0
  31. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/ffmpeg/test.py +0 -0
  32. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/utils/__init__.py +0 -0
  33. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec/utils/compat.py +0 -0
  34. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec.egg-info/dependency_links.txt +0 -0
  35. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/src/nvidia_codec.egg-info/requires.txt +0 -0
  36. {nvidia_codec-0.0.21 → nvidia_codec-0.0.23}/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.23
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
@@ -119,6 +124,23 @@ Codec support depends on your GPU:
119
124
  - VPF allocates the bare minimum of resources needed for most decoding tasks. PNC allows the user to specify the amount of resources to be allocated for advanced applications. Users own the resources and decide when and whether to deal with them.
120
125
  - Managing resources is not painful: similar to `pycuda`, we shift the burden of managing host/device resources to the Python garbage collector. Resources (such as `Picture` and `Surface`) are automatically freed when the user drops the reference.
121
126
 
127
+ ## Limitations
128
+
129
+ - **One consumer per `Decoder`.** A `Decoder`/`Player` instance serves one
130
+ `frames()`, `screenshots()` or `screenshot()` call at a time. Use separate
131
+ instances (from any threads) for concurrent decodes of the same or different
132
+ files.
133
+ - **No mid-stream resolution or codec changes.** The NVDEC decoder is created
134
+ from the first sequence header; a later sequence with a different coded
135
+ size, chroma format or bit depth raises `decoder already initialized with
136
+ different format`. Split such files per segment.
137
+ - **Timestamp-less elementary streams** (e.g. raw Annex-B `.h264`/`.hevc`)
138
+ carry no pts: `frames()` yields `AV_NOPTS_VALUE` for every frame,
139
+ `pts2time` returns `None`, and `screenshot()` cannot locate a time. Mux
140
+ into a container first.
141
+ - **Embedded cover art** (`attached_pic` streams) is skipped by `Player`
142
+ and reported by `VideoTrack.attached_pic` from `parse()`.
143
+
122
144
  ## Roadmap
123
145
 
124
146
  - [x] Decoding
@@ -128,7 +150,7 @@ Codec support depends on your GPU:
128
150
  - Color Ranges: MPEG (limited), JPEG (full)
129
151
  - Color Space: BT.601, BT.709, BT.2020
130
152
  - [x] Built-in cropping and scaling via NV decoder
131
- - [x] Thread-safe and thread-friendly
153
+ - [x] Thread-friendly (one consumer per decoder; decoders are independent)
132
154
  - [ ] Encoder
133
155
  - [x] PyTorch integration
134
156
 
@@ -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
@@ -102,6 +107,23 @@ Codec support depends on your GPU:
102
107
  - VPF allocates the bare minimum of resources needed for most decoding tasks. PNC allows the user to specify the amount of resources to be allocated for advanced applications. Users own the resources and decide when and whether to deal with them.
103
108
  - Managing resources is not painful: similar to `pycuda`, we shift the burden of managing host/device resources to the Python garbage collector. Resources (such as `Picture` and `Surface`) are automatically freed when the user drops the reference.
104
109
 
110
+ ## Limitations
111
+
112
+ - **One consumer per `Decoder`.** A `Decoder`/`Player` instance serves one
113
+ `frames()`, `screenshots()` or `screenshot()` call at a time. Use separate
114
+ instances (from any threads) for concurrent decodes of the same or different
115
+ files.
116
+ - **No mid-stream resolution or codec changes.** The NVDEC decoder is created
117
+ from the first sequence header; a later sequence with a different coded
118
+ size, chroma format or bit depth raises `decoder already initialized with
119
+ different format`. Split such files per segment.
120
+ - **Timestamp-less elementary streams** (e.g. raw Annex-B `.h264`/`.hevc`)
121
+ carry no pts: `frames()` yields `AV_NOPTS_VALUE` for every frame,
122
+ `pts2time` returns `None`, and `screenshot()` cannot locate a time. Mux
123
+ into a container first.
124
+ - **Embedded cover art** (`attached_pic` streams) is skipped by `Player`
125
+ and reported by `VideoTrack.attached_pic` from `parse()`.
126
+
105
127
  ## Roadmap
106
128
 
107
129
  - [x] Decoding
@@ -111,7 +133,7 @@ Codec support depends on your GPU:
111
133
  - Color Ranges: MPEG (limited), JPEG (full)
112
134
  - Color Space: BT.601, BT.709, BT.2020
113
135
  - [x] Built-in cropping and scaling via NV decoder
114
- - [x] Thread-safe and thread-friendly
136
+ - [x] Thread-friendly (one consumer per decoder; decoders are independent)
115
137
  - [ ] Encoder
116
138
  - [x] PyTorch integration
117
139
 
@@ -1,6 +1,6 @@
1
1
  [metadata]
2
2
  name = nvidia-codec
3
- version = 0.0.21
3
+ version = 0.0.23
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()
@@ -19,6 +19,7 @@ class AVIndexEntry(Structure):
19
19
  return (self.flags_size >> 2) & 0x3FFFFFFF
20
20
 
21
21
  AVINDEX_KEYFRAME = 0x0001
22
+ AV_DISPOSITION_ATTACHED_PIC = 0x0400
22
23
 
23
24
  class AVStream(Structure):
24
25
  _fields_ = [
@@ -63,10 +63,12 @@ class AVColorRange(TypedCEnumeration(c_int)):
63
63
  # NB = 3
64
64
 
65
65
  class AVPixelFormat(TypedCEnumeration(c_int)):
66
- NONE = -1,
67
- RGB24 = 0,
68
- YUV420P = 1,
69
- YUV444P = 5,
70
- YUV420P16LE = 45,
66
+ # values per libavutil/pixfmt.h (ffmpeg 9); no trailing commas — they
67
+ # would make the members 1-tuples, not ints
68
+ NONE = -1
69
+ YUV420P = 0
70
+ RGB24 = 2
71
+ YUV444P = 5
72
+ YUV420P16LE = 45
71
73
  YUV444P16LE = 49
72
74
  # other omitted
@@ -3,7 +3,6 @@
3
3
  # not pyav
4
4
 
5
5
  from ctypes import *
6
- from fractions import Fraction
7
6
 
8
7
  from .include.libavutil import *
9
8
  from .include.libavcodec import *
@@ -17,14 +16,16 @@ import logging
17
16
 
18
17
  log = logging.getLogger(__name__)
19
18
 
20
- lib = cdll.LoadLibrary("libavformat.so")
19
+ # Pin the ABI major to ffmpeg 9 (libavformat 63), matching the libavcodec.so.63
20
+ # pin in ffmpeg/libavcodec.py: struct layouts + enum values in include/ track
21
+ # ffmpeg 9, so a system upgrade fails loudly at load instead of mis-reading.
22
+ lib = cdll.LoadLibrary("libavformat.so.63")
21
23
 
22
- # Check FFmpeg version compatibility (requires FFmpeg 8.x / libavformat 62+)
23
24
  _version = lib.avformat_version()
24
25
  _major = _version >> 16
25
- if _major < 62:
26
+ if _major != 63:
26
27
  raise ImportError(
27
- f"nvidia-codec requires FFmpeg 8.x (libavformat 62+), "
28
+ f"nvidia-codec requires FFmpeg 9.x (libavformat 63), "
28
29
  f"but found libavformat {_major}.{(_version >> 8) & 0xff}.{_version & 0xff}"
29
30
  )
30
31
 
@@ -102,15 +103,6 @@ class FormatContext:
102
103
  return result.contents
103
104
  return None
104
105
 
105
- def infer_start_time(self):
106
- # log.warning('infering start time from first packet')
107
- # in this case, we get the first packet
108
- pkt = Packet()
109
- self.read_packet(pkt)
110
- start_time = pkt.av.pts
111
- time_base = pkt.av.time_base
112
- return start_time * Fraction(time_base.num, time_base.den)
113
-
114
106
  @property
115
107
  def streams(self):
116
108
  for i in range(self.av.nb_streams):
@@ -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)
@@ -2,6 +2,7 @@
2
2
 
3
3
  import collections
4
4
  from datetime import timedelta
5
+ from fractions import Fraction
5
6
  import numpy as np
6
7
  from queue import Queue, ShutDown
7
8
  import threading
@@ -9,15 +10,14 @@ import weakref
9
10
 
10
11
  from ..ffmpeg.libavcodec import BSFContext
11
12
  from ..ffmpeg.libavformat import FormatContext, AVCodecID
13
+ from ..ffmpeg.include.libavcodec import AV_PKT_FLAG_KEY
12
14
  from ..ffmpeg.include.libavutil import (
13
15
  AV_NOPTS_VALUE,
14
16
  AVColorRange,
15
17
  AVColorSpace,
16
- AVDISCARD_NONE,
17
- AVDISCARD_NONKEY,
18
18
  )
19
19
  from .compat import av2cuda, extract_stream_ptr
20
- from ..core.decode import BaseDecoder, Picture, Surface
20
+ from ..core.decode import BaseDecoder, Picture, Surface, decide_surface_format
21
21
  from ..ffmpeg.common import AVException
22
22
  from .color import convert, to_yuv420p
23
23
  from .demux import VideoTrack
@@ -33,11 +33,21 @@ log = logging.getLogger(__name__)
33
33
 
34
34
 
35
35
  class OwnedSurface:
36
- """Wraps a Surface with semaphore ownership. Releases on free."""
36
+ """Maps a Picture under the decoder's output-surface semaphore.
37
37
 
38
- def __init__(self, surface):
39
- surface.decoder._surface_sem.acquire()
40
- self._surface = surface
38
+ The semaphore (sized to ulNumOutputSurfaces) must be held BEFORE
39
+ cuvidMapVideoFrame64 runs — two concurrent maps with one output
40
+ surface would overwrite each other — so the map happens here, not
41
+ in the caller. free() unmaps and releases.
42
+ """
43
+
44
+ def __init__(self, decoder, pic, stream):
45
+ decoder._surface_sem.acquire()
46
+ try:
47
+ self._surface = decoder.map(pic, stream)
48
+ except BaseException:
49
+ decoder._surface_sem.release()
50
+ raise
41
51
 
42
52
  def __getattr__(self, name):
43
53
  return getattr(self._surface, name)
@@ -83,7 +93,7 @@ class OwnedPicture:
83
93
 
84
94
  def map(self, stream=0):
85
95
  assert self._pic is not None, "Picture already freed"
86
- return OwnedSurface(self._decoder.map(self._pic, stream))
96
+ return OwnedSurface(self._decoder, self._pic, stream)
87
97
 
88
98
 
89
99
  class Decoder(BaseDecoder):
@@ -100,7 +110,10 @@ class Decoder(BaseDecoder):
100
110
  num_pictures: Callable(min_pics) -> int, returns the total number of
101
111
  picture slots (must be >= min_pics and <= 32). Higher values allow
102
112
  better pipelining between decode and consumer threads.
103
- Default: lambda min_pics: min_pics * 2 (double the minimum).
113
+ Default: lambda min_pics: min_pics + 4.
114
+ allow_high: Keep >8-bit sources at 16 bits (P016 / YUV444_16Bit
115
+ surfaces) instead of letting NVDEC round them to 8-bit NV12 /
116
+ YUV444. convert() handles both; to_yuv420p() is 8-bit only.
104
117
  """
105
118
 
106
119
  def __init__(
@@ -111,6 +124,7 @@ class Decoder(BaseDecoder):
111
124
  cropping=None,
112
125
  target_rect=None,
113
126
  device=None,
127
+ allow_high=False,
114
128
  ):
115
129
  self.track = track
116
130
  self._slot_locks = None # initialized in decide(), after we know num_pictures
@@ -147,6 +161,10 @@ class Decoder(BaseDecoder):
147
161
  "num_pictures": num_pics,
148
162
  "num_surfaces": 1,
149
163
  }
164
+ if allow_high:
165
+ d["surface_format"] = decide_surface_format(
166
+ p["chroma_format"], p["bit_depth"],
167
+ p["supported_surface_formats"], allow_high=True)
150
168
  if target_size is not None:
151
169
  d["target_size"] = target_size
152
170
  if cropping is not None:
@@ -164,6 +182,11 @@ class Decoder(BaseDecoder):
164
182
  coded_height=track.height,
165
183
  )
166
184
 
185
+ @property
186
+ def duration(self):
187
+ """Track duration as a timedelta, or None if the container doesn't say."""
188
+ return self.track.duration
189
+
167
190
  def pre_decode(self, idx):
168
191
  """NVDEC reclaims slot idx. Clear tracking, wait for consumer if held."""
169
192
  self._slots.pop(idx, None)
@@ -231,8 +254,7 @@ class Decoder(BaseDecoder):
231
254
  packet/frame traversal. Only if avformat cannot seek the container at all do we
232
255
  fall back to the in-memory-index path (then BOF)."""
233
256
  try:
234
- self.track.fc.seek_file(self.track.stream, target_pts,
235
- min_ts=-(2**63), max_ts=target_pts)
257
+ self._native_seek_to_keyframe(target_pts)
236
258
  except AVException:
237
259
  kts = self._keyframe_before(target_pts)
238
260
  if kts is None:
@@ -242,24 +264,39 @@ class Decoder(BaseDecoder):
242
264
  self.with_callbacks(lambda: self.send(None), post_decode=lambda pic: None)
243
265
  self._recent.clear()
244
266
 
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.
267
+ def _native_seek_to_keyframe(self, target_pts, max_tries=16):
268
+ """avformat_seek_file(max_ts=target) positions on the last packet with
269
+ dts <= target. Indexed containers (MP4 / MKV) snap that to a keyframe.
270
+ Index-less ones (MPEG-TS / M2TS) use the generic dts bisection — which
271
+ is NOT keyframe-aware ("FIXME keyframe?" in mpegts.c) — and land
272
+ mid-GOP; NVDEC then discards everything up to the NEXT IDR, so the
273
+ frame we hand back is up to a GOP late. (ffmpeg's CLI papers over this
274
+ with a fixed 3/23 s dts heuristic.)
275
+
276
+ Peek at the landing packet; if it isn't a keyframe, back off by a
277
+ growing number of frame-durations and retry. Landing on an earlier
278
+ keyframe than strictly needed only costs decode-forward time, so
279
+ overshooting backwards is safe. The peek consumes a packet, so the
280
+ final seek re-runs to restore the position. Raises AVException if the
281
+ container can't seek at all (caller falls back to the in-memory index).
249
282
  """
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)
283
+ fc, stream = self.track.fc, self.track.stream
284
+ tb = self.track._time_base
285
+ afr = stream.avg_frame_rate
286
+ fps = Fraction(afr.num, afr.den) if afr.num > 0 and afr.den > 0 else Fraction(25)
287
+ step = max(1, int(1 / (fps * tb))) # one frame, in pts units
288
+ ts = target_pts
289
+ for _ in range(max_tries):
290
+ fc.seek_file(stream, ts, min_ts=-(2**63), max_ts=ts)
291
+ first = next(fc.read_packets(stream), None)
292
+ if first is None or first.av.flags & AV_PKT_FLAG_KEY:
293
+ break # keyframe (or EOF — nothing to fix)
294
+ landed = first.av.dts if first.av.dts != AV_NOPTS_VALUE else first.av.pts
295
+ if landed == AV_NOPTS_VALUE:
296
+ break # can't reason about position
297
+ ts = min(ts, landed) - step
298
+ step *= 2
299
+ fc.seek_file(stream, ts, min_ts=-(2**63), max_ts=ts)
263
300
 
264
301
  def _map_and_convert(self, pic, dtype, pixel_format='rgb'):
265
302
  """Map a Picture (or OwnedPicture) to a tensor. Frees pic and surface.
@@ -268,16 +305,16 @@ class Decoder(BaseDecoder):
268
305
  pixel_format='rgb' (default) → (3,H,W) RGB via convert() (color matrix).
269
306
  pixel_format='yuv420p' → 1D planar I420 uint8 via to_yuv420p() (no color
270
307
  math, half the bytes) — for piping straight to an encoder."""
308
+ t = None
271
309
  try:
272
310
  t = self.track.pts2time(pic.pts)
273
311
  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().
312
+ # Both OwnedPicture (owns its slot) and bare Picture map under
313
+ # the output-surface semaphore; .free() on either is safe.
277
314
  if isinstance(pic, OwnedPicture):
278
315
  surface = pic.map(stream)
279
316
  else:
280
- surface = self.map(pic, stream)
317
+ surface = OwnedSurface(self, pic, stream)
281
318
  except Exception as e:
282
319
  return (t, e)
283
320
  finally:
@@ -683,6 +720,13 @@ class Decoder(BaseDecoder):
683
720
  """Yield (int_pts, frame_or_exception) for every decoded frame in
684
721
  display order. Decodes ALL frames (B-frames included).
685
722
 
723
+ POSITION: without `start` / `start_pts` this does NOT rewind. It
724
+ decodes from wherever the demuxer currently is — the beginning on
725
+ a fresh Decoder, but after a `screenshot()`, `screenshots()` or an
726
+ exhausted `frames()` it resumes from there (and yields nothing if
727
+ that was EOF). Pass `start=timedelta(0)` to decode from the start
728
+ of the file regardless of what ran before.
729
+
686
730
  The yielded `int_pts` is the raw track pts (same units as
687
731
  `pic.pts`, including the demuxer's `start_time` offset). Pass
688
732
  through `track.pts2time(pts)` if you need a `timedelta`. Yielding
@@ -10,6 +10,7 @@ from fractions import Fraction
10
10
  import functools
11
11
 
12
12
  from ..ffmpeg.libavformat import FormatContext, AVMediaType, AVCodecID
13
+ from ..ffmpeg.include.libavformat import AV_DISPOSITION_ATTACHED_PIC
13
14
  from ..ffmpeg.include.libavutil import (
14
15
  AV_NOPTS_VALUE,
15
16
  AV_TIME_BASE,
@@ -36,6 +37,7 @@ class VideoTrack:
36
37
  extradata: Codec extradata bytes, or None.
37
38
  color_space: AVColorSpace enum value.
38
39
  color_range: AVColorRange enum value.
40
+ attached_pic: True if this is embedded cover art, not real video.
39
41
  """
40
42
 
41
43
  def __init__(self, fc, stream):
@@ -51,20 +53,33 @@ class VideoTrack:
51
53
  # Time base
52
54
  self._time_base = Fraction(stream.time_base.num, stream.time_base.den)
53
55
 
54
- # Start time
56
+ # Start time. Lazy containers (MPEG-TS, raw elementary streams) leave
57
+ # both the stream and the format start_time unset until probed;
58
+ # probing buffers the packets it reads, so nothing is lost. (The old
59
+ # path read the first packet directly — which both computed 0 via the
60
+ # packet's unset 0/1 time_base AND swallowed the IDR, dropping the
61
+ # whole first GOP from decode.)
55
62
  self._start_time = stream.start_time
56
63
  if self._start_time == AV_NOPTS_VALUE:
57
- self._start_time = fc.av.start_time
58
- if self._start_time == AV_NOPTS_VALUE:
59
- start_time = fc.infer_start_time()
60
- self._start_time = int(start_time / self._time_base)
61
- else:
64
+ fc.probe()
65
+ self._start_time = stream.start_time
66
+ if self._start_time == AV_NOPTS_VALUE:
67
+ if fc.av.start_time != AV_NOPTS_VALUE:
62
68
  self._start_time = int(
63
- self._start_time / AV_TIME_BASE / self._time_base
69
+ fc.av.start_time / AV_TIME_BASE / self._time_base
64
70
  )
71
+ else:
72
+ # Timestamp-less stream (e.g. raw Annex-B .h264): packets carry
73
+ # no pts at all; frames() yields AV_NOPTS_VALUE and pts2time
74
+ # returns None. Anchor at 0 so time2pts is at least defined.
75
+ self._start_time = 0
65
76
 
66
- # Duration (lazy — probes only if needed)
67
77
 
78
+ @property
79
+ def attached_pic(self):
80
+ """True for cover-art / thumbnail "video" streams (AV_DISPOSITION_ATTACHED_PIC):
81
+ a single still image, not a decodable video track."""
82
+ return bool(self.stream.disposition & AV_DISPOSITION_ATTACHED_PIC)
68
83
 
69
84
  @functools.cached_property
70
85
  def width(self):
@@ -20,8 +20,12 @@ class Player(Decoder):
20
20
  target_rect=None,
21
21
  device=None,
22
22
  track_idx=None,
23
+ **decoder_kwargs,
23
24
  ):
24
- tracks = [t for t in parse(url) if isinstance(t, VideoTrack)]
25
+ # Cover art (attached_pic) is a "video" stream to ffmpeg but a single
26
+ # still image; never pick it, and don't count it as ambiguity.
27
+ tracks = [t for t in parse(url)
28
+ if isinstance(t, VideoTrack) and not t.attached_pic]
25
29
  if not tracks:
26
30
  raise ValueError(f"{url} has no video stream")
27
31
  if track_idx is not None:
@@ -38,5 +42,6 @@ class Player(Decoder):
38
42
  cropping=cropping,
39
43
  target_rect=target_rect,
40
44
  device=device,
45
+ **decoder_kwargs,
41
46
  )
42
47
  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.23
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
@@ -119,6 +124,23 @@ Codec support depends on your GPU:
119
124
  - VPF allocates the bare minimum of resources needed for most decoding tasks. PNC allows the user to specify the amount of resources to be allocated for advanced applications. Users own the resources and decide when and whether to deal with them.
120
125
  - Managing resources is not painful: similar to `pycuda`, we shift the burden of managing host/device resources to the Python garbage collector. Resources (such as `Picture` and `Surface`) are automatically freed when the user drops the reference.
121
126
 
127
+ ## Limitations
128
+
129
+ - **One consumer per `Decoder`.** A `Decoder`/`Player` instance serves one
130
+ `frames()`, `screenshots()` or `screenshot()` call at a time. Use separate
131
+ instances (from any threads) for concurrent decodes of the same or different
132
+ files.
133
+ - **No mid-stream resolution or codec changes.** The NVDEC decoder is created
134
+ from the first sequence header; a later sequence with a different coded
135
+ size, chroma format or bit depth raises `decoder already initialized with
136
+ different format`. Split such files per segment.
137
+ - **Timestamp-less elementary streams** (e.g. raw Annex-B `.h264`/`.hevc`)
138
+ carry no pts: `frames()` yields `AV_NOPTS_VALUE` for every frame,
139
+ `pts2time` returns `None`, and `screenshot()` cannot locate a time. Mux
140
+ into a container first.
141
+ - **Embedded cover art** (`attached_pic` streams) is skipped by `Player`
142
+ and reported by `VideoTrack.attached_pic` from `parse()`.
143
+
122
144
  ## Roadmap
123
145
 
124
146
  - [x] Decoding
@@ -128,7 +150,7 @@ Codec support depends on your GPU:
128
150
  - Color Ranges: MPEG (limited), JPEG (full)
129
151
  - Color Space: BT.601, BT.709, BT.2020
130
152
  - [x] Built-in cropping and scaling via NV decoder
131
- - [x] Thread-safe and thread-friendly
153
+ - [x] Thread-friendly (one consumer per decoder; decoders are independent)
132
154
  - [ ] Encoder
133
155
  - [x] PyTorch integration
134
156
 
@@ -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,188 @@
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 test_duration_property(clip):
90
+ with Player(clip) as p:
91
+ assert abs(p.duration.total_seconds() - SECONDS) < 0.2
92
+
93
+
94
+ def test_cover_art_skipped(clip):
95
+ """A PNG attached_pic stream must not be selected or count as ambiguity."""
96
+ d = tempfile.mkdtemp(prefix="nvc-smoke-cover-")
97
+ try:
98
+ png = os.path.join(d, "cover.png")
99
+ out = os.path.join(d, "cover.mp4")
100
+ subprocess.run(["ffmpeg", "-v", "error", "-y", "-f", "lavfi",
101
+ "-i", "color=c=red:size=64x64:duration=1", "-frames:v", "1", png],
102
+ check=True, timeout=60)
103
+ subprocess.run(["ffmpeg", "-v", "error", "-y", "-i", clip, "-i", png,
104
+ "-map", "0", "-map", "1", "-c", "copy",
105
+ "-disposition:v:1", "attached_pic", out], check=True, timeout=60)
106
+ vids = [t for t in parse(out) if isinstance(t, VideoTrack)]
107
+ assert [t.attached_pic for t in vids] == [False, True]
108
+ with Player(out) as p:
109
+ assert p.track.index == 0
110
+ n = sum(1 for _ in p.frames(torch.uint8))
111
+ assert n == FPS * SECONDS
112
+ finally:
113
+ shutil.rmtree(d, ignore_errors=True)
114
+
115
+
116
+ def test_mpegts_start_time_and_first_gop(clip):
117
+ """MPEG-TS leaves start_time unset until probed. Probing must not swallow
118
+ the IDR (which used to drop the first GOP) and start_time must be real."""
119
+ d = tempfile.mkdtemp(prefix="nvc-smoke-ts-")
120
+ try:
121
+ out = os.path.join(d, "clip.ts")
122
+ subprocess.run(["ffmpeg", "-v", "error", "-y", "-i", clip, "-c", "copy", out],
123
+ check=True, timeout=60)
124
+ with Player(out) as p:
125
+ assert p.track._start_time > 0 # ffmpeg's TS muxer starts at 1.4 s
126
+ frames = list(p.frames(torch.uint8))
127
+ assert len(frames) == FPS * SECONDS
128
+ first_t = p.track.pts2time(frames[0][0])
129
+ assert abs(first_t.total_seconds()) < 1 / FPS
130
+ # TS seeks by dts with no keyframe awareness; the decoder must still
131
+ # produce the frame actually on screen at 1.0 s, not the next IDR.
132
+ t, _ = p.screenshot(timedelta(seconds=1), torch.uint8, accurate=True)
133
+ assert abs(t.total_seconds() - 1.0) < 1 / FPS + 1e-3
134
+ pts0, _ = next(iter(p.frames(torch.uint8, start=timedelta(seconds=1))))
135
+ assert abs(p.track.pts2time(pts0).total_seconds() - 1.0) < 1 / FPS + 1e-3
136
+ finally:
137
+ shutil.rmtree(d, ignore_errors=True)
138
+
139
+
140
+ def test_pixel_format_enum_values():
141
+ from nvidia_codec.ffmpeg.include.libavutil import AVPixelFormat
142
+ assert AVPixelFormat.YUV420P == 0 and AVPixelFormat.RGB24 == 2
143
+ assert isinstance(AVPixelFormat.YUV444P, int)
144
+
145
+
146
+ def _has_encoder(name):
147
+ out = subprocess.run(["ffmpeg", "-hide_banner", "-encoders"],
148
+ capture_output=True, text=True, timeout=30).stdout
149
+ return f" {name} " in out
150
+
151
+
152
+ @pytest.fixture(scope="module")
153
+ def clip10():
154
+ if not _has_encoder("libx265"):
155
+ pytest.skip("ffmpeg lacks libx265")
156
+ d = tempfile.mkdtemp(prefix="nvc-smoke10-")
157
+ path = os.path.join(d, "clip10.mp4")
158
+ subprocess.run(
159
+ ["ffmpeg", "-v", "error", "-y", "-f", "lavfi",
160
+ "-i", f"testsrc=size={W}x{H}:rate={FPS}:duration=1",
161
+ "-pix_fmt", "yuv420p10le", "-c:v", "libx265",
162
+ "-x265-params", "log-level=none", "-g", "12", path],
163
+ check=True, timeout=120,
164
+ )
165
+ yield path
166
+ shutil.rmtree(d, ignore_errors=True)
167
+
168
+
169
+ def test_10bit_allow_high_matches_8bit(clip10):
170
+ """P016 (allow_high) and NV12 paths must agree on the same 10-bit source."""
171
+ with Player(clip10) as p:
172
+ _, lo = p.screenshot(timedelta(seconds=0.5), torch.float32, accurate=True)
173
+ assert p.surface_format == cudaVideoSurfaceFormat.NV12
174
+ with Player(clip10, allow_high=True) as p:
175
+ _, hi = p.screenshot(timedelta(seconds=0.5), torch.float32, accurate=True)
176
+ assert p.surface_format == cudaVideoSurfaceFormat.P016
177
+ assert hi.shape == (3, H, W)
178
+ assert 0.0 <= hi.min() and hi.max() <= 1.0
179
+ # yuv420p export is 8-bit only: the error is reported in-band
180
+ _, r = next(iter(p.frames(torch.uint8, start=timedelta(0),
181
+ pixel_format="yuv420p")))
182
+ assert isinstance(r, ValueError)
183
+ # 10-bit rounds to 8-bit at most 1/255 off, plus matrix rounding
184
+ assert (hi - lo).abs().max() < 0.02
185
+
186
+
187
+ if __name__ == "__main__":
188
+ raise SystemExit(pytest.main([__file__, "-v"]))
File without changes