Skip to content

Commit d8d7602

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

File tree

1 file changed

+84
-59
lines changed

1 file changed

+84
-59
lines changed

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

Lines changed: 84 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
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+
109

1110
class Interpolation(IntEnum):
1211
FAST_BILINEAR: "Fast bilinear" = lib.SWS_FAST_BILINEAR
@@ -38,6 +37,7 @@ class Colorspace(IntEnum):
3837
smpte240m = lib.SWS_CS_SMPTE240M
3938
default = lib.SWS_CS_DEFAULT
4039

40+
4141
class ColorRange(IntEnum):
4242
UNSPECIFIED: "Unspecified" = lib.AVCOL_RANGE_UNSPECIFIED
4343
MPEG: "MPEG (limited) YUV range, 219*2^(n-8)" = lib.AVCOL_RANGE_MPEG
@@ -58,7 +58,8 @@ def _resolve_enum_value(value, enum_class, default):
5858
raise ValueError(f"Cannot convert {value} to {enum_class.__name__}")
5959

6060

61-
cdef class VideoReformatter:
61+
@cython.cclass
62+
class VideoReformatter:
6263
"""An object for reformatting size and pixel format of :class:`.VideoFrame`.
6364
6465
It is most efficient to have a reformatter object for each set of parameters
@@ -67,13 +68,21 @@ def _resolve_enum_value(value, enum_class, default):
6768
"""
6869

6970
def __dealloc__(self):
70-
with nogil:
71+
with cython.nogil:
7172
lib.sws_freeContext(self.ptr)
7273

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):
74+
def reformat(
75+
self,
76+
frame: VideoFrame,
77+
width=None,
78+
height=None,
79+
format=None,
80+
src_colorspace=None,
81+
dst_colorspace=None,
82+
interpolation=None,
83+
src_color_range=None,
84+
dst_color_range=None,
85+
):
7786
"""Create a new :class:`VideoFrame` with the given width/height/format/colorspace.
7887
7988
Returns the same frame untouched if nothing needs to be done to it.
@@ -95,13 +104,24 @@ def reformat(self, VideoFrame frame not None, width=None, height=None,
95104
96105
"""
97106

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)
107+
video_format: VideoFormat = VideoFormat(
108+
format if format is not None else frame.format
109+
)
110+
c_src_colorspace: cython.int = _resolve_enum_value(
111+
src_colorspace, Colorspace, frame.colorspace
112+
)
113+
c_dst_colorspace: cython.int = _resolve_enum_value(
114+
dst_colorspace, Colorspace, frame.colorspace
115+
)
116+
c_interpolation: cython.int = _resolve_enum_value(
117+
interpolation, Interpolation, int(Interpolation.BILINEAR)
118+
)
119+
c_src_color_range: cython.int = _resolve_enum_value(
120+
src_color_range, ColorRange, 0
121+
)
122+
c_dst_color_range: cython.int = _resolve_enum_value(
123+
dst_color_range, ColorRange, 0
124+
)
105125

106126
return self._reformat(
107127
frame,
@@ -115,31 +135,39 @@ def reformat(self, VideoFrame frame not None, width=None, height=None,
115135
c_dst_color_range,
116136
)
117137

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-
138+
@cython.cfunc
139+
def _reformat(
140+
self,
141+
frame: VideoFrame,
142+
width: cython.int,
143+
height: cython.int,
144+
dst_format: lib.AVPixelFormat,
145+
src_colorspace: cython.int,
146+
dst_colorspace: cython.int,
147+
interpolation: cython.int,
148+
src_color_range: cython.int,
149+
dst_color_range: cython.int,
150+
):
123151
if frame.ptr.format < 0:
124152
raise ValueError("Frame does not have format set.")
125153

126154
# The definition of color range in pixfmt.h and swscale.h is different.
127155
src_color_range = 1 if src_color_range == ColorRange.JPEG.value else 0
128156
dst_color_range = 1 if dst_color_range == ColorRange.JPEG.value else 0
129157

130-
cdef lib.AVPixelFormat src_format = <lib.AVPixelFormat> frame.ptr.format
158+
src_format = cython.cast(lib.AVPixelFormat, frame.ptr.format)
131159

132160
# Shortcut!
133161
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
162+
dst_format == src_format
163+
and width == frame.ptr.width
164+
and height == frame.ptr.height
165+
and dst_colorspace == src_colorspace
166+
and src_color_range == dst_color_range
139167
):
140168
return frame
141169

142-
with nogil:
170+
with cython.nogil:
143171
self.ptr = lib.sws_getCachedContext(
144172
self.ptr,
145173
frame.ptr.width,
@@ -149,44 +177,44 @@ def reformat(self, VideoFrame frame not None, width=None, height=None,
149177
height,
150178
dst_format,
151179
interpolation,
152-
NULL,
153-
NULL,
154-
NULL
180+
cython.NULL,
181+
cython.NULL,
182+
cython.NULL,
155183
)
156184

157185
# We want to change the colorspace/color_range transforms.
158-
# We do that by grabbing all of the current settings, changing a
186+
# We do that by grabbing all the current settings, changing a
159187
# 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
188+
inv_tbl: cython.p_int
189+
tbl: cython.p_int
190+
src_colorspace_range: cython.int
191+
dst_colorspace_range: cython.int
192+
brightness: cython.int
193+
contrast: cython.int
194+
saturation: cython.int
165195

166196
if src_colorspace != dst_colorspace or src_color_range != dst_color_range:
167-
with nogil:
197+
with cython.nogil:
168198
ret = lib.sws_getColorspaceDetails(
169199
self.ptr,
170-
&inv_tbl,
171-
&src_colorspace_range,
172-
&tbl,
173-
&dst_colorspace_range,
174-
&brightness,
175-
&contrast,
176-
&saturation
200+
cython.address(inv_tbl),
201+
cython.address(src_colorspace_range),
202+
cython.address(tbl),
203+
cython.address(dst_colorspace_range),
204+
cython.address(brightness),
205+
cython.address(contrast),
206+
cython.address(saturation),
177207
)
178-
179208
err_check(ret)
180209

181-
with nogil:
210+
with cython.nogil:
182211
# Grab the coefficients for the requested transforms.
183212
# The inv_table brings us to linear, and `tbl` to the new space.
184213
if src_colorspace != lib.SWS_CS_DEFAULT:
185214
inv_tbl = lib.sws_getCoefficients(src_colorspace)
186215
if dst_colorspace != lib.SWS_CS_DEFAULT:
187216
tbl = lib.sws_getCoefficients(dst_colorspace)
188217

189-
# Apply!
190218
ret = lib.sws_setColorspaceDetails(
191219
self.ptr,
192220
inv_tbl,
@@ -195,21 +223,18 @@ def reformat(self, VideoFrame frame not None, width=None, height=None,
195223
dst_color_range,
196224
brightness,
197225
contrast,
198-
saturation
226+
saturation,
199227
)
200-
201228
err_check(ret)
202229

203-
# Create a new VideoFrame.
204-
cdef VideoFrame new_frame = alloc_video_frame()
230+
new_frame: VideoFrame = alloc_video_frame()
205231
new_frame._copy_internal_attributes(frame)
206232
new_frame._init(dst_format, width, height)
207233

208-
# Finally, scale the image.
209-
with nogil:
234+
with cython.nogil:
210235
lib.sws_scale(
211236
self.ptr,
212-
<const uint8_t *const *>frame.ptr.data,
237+
frame.ptr.data,
213238
frame.ptr.linesize,
214239
0, # slice Y
215240
frame.ptr.height,

0 commit comments

Comments
 (0)