nvidia-codec 0.0.20__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.20/src/nvidia_codec.egg-info → nvidia_codec-0.0.22}/PKG-INFO +7 -2
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/README.md +6 -1
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/setup.cfg +1 -1
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/__init__.py +1 -1
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/core/cuda.py +1 -1
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/core/decode.py +73 -19
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/core/nvcuvid.py +0 -17
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/libavcodec.py +9 -4
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/libavcodec.py +4 -1
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/libavformat.py +6 -4
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/color.py +40 -3
- nvidia_codec-0.0.22/src/nvidia_codec/utils/decode.py +772 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/demux.py +83 -22
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/player.py +2 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22/src/nvidia_codec.egg-info}/PKG-INFO +7 -2
- {nvidia_codec-0.0.20 → 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.20/src/nvidia_codec/utils/decode.py +0 -474
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/LICENSE +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/pyproject.toml +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/common.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/core/__init__.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/core/common.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/core/cuviddec.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/__init__.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/adapt.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/common.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/__init__.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/libavformat.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/include/libavutil.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/libavutil.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/ffmpeg/test.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/__init__.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec/utils/compat.py +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec.egg-info/dependency_links.txt +0 -0
- {nvidia_codec-0.0.20 → nvidia_codec-0.0.22}/src/nvidia_codec.egg-info/requires.txt +0 -0
- {nvidia_codec-0.0.20 → 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
|
|
@@ -17,6 +17,7 @@ from .common import *
|
|
|
17
17
|
from .. import CodecNotSupportedError
|
|
18
18
|
import numpy as np
|
|
19
19
|
from ctypes import *
|
|
20
|
+
import weakref
|
|
20
21
|
|
|
21
22
|
import logging
|
|
22
23
|
|
|
@@ -36,7 +37,7 @@ class Surface:
|
|
|
36
37
|
self.pts = pts
|
|
37
38
|
|
|
38
39
|
self.params = CUVIDPROCPARAMS()
|
|
39
|
-
memmove(byref(self.params), byref(params), sizeof(
|
|
40
|
+
memmove(byref(self.params), byref(params), sizeof(CUVIDPROCPARAMS))
|
|
40
41
|
|
|
41
42
|
self.params.output_stream = stream
|
|
42
43
|
|
|
@@ -116,19 +117,28 @@ class Surface:
|
|
|
116
117
|
"strides": strides,
|
|
117
118
|
"version": 3,
|
|
118
119
|
"data": (self.c_devptr.value, False),
|
|
119
|
-
|
|
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,
|
|
120
123
|
}
|
|
121
124
|
|
|
122
125
|
|
|
123
126
|
class Picture:
|
|
124
127
|
"""Decoded frame before post-processing.
|
|
125
128
|
|
|
126
|
-
A data object —
|
|
127
|
-
|
|
129
|
+
A pure data object — index, proc params, pts. Deliberately holds NO
|
|
130
|
+
reference to the decoder: the decoder tracks live Pictures in
|
|
131
|
+
_slots/_recent, and a back-reference here would form a cycle that
|
|
132
|
+
keeps the decoder (and its large NVDEC context) alive past the
|
|
133
|
+
refcount-zero point, deferring cleanup to the cyclic GC. To turn a
|
|
134
|
+
Picture into a Surface use Decoder.map(pic) (the decoder is always
|
|
135
|
+
in scope where mapping happens).
|
|
136
|
+
|
|
137
|
+
Slot ownership is managed by the caller via pre_decode/pic_release
|
|
138
|
+
callbacks.
|
|
128
139
|
"""
|
|
129
140
|
|
|
130
|
-
def __init__(self,
|
|
131
|
-
self.decoder = decoder
|
|
141
|
+
def __init__(self, index, proc_params, pts):
|
|
132
142
|
self.index = index
|
|
133
143
|
self.params = proc_params
|
|
134
144
|
self.pts = pts
|
|
@@ -137,9 +147,6 @@ class Picture:
|
|
|
137
147
|
"""No-op. Slot ownership is managed by the caller (player.py)."""
|
|
138
148
|
pass
|
|
139
149
|
|
|
140
|
-
def map(self, stream: int = 0):
|
|
141
|
-
return Surface(self.decoder, self.index, self.params, stream, self.pts)
|
|
142
|
-
|
|
143
150
|
|
|
144
151
|
def decide_surface_format(
|
|
145
152
|
chroma_format, bit_depth, supported_surface_formats, allow_high=False
|
|
@@ -192,12 +199,27 @@ class BaseDecoder:
|
|
|
192
199
|
"""
|
|
193
200
|
|
|
194
201
|
def catch_exception(self, func, return_on_error=0):
|
|
202
|
+
# `func` is a bound method of self, and the wrapper records errors
|
|
203
|
+
# on self.exception — both strong refs to self. The wrapper is
|
|
204
|
+
# wrapped in a CFUNCTYPE and stored on self (self.handleXxxCallback),
|
|
205
|
+
# so capturing self strongly here forms a self-cycle that makes the
|
|
206
|
+
# decoder uncollectable by refcounting (only the cyclic GC can take
|
|
207
|
+
# it — which under streaming load lags and balloons VRAM). Hold
|
|
208
|
+
# self weakly instead; the callback only ever fires during a
|
|
209
|
+
# send()/cuvidParseVideoData while a caller keeps the decoder
|
|
210
|
+
# alive, so the weakref always resolves then. The None branch is
|
|
211
|
+
# just belt-and-suspenders for a callback racing finalization.
|
|
212
|
+
selfref = weakref.ref(self)
|
|
213
|
+
fn = func.__func__
|
|
195
214
|
def wrapper(*args, **kwargs):
|
|
215
|
+
s = selfref()
|
|
216
|
+
if s is None:
|
|
217
|
+
return return_on_error
|
|
196
218
|
try:
|
|
197
|
-
return
|
|
219
|
+
return fn(s, *args, **kwargs)
|
|
198
220
|
except BaseException as e:
|
|
199
221
|
log.debug(f"callback exception logged {e}")
|
|
200
|
-
|
|
222
|
+
s.exception = e
|
|
201
223
|
log.debug(f"returning error code {return_on_error}")
|
|
202
224
|
return return_on_error
|
|
203
225
|
|
|
@@ -410,10 +432,16 @@ class BaseDecoder:
|
|
|
410
432
|
unpaired_field=di.repeat_first_field < 0,
|
|
411
433
|
)
|
|
412
434
|
|
|
413
|
-
picture = Picture(
|
|
435
|
+
picture = Picture(di.picture_index, params, di.timestamp)
|
|
414
436
|
self.post_decode(picture)
|
|
415
437
|
return 1
|
|
416
438
|
|
|
439
|
+
def map(self, pic, stream: int = 0):
|
|
440
|
+
"""Map a decoded Picture to a Surface (post-processed frame).
|
|
441
|
+
Lives on the decoder (not Picture.map) so Picture needn't hold a
|
|
442
|
+
decoder back-reference — see Picture's docstring."""
|
|
443
|
+
return Surface(self, pic.index, pic.params, stream, pic.pts)
|
|
444
|
+
|
|
417
445
|
def handleOperatingPoint(self, pUserData, pOPInfo):
|
|
418
446
|
opi = pOPInfo.contents
|
|
419
447
|
if opi.codec == cudaVideoCodec.AV1:
|
|
@@ -528,18 +556,41 @@ class BaseDecoder:
|
|
|
528
556
|
payload=data.ctypes.data_as(POINTER(c_uint8)),
|
|
529
557
|
timestamp=pts,
|
|
530
558
|
)
|
|
559
|
+
# self.exception relays an error out of a ctypes callback (which
|
|
560
|
+
# can't propagate through the C parser). Clear it as we consume
|
|
561
|
+
# it: a retained exception keeps its __traceback__ alive, whose
|
|
562
|
+
# frames reference self, pinning the decoder. Raise it WITHOUT
|
|
563
|
+
# binding it to a local here — a `e = self.exception; raise e`
|
|
564
|
+
# would put `e` in this frame, and the raised exception's
|
|
565
|
+
# traceback includes this same frame → e ↔ frame cycle, defeating
|
|
566
|
+
# the point. `try: raise self.exception finally: clear` keeps the
|
|
567
|
+
# exception out of this frame's locals.
|
|
531
568
|
self.exception = None
|
|
532
|
-
|
|
533
|
-
|
|
569
|
+
try:
|
|
570
|
+
with cuda.Device(self.device):
|
|
534
571
|
cuda.check(nvcuvid.cuvidParseVideoData(self.cuvid_parser, byref(p)))
|
|
535
|
-
|
|
536
|
-
|
|
537
|
-
raise self.exception
|
|
572
|
+
except cuda.CUError:
|
|
573
|
+
if self.exception is None:
|
|
538
574
|
raise
|
|
575
|
+
# else: a callback captured a more specific exception — it's
|
|
576
|
+
# raised (and cleared) below in preference to the CUError.
|
|
539
577
|
if self.exception is not None:
|
|
540
|
-
|
|
578
|
+
try:
|
|
579
|
+
raise self.exception
|
|
580
|
+
finally:
|
|
581
|
+
self.exception = None
|
|
541
582
|
|
|
542
583
|
def free(self):
|
|
584
|
+
# __del__ can call this during interpreter shutdown, by which
|
|
585
|
+
# point module globals (nvcuvid, cuda, the ctypes types) may be
|
|
586
|
+
# torn down to None — calling through them then raises
|
|
587
|
+
# "NoneType is not callable". The OS reclaims the GPU context on
|
|
588
|
+
# process exit regardless, so just bail if our entry points are
|
|
589
|
+
# gone. (Now reachable because the decoder is refcount-collectable
|
|
590
|
+
# — its self-reference cycles were removed — so __del__ actually
|
|
591
|
+
# runs at shutdown instead of the object leaking.)
|
|
592
|
+
if nvcuvid is None or cuda is None:
|
|
593
|
+
return
|
|
543
594
|
if getattr(self, 'cuvid_parser', None):
|
|
544
595
|
cuda.check(nvcuvid.cuvidDestroyVideoParser(self.cuvid_parser))
|
|
545
596
|
self.cuvid_parser = CUvideoparser()
|
|
@@ -549,7 +600,10 @@ class BaseDecoder:
|
|
|
549
600
|
self.cuvid_decoder = CUvideodecoder()
|
|
550
601
|
|
|
551
602
|
def __del__(self):
|
|
552
|
-
|
|
603
|
+
try:
|
|
604
|
+
self.free()
|
|
605
|
+
except Exception:
|
|
606
|
+
pass # best-effort during finalization / shutdown
|
|
553
607
|
|
|
554
608
|
def __enter__(self):
|
|
555
609
|
return self
|
|
@@ -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()
|
|
@@ -15,6 +15,10 @@ class AVMediaType(TypedCEnumeration(c_int)):
|
|
|
15
15
|
NB = 5
|
|
16
16
|
|
|
17
17
|
class AVCodecID(TypedCEnumeration(c_int)):
|
|
18
|
+
# Values track the ffmpeg 9 enum (libavcodec 63, pinned via the versioned
|
|
19
|
+
# libavcodec.so.63 load in ffmpeg/libavcodec.py — ffmpeg 8 removed entries,
|
|
20
|
+
# shifting WEBP/VP9/HEVC/HQ_HQA/AV1 down). A mismatch surfaces as
|
|
21
|
+
# CodecNotSupportedError ("<AVCodecID.(unknown): N>") on NVDEC-supported files.
|
|
18
22
|
MPEG1 = 1
|
|
19
23
|
MPEG2 = 2
|
|
20
24
|
MPEG4 = 12
|
|
@@ -30,11 +34,12 @@ class AVCodecID(TypedCEnumeration(c_int)):
|
|
|
30
34
|
VP6F = 92
|
|
31
35
|
GIF = 97
|
|
32
36
|
PRORES = 147
|
|
33
|
-
WEBP =
|
|
37
|
+
WEBP = 170
|
|
34
38
|
VP8 = 139
|
|
35
|
-
VP9 =
|
|
36
|
-
HEVC =
|
|
37
|
-
|
|
39
|
+
VP9 = 166
|
|
40
|
+
HEVC = 172
|
|
41
|
+
HQ_HQA = 185
|
|
42
|
+
AV1 = 222
|
|
38
43
|
# Audio codecs — values must match ffmpeg's avcodec.h
|
|
39
44
|
PCM_S16LE = 0x10000
|
|
40
45
|
PCM_S16BE = 0x10001
|
|
@@ -8,7 +8,10 @@ from errno import EAGAIN
|
|
|
8
8
|
import logging
|
|
9
9
|
log = logging.getLogger(__name__)
|
|
10
10
|
|
|
11
|
-
|
|
11
|
+
# Pin the ABI major: enum values + struct layouts in include/libavcodec.py
|
|
12
|
+
# track ffmpeg 9 (libavcodec 63). Loading the versioned .so makes a system
|
|
13
|
+
# ffmpeg upgrade fail loudly here instead of mis-decoding with shifted enums.
|
|
14
|
+
lib = cdll.LoadLibrary('libavcodec.so.63')
|
|
12
15
|
|
|
13
16
|
|
|
14
17
|
av_packet_alloc = lib.av_packet_alloc
|
|
@@ -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
|
|
|
@@ -23,6 +23,44 @@ from ..ffmpeg.include.libavutil import AVColorRange, AVColorSpace
|
|
|
23
23
|
from ..core.cuviddec import cudaVideoSurfaceFormat
|
|
24
24
|
from ..core.decode import Surface
|
|
25
25
|
|
|
26
|
+
def to_yuv420p(surface : Surface):
|
|
27
|
+
"""De-interleave/subsample a decoded surface into a contiguous planar I420
|
|
28
|
+
(yuv420p) uint8 CUDA tensor: Y (H·W) ++ U (H·W/4) ++ V (H·W/4).
|
|
29
|
+
|
|
30
|
+
Unlike convert() — which applies the YUV→RGB color matrix — this does NO
|
|
31
|
+
color math, only plane layout: for NV12 (semi-planar 4:2:0) it splits the
|
|
32
|
+
interleaved chroma; for YUV444 (planar 4:4:4) it box-averages the full-res
|
|
33
|
+
chroma 2×2 down to 4:2:0. 8-bit only (NV12 / YUV444). Half the bytes of rgb24
|
|
34
|
+
and no matrix round-trip; the encoder consumes yuv420p directly. Returns a 1D
|
|
35
|
+
tensor of length H·W·3//2, ready for `.cpu().numpy().tobytes()`."""
|
|
36
|
+
fmt = surface.format
|
|
37
|
+
width, height = surface.width, surface.height
|
|
38
|
+
YUV = torch.as_tensor(surface, device=f'cuda:{surface.decoder.device}')
|
|
39
|
+
if fmt == cudaVideoSurfaceFormat.NV12:
|
|
40
|
+
Y = YUV[:height, :] # (H, W)
|
|
41
|
+
UV = YUV[height:, :] # (H/2, W): U,V interleaved
|
|
42
|
+
U = UV[:, 0::2] # (H/2, W/2)
|
|
43
|
+
V = UV[:, 1::2] # (H/2, W/2)
|
|
44
|
+
elif fmt == cudaVideoSurfaceFormat.YUV444:
|
|
45
|
+
# planar 4:4:4 → 4:2:0: keep Y full-res, box-average chroma 2×2. Plane
|
|
46
|
+
# layout mirrors convert() (3-plane tensor, or 3 planes stacked H·3).
|
|
47
|
+
if YUV.ndim == 3:
|
|
48
|
+
Y, U444, V444 = YUV[0], YUV[1], YUV[2] # each (H, W)
|
|
49
|
+
else:
|
|
50
|
+
Y = YUV[:height, :]
|
|
51
|
+
U444 = YUV[height:height * 2, :]
|
|
52
|
+
V444 = YUV[height * 2:height * 3, :]
|
|
53
|
+
eh, ew = height // 2 * 2, width // 2 * 2
|
|
54
|
+
def _sub(c): # (H, W) 4:4:4 → (H/2, W/2) box-average, uint8
|
|
55
|
+
c = c[:eh, :ew].contiguous().to(torch.float32).reshape(eh // 2, 2, ew // 2, 2)
|
|
56
|
+
return c.mean(dim=(1, 3)).round().clamp(0, 255).to(torch.uint8)
|
|
57
|
+
U, V = _sub(U444), _sub(V444)
|
|
58
|
+
else:
|
|
59
|
+
raise ValueError(f"to_yuv420p: NV12/YUV444 8-bit only, got {fmt}")
|
|
60
|
+
# reshape(-1) copies the (possibly pitch-strided) planes to contiguous.
|
|
61
|
+
return torch.cat([Y.reshape(-1), U.reshape(-1), V.reshape(-1)])
|
|
62
|
+
|
|
63
|
+
|
|
26
64
|
def convert(surface : Surface,
|
|
27
65
|
source_space : AVColorSpace,
|
|
28
66
|
source_range : AVColorRange,
|
|
@@ -65,7 +103,6 @@ def convert(surface : Surface,
|
|
|
65
103
|
assert height % 2 == 0
|
|
66
104
|
assert width % 2 == 0
|
|
67
105
|
assert height // 2 * 3 == h
|
|
68
|
-
assert YUV.dtype == torch.uint8
|
|
69
106
|
|
|
70
107
|
Y = YUV[:height,:]
|
|
71
108
|
UV = YUV[height:,:]
|
|
@@ -96,8 +133,8 @@ def convert(surface : Surface,
|
|
|
96
133
|
#normalize
|
|
97
134
|
if source_range == AVColorRange.MPEG:
|
|
98
135
|
y = (Y.type(torch.float) - (16 << bit8)) / (220 << bit8)
|
|
99
|
-
u = (U.type(torch.float) - (128 << bit8)) / (
|
|
100
|
-
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)
|
|
101
138
|
elif source_range == AVColorRange.JPEG:
|
|
102
139
|
y = Y.type(torch.float) / (256 << bit8)
|
|
103
140
|
u = (U.type(torch.float) - (128 << bit8)) / (256 << bit8)
|