1from libc.stdint cimport uint8_t
2cimport libav as lib
3
4from av.enum cimport define_enum
5from av.error cimport err_check
6from av.video.format cimport VideoFormat
7from av.video.frame cimport alloc_video_frame
8
9
10Interpolation = define_enum('Interpolation', __name__, (
11    ('FAST_BILINEAR', lib.SWS_FAST_BILINEAR, "Fast bilinear"),
12    ('BILINEAR', lib.SWS_BILINEAR, "Bilinear"),
13    ('BICUBIC', lib.SWS_BICUBIC, "Bicubic"),
14    ('X', lib.SWS_X, "Experimental"),
15    ('POINT', lib.SWS_POINT, "Nearest neighbor / point"),
16    ('AREA', lib.SWS_AREA, "Area averaging"),
17    ('BICUBLIN', lib.SWS_BICUBLIN, "Luma bicubic / chroma bilinear"),
18    ('GAUSS', lib.SWS_GAUSS, "Gaussian"),
19    ('SINC', lib.SWS_SINC, "Sinc"),
20    ('LANCZOS', lib.SWS_LANCZOS, "Lanczos"),
21    ('SPLINE', lib.SWS_SPLINE, "Bicubic spline"),
22))
23
24Colorspace = define_enum('Colorspace', __name__, (
25
26    ('ITU709', lib.SWS_CS_ITU709),
27    ('FCC', lib.SWS_CS_FCC),
28    ('ITU601', lib.SWS_CS_ITU601),
29    ('ITU624', lib.SWS_CS_ITU624),
30    ('SMPTE170M', lib.SWS_CS_SMPTE170M),
31    ('SMPTE240M', lib.SWS_CS_SMPTE240M),
32    ('DEFAULT', lib.SWS_CS_DEFAULT),
33
34    # Lowercase for b/c.
35    ('itu709', lib.SWS_CS_ITU709),
36    ('fcc', lib.SWS_CS_FCC),
37    ('itu601', lib.SWS_CS_ITU601),
38    ('itu624', lib.SWS_CS_SMPTE170M),
39    ('smpte240', lib.SWS_CS_SMPTE240M),
40    ('default', lib.SWS_CS_DEFAULT),
41
42))
43
44
45cdef class VideoReformatter(object):
46
47    """An object for reformatting size and pixel format of :class:`.VideoFrame`.
48
49    It is most efficient to have a reformatter object for each set of parameters
50    you will use as calling :meth:`reformat` will reconfigure the internal object.
51
52    """
53
54    def __dealloc__(self):
55        with nogil:
56            lib.sws_freeContext(self.ptr)
57
58    def reformat(self, VideoFrame frame not None, width=None, height=None,
59                 format=None, src_colorspace=None, dst_colorspace=None,
60                 interpolation=None):
61        """Create a new :class:`VideoFrame` with the given width/height/format/colorspace.
62
63        Returns the same frame untouched if nothing needs to be done to it.
64
65        :param int width: New width, or ``None`` for the same width.
66        :param int height: New height, or ``None`` for the same height.
67        :param format: New format, or ``None`` for the same format.
68        :type  format: :class:`.VideoFormat` or ``str``
69        :param src_colorspace: Current colorspace, or ``None`` for ``DEFAULT``.
70        :type  src_colorspace: :class:`Colorspace` or ``str``
71        :param dst_colorspace: Desired colorspace, or ``None`` for ``DEFAULT``.
72        :type  dst_colorspace: :class:`Colorspace` or ``str``
73        :param interpolation: The interpolation method to use, or ``None`` for ``BILINEAR``.
74        :type  interpolation: :class:`Interpolation` or ``str``
75
76        """
77
78        cdef VideoFormat video_format = VideoFormat(format if format is not None else frame.format)
79        cdef int c_src_colorspace = (Colorspace[src_colorspace] if src_colorspace is not None else Colorspace.DEFAULT).value
80        cdef int c_dst_colorspace = (Colorspace[dst_colorspace] if dst_colorspace is not None else Colorspace.DEFAULT).value
81        cdef int c_interpolation = (Interpolation[interpolation] if interpolation is not None else Interpolation.BILINEAR).value
82
83        return self._reformat(
84            frame,
85            width or frame.ptr.width,
86            height or frame.ptr.height,
87            video_format.pix_fmt,
88            c_src_colorspace,
89            c_dst_colorspace,
90            c_interpolation,
91        )
92
93    cdef _reformat(self, VideoFrame frame, int width, int height,
94                   lib.AVPixelFormat dst_format, int src_colorspace,
95                   int dst_colorspace, int interpolation):
96
97        if frame.ptr.format < 0:
98            raise ValueError("Frame does not have format set.")
99
100        cdef lib.AVPixelFormat src_format = <lib.AVPixelFormat> frame.ptr.format
101
102        # Shortcut!
103        if (
104            dst_format == src_format and
105            width == frame.ptr.width and
106            height == frame.ptr.height and
107            dst_colorspace == src_colorspace
108        ):
109            return frame
110
111        # Try and reuse existing SwsContextProxy
112        # VideoStream.decode will copy its SwsContextProxy to VideoFrame
113        # So all Video frames from the same VideoStream should have the same one
114        with nogil:
115            self.ptr = lib.sws_getCachedContext(
116                self.ptr,
117                frame.ptr.width,
118                frame.ptr.height,
119                src_format,
120                width,
121                height,
122                dst_format,
123                interpolation,
124                NULL,
125                NULL,
126                NULL
127            )
128
129        # We want to change the colorspace transforms. We do that by grabbing
130        # all of the current settings, changing a couple, and setting them all.
131        # We need a lot of state here.
132        cdef const int *inv_tbl
133        cdef const int *tbl
134        cdef int src_range, dst_range, brightness, contrast, saturation
135        cdef int ret
136        if src_colorspace != dst_colorspace:
137
138            with nogil:
139
140                # Casts for const-ness, because Cython isn't expressive enough.
141                ret = lib.sws_getColorspaceDetails(
142                    self.ptr,
143                    <int**>&inv_tbl,
144                    &src_range,
145                    <int**>&tbl,
146                    &dst_range,
147                    &brightness,
148                    &contrast,
149                    &saturation
150                )
151
152            err_check(ret)
153
154            with nogil:
155
156                # Grab the coefficients for the requested transforms.
157                # The inv_table brings us to linear, and `tbl` to the new space.
158                if src_colorspace != lib.SWS_CS_DEFAULT:
159                    inv_tbl = lib.sws_getCoefficients(src_colorspace)
160                if dst_colorspace != lib.SWS_CS_DEFAULT:
161                    tbl = lib.sws_getCoefficients(dst_colorspace)
162
163                # Apply!
164                ret = lib.sws_setColorspaceDetails(
165                    self.ptr,
166                    inv_tbl,
167                    src_range,
168                    tbl,
169                    dst_range,
170                    brightness,
171                    contrast,
172                    saturation
173                )
174
175            err_check(ret)
176
177        # Create a new VideoFrame.
178        cdef VideoFrame new_frame = alloc_video_frame()
179        new_frame._copy_internal_attributes(frame)
180        new_frame._init(dst_format, width, height)
181
182        # Finally, scale the image.
183        with nogil:
184            lib.sws_scale(
185                self.ptr,
186                # Cast for const-ness, because Cython isn't expressive enough.
187                <const uint8_t**>frame.ptr.data,
188                frame.ptr.linesize,
189                0,  # slice Y
190                frame.ptr.height,
191                new_frame.ptr.data,
192                new_frame.ptr.linesize,
193            )
194
195        return new_frame
196