Skip to content

Commit 8bd28eb

Browse files
WyattBlueWyattBlue
committed
Make video/reformatter pure
1 parent 76d2cb7 commit 8bd28eb

File tree

1 file changed

+85
-57
lines changed

1 file changed

+85
-57
lines changed

av/video/reformatter.pyx renamed to av/video/reformatter.py

Lines changed: 85 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
1-
cimport libav as lib
2-
from libc.stdint cimport uint8_t
3-
4-
from av.error cimport err_check
5-
from av.video.format cimport VideoFormat
6-
from av.video.frame cimport alloc_video_frame
7-
81
from enum import IntEnum
92

3+
import cython
4+
import cython.cimports.libav as lib
5+
from cython.cimports.av.error import err_check
6+
from cython.cimports.av.video.format import VideoFormat
7+
from cython.cimports.av.video.frame import alloc_video_frame
8+
from cython.cimports.libc.stdint import uint8_t
9+
1010

1111
class Interpolation(IntEnum):
1212
FAST_BILINEAR: "Fast bilinear" = lib.SWS_FAST_BILINEAR
@@ -38,6 +38,7 @@ class Colorspace(IntEnum):
3838
smpte240m = lib.SWS_CS_SMPTE240M
3939
default = lib.SWS_CS_DEFAULT
4040

41+
4142
class ColorRange(IntEnum):
4243
UNSPECIFIED: "Unspecified" = lib.AVCOL_RANGE_UNSPECIFIED
4344
MPEG: "MPEG (limited) YUV range, 219*2^(n-8)" = lib.AVCOL_RANGE_MPEG
@@ -58,7 +59,8 @@ def _resolve_enum_value(value, enum_class, default):
5859
raise ValueError(f"Cannot convert {value} to {enum_class.__name__}")
5960

6061

61-
cdef class VideoReformatter:
62+
@cython.cclass
63+
class VideoReformatter:
6264
"""An object for reformatting size and pixel format of :class:`.VideoFrame`.
6365
6466
It is most efficient to have a reformatter object for each set of parameters
@@ -67,13 +69,21 @@ def _resolve_enum_value(value, enum_class, default):
6769
"""
6870

6971
def __dealloc__(self):
70-
with nogil:
72+
with cython.nogil:
7173
lib.sws_freeContext(self.ptr)
7274

73-
def reformat(self, VideoFrame frame not None, width=None, height=None,
74-
format=None, src_colorspace=None, dst_colorspace=None,
75-
interpolation=None, src_color_range=None,
76-
dst_color_range=None):
75+
def reformat(
76+
self,
77+
frame: VideoFrame,
78+
width=None,
79+
height=None,
80+
format=None,
81+
src_colorspace=None,
82+
dst_colorspace=None,
83+
interpolation=None,
84+
src_color_range=None,
85+
dst_color_range=None,
86+
):
7787
"""Create a new :class:`VideoFrame` with the given width/height/format/colorspace.
7888
7989
Returns the same frame untouched if nothing needs to be done to it.
@@ -95,13 +105,24 @@ def reformat(self, VideoFrame frame not None, width=None, height=None,
95105
96106
"""
97107

98-
cdef VideoFormat video_format = VideoFormat(format if format is not None else frame.format)
99-
100-
cdef int c_src_colorspace = _resolve_enum_value(src_colorspace, Colorspace, frame.colorspace)
101-
cdef int c_dst_colorspace = _resolve_enum_value(dst_colorspace, Colorspace, frame.colorspace)
102-
cdef int c_interpolation = _resolve_enum_value(interpolation, Interpolation, int(Interpolation.BILINEAR))
103-
cdef int c_src_color_range = _resolve_enum_value(src_color_range, ColorRange, 0)
104-
cdef int c_dst_color_range = _resolve_enum_value(dst_color_range, ColorRange, 0)
108+
video_format: VideoFormat = VideoFormat(
109+
format if format is not None else frame.format
110+
)
111+
c_src_colorspace: cython.int = _resolve_enum_value(
112+
src_colorspace, Colorspace, frame.colorspace
113+
)
114+
c_dst_colorspace: cython.int = _resolve_enum_value(
115+
dst_colorspace, Colorspace, frame.colorspace
116+
)
117+
c_interpolation: cython.int = _resolve_enum_value(
118+
interpolation, Interpolation, int(Interpolation.BILINEAR)
119+
)
120+
c_src_color_range: cython.int = _resolve_enum_value(
121+
src_color_range, ColorRange, 0
122+
)
123+
c_dst_color_range: cython.int = _resolve_enum_value(
124+
dst_color_range, ColorRange, 0
125+
)
105126

106127
return self._reformat(
107128
frame,
@@ -115,31 +136,39 @@ def reformat(self, VideoFrame frame not None, width=None, height=None,
115136
c_dst_color_range,
116137
)
117138

118-
cdef _reformat(self, VideoFrame frame, int width, int height,
119-
lib.AVPixelFormat dst_format, int src_colorspace,
120-
int dst_colorspace, int interpolation,
121-
int src_color_range, int dst_color_range):
122-
139+
@cython.cfunc
140+
def _reformat(
141+
self,
142+
frame: VideoFrame,
143+
width: cython.int,
144+
height: cython.int,
145+
dst_format: lib.AVPixelFormat,
146+
src_colorspace: cython.int,
147+
dst_colorspace: cython.int,
148+
interpolation: cython.int,
149+
src_color_range: cython.int,
150+
dst_color_range: cython.int,
151+
):
123152
if frame.ptr.format < 0:
124153
raise ValueError("Frame does not have format set.")
125154

126155
# The definition of color range in pixfmt.h and swscale.h is different.
127156
src_color_range = 1 if src_color_range == ColorRange.JPEG.value else 0
128157
dst_color_range = 1 if dst_color_range == ColorRange.JPEG.value else 0
129158

130-
cdef lib.AVPixelFormat src_format = <lib.AVPixelFormat> frame.ptr.format
159+
src_format: lib.AVPixelFormat = cython.cast(lib.AVPixelFormat, frame.ptr.format)
131160

132161
# Shortcut!
133162
if (
134-
dst_format == src_format and
135-
width == frame.ptr.width and
136-
height == frame.ptr.height and
137-
dst_colorspace == src_colorspace and
138-
src_color_range == dst_color_range
163+
dst_format == src_format
164+
and width == frame.ptr.width
165+
and height == frame.ptr.height
166+
and dst_colorspace == src_colorspace
167+
and src_color_range == dst_color_range
139168
):
140169
return frame
141170

142-
with nogil:
171+
with cython.nogil:
143172
self.ptr = lib.sws_getCachedContext(
144173
self.ptr,
145174
frame.ptr.width,
@@ -149,44 +178,45 @@ def reformat(self, VideoFrame frame not None, width=None, height=None,
149178
height,
150179
dst_format,
151180
interpolation,
152-
NULL,
153-
NULL,
154-
NULL
181+
cython.NULL,
182+
cython.NULL,
183+
cython.NULL,
155184
)
156185

157186
# We want to change the colorspace/color_range transforms.
158-
# We do that by grabbing all of the current settings, changing a
187+
# We do that by grabbing all the current settings, changing a
159188
# couple, and setting them all. We need a lot of state here.
160-
cdef int *inv_tbl
161-
cdef int *tbl
162-
cdef int src_colorspace_range, dst_colorspace_range
163-
cdef int brightness, contrast, saturation
164-
cdef int ret
189+
inv_tbl: cython.p_int
190+
tbl: cython.p_int
191+
src_colorspace_range: cython.int
192+
dst_colorspace_range: cython.int
193+
brightness: cython.int
194+
contrast: cython.int
195+
saturation: cython.int
165196

166197
if src_colorspace != dst_colorspace or src_color_range != dst_color_range:
167-
with nogil:
198+
with cython.nogil:
168199
ret = lib.sws_getColorspaceDetails(
169200
self.ptr,
170-
&inv_tbl,
171-
&src_colorspace_range,
172-
&tbl,
173-
&dst_colorspace_range,
174-
&brightness,
175-
&contrast,
176-
&saturation
201+
cython.address(inv_tbl),
202+
cython.address(src_colorspace_range),
203+
cython.address(tbl),
204+
cython.address(dst_colorspace_range),
205+
cython.address(brightness),
206+
cython.address(contrast),
207+
cython.address(saturation),
177208
)
178209

179210
err_check(ret)
180211

181-
with nogil:
212+
with cython.nogil:
182213
# Grab the coefficients for the requested transforms.
183214
# The inv_table brings us to linear, and `tbl` to the new space.
184215
if src_colorspace != lib.SWS_CS_DEFAULT:
185216
inv_tbl = lib.sws_getCoefficients(src_colorspace)
186217
if dst_colorspace != lib.SWS_CS_DEFAULT:
187218
tbl = lib.sws_getCoefficients(dst_colorspace)
188219

189-
# Apply!
190220
ret = lib.sws_setColorspaceDetails(
191221
self.ptr,
192222
inv_tbl,
@@ -195,21 +225,19 @@ def reformat(self, VideoFrame frame not None, width=None, height=None,
195225
dst_color_range,
196226
brightness,
197227
contrast,
198-
saturation
228+
saturation,
199229
)
200230

201231
err_check(ret)
202232

203-
# Create a new VideoFrame.
204-
cdef VideoFrame new_frame = alloc_video_frame()
233+
new_frame: VideoFrame = alloc_video_frame()
205234
new_frame._copy_internal_attributes(frame)
206235
new_frame._init(dst_format, width, height)
207236

208-
# Finally, scale the image.
209-
with nogil:
237+
with cython.nogil:
210238
lib.sws_scale(
211239
self.ptr,
212-
<const uint8_t *const *>frame.ptr.data,
240+
frame.ptr.data,
213241
frame.ptr.linesize,
214242
0, # slice Y
215243
frame.ptr.height,

0 commit comments

Comments
 (0)