1from libc.errno cimport EAGAIN
2from libc.stdint cimport int64_t, uint8_t
3from libc.stdlib cimport free, malloc, realloc
4from libc.string cimport memcpy
5cimport libav as lib
6
7from av.bytesource cimport ByteSource, bytesource
8from av.codec.codec cimport Codec, wrap_codec
9from av.dictionary cimport _Dictionary
10from av.enum cimport define_enum
11from av.error cimport err_check
12from av.packet cimport Packet
13from av.utils cimport avrational_to_fraction, to_avrational
14
15from av.dictionary import Dictionary
16
17
18cdef object _cinit_sentinel = object()
19
20
21cdef CodecContext wrap_codec_context(lib.AVCodecContext *c_ctx, const lib.AVCodec *c_codec, bint allocated):
22    """Build an av.CodecContext for an existing AVCodecContext."""
23
24    cdef CodecContext py_ctx
25
26    # TODO: This.
27    if c_ctx.codec_type == lib.AVMEDIA_TYPE_VIDEO:
28        from av.video.codeccontext import VideoCodecContext
29        py_ctx = VideoCodecContext(_cinit_sentinel)
30    elif c_ctx.codec_type == lib.AVMEDIA_TYPE_AUDIO:
31        from av.audio.codeccontext import AudioCodecContext
32        py_ctx = AudioCodecContext(_cinit_sentinel)
33    elif c_ctx.codec_type == lib.AVMEDIA_TYPE_SUBTITLE:
34        from av.subtitles.codeccontext import SubtitleCodecContext
35        py_ctx = SubtitleCodecContext(_cinit_sentinel)
36    else:
37        py_ctx = CodecContext(_cinit_sentinel)
38
39    py_ctx.allocated = allocated
40    py_ctx._init(c_ctx, c_codec)
41
42    return py_ctx
43
44
45ThreadType = define_enum('ThreadType', __name__, (
46    ('NONE', 0),
47    ('FRAME', lib.FF_THREAD_FRAME,
48        """Decode more than one frame at once"""),
49    ('SLICE', lib.FF_THREAD_SLICE,
50        """Decode more than one part of a single frame at once"""),
51    ('AUTO', lib.FF_THREAD_SLICE | lib.FF_THREAD_FRAME,
52        """Either method."""),
53), is_flags=True)
54
55SkipType = define_enum('SkipType', __name__, (
56    ('NONE', lib.AVDISCARD_NONE,
57        """Discard nothing"""),
58    ('DEFAULT', lib.AVDISCARD_DEFAULT,
59        """Discard useless packets like 0 size packets in AVI"""),
60    ('NONREF', lib.AVDISCARD_NONREF,
61        """Discard all non reference"""),
62    ('BIDIR', lib.AVDISCARD_BIDIR,
63        """Discard all bidirectional frames"""),
64    ('NONINTRA', lib.AVDISCARD_NONINTRA,
65        """Discard all non intra frames"""),
66    ('NONKEY', lib.AVDISCARD_NONKEY,
67        """Discard all frames except keyframes"""),
68    ('ALL', lib.AVDISCARD_ALL,
69        """Discard all"""),
70))
71
72Flags = define_enum('Flags', __name__, (
73    ('NONE', 0),
74    ('UNALIGNED', lib.AV_CODEC_FLAG_UNALIGNED,
75        """Allow decoders to produce frames with data planes that are not aligned
76        to CPU requirements (e.g. due to cropping)."""),
77    ('QSCALE', lib.AV_CODEC_FLAG_QSCALE,
78        """Use fixed qscale."""),
79    ('4MV', lib.AV_CODEC_FLAG_4MV,
80        """4 MV per MB allowed / advanced prediction for H.263."""),
81    ('OUTPUT_CORRUPT', lib.AV_CODEC_FLAG_OUTPUT_CORRUPT,
82        """Output even those frames that might be corrupted."""),
83    ('QPEL', lib.AV_CODEC_FLAG_QPEL,
84        """Use qpel MC."""),
85    ('DROPCHANGED', 1 << 5,
86        """Don't output frames whose parameters differ from first
87        decoded frame in stream."""),
88    ('PASS1', lib.AV_CODEC_FLAG_PASS1,
89        """Use internal 2pass ratecontrol in first pass mode."""),
90    ('PASS2', lib.AV_CODEC_FLAG_PASS2,
91        """Use internal 2pass ratecontrol in second pass mode."""),
92    ('LOOP_FILTER', lib.AV_CODEC_FLAG_LOOP_FILTER,
93        """loop filter."""),
94    ('GRAY', lib.AV_CODEC_FLAG_GRAY,
95        """Only decode/encode grayscale."""),
96    ('PSNR', lib.AV_CODEC_FLAG_PSNR,
97        """error[?] variables will be set during encoding."""),
98    ('TRUNCATED', lib.AV_CODEC_FLAG_TRUNCATED,
99        """Input bitstream might be truncated at a random location
100        instead of only at frame boundaries."""),
101    ('INTERLACED_DCT', lib.AV_CODEC_FLAG_INTERLACED_DCT,
102        """Use interlaced DCT."""),
103    ('LOW_DELAY', lib.AV_CODEC_FLAG_LOW_DELAY,
104        """Force low delay."""),
105    ('GLOBAL_HEADER', lib.AV_CODEC_FLAG_GLOBAL_HEADER,
106        """Place global headers in extradata instead of every keyframe."""),
107    ('BITEXACT', lib.AV_CODEC_FLAG_BITEXACT,
108        """Use only bitexact stuff (except (I)DCT)."""),
109    ('AC_PRED', lib.AV_CODEC_FLAG_AC_PRED,
110        """H.263 advanced intra coding / MPEG-4 AC prediction"""),
111    ('INTERLACED_ME', lib.AV_CODEC_FLAG_INTERLACED_ME,
112        """Interlaced motion estimation"""),
113    ('CLOSED_GOP', lib.AV_CODEC_FLAG_CLOSED_GOP),
114), is_flags=True)
115
116Flags2 = define_enum('Flags2', __name__, (
117    ('NONE', 0),
118    ('FAST', lib.AV_CODEC_FLAG2_FAST,
119        """Allow non spec compliant speedup tricks."""),
120    ('NO_OUTPUT', lib.AV_CODEC_FLAG2_NO_OUTPUT,
121        """Skip bitstream encoding."""),
122    ('LOCAL_HEADER', lib.AV_CODEC_FLAG2_LOCAL_HEADER,
123        """Place global headers at every keyframe instead of in extradata."""),
124    ('DROP_FRAME_TIMECODE', lib.AV_CODEC_FLAG2_DROP_FRAME_TIMECODE,
125        """Timecode is in drop frame format. DEPRECATED!!!!"""),
126    ('CHUNKS', lib.AV_CODEC_FLAG2_CHUNKS,
127        """Input bitstream might be truncated at a packet boundaries
128        instead of only at frame boundaries."""),
129    ('IGNORE_CROP', lib.AV_CODEC_FLAG2_IGNORE_CROP,
130        """Discard cropping information from SPS."""),
131    ('SHOW_ALL', lib.AV_CODEC_FLAG2_SHOW_ALL,
132        """Show all frames before the first keyframe"""),
133    ('EXPORT_MVS', lib.AV_CODEC_FLAG2_EXPORT_MVS,
134        """Export motion vectors through frame side data"""),
135    ('SKIP_MANUAL', lib.AV_CODEC_FLAG2_SKIP_MANUAL,
136        """Do not skip samples and export skip information as frame side data"""),
137    ('RO_FLUSH_NOOP', lib.AV_CODEC_FLAG2_RO_FLUSH_NOOP,
138        """Do not reset ASS ReadOrder field on flush (subtitles decoding)"""),
139), is_flags=True)
140
141
142cdef class CodecContext(object):
143
144    @staticmethod
145    def create(codec, mode=None):
146        cdef Codec cy_codec = codec if isinstance(codec, Codec) else Codec(codec, mode)
147        cdef lib.AVCodecContext *c_ctx = lib.avcodec_alloc_context3(cy_codec.ptr)
148        return wrap_codec_context(c_ctx, cy_codec.ptr, True)
149
150    def __cinit__(self, sentinel=None, *args, **kwargs):
151        if sentinel is not _cinit_sentinel:
152            raise RuntimeError('Cannot instantiate CodecContext')
153
154        self.options = {}
155        self.stream_index = -1  # This is set by the container immediately.
156
157    cdef _init(self, lib.AVCodecContext *ptr, const lib.AVCodec *codec):
158
159        self.ptr = ptr
160        if self.ptr.codec and codec and self.ptr.codec != codec:
161            raise RuntimeError('Wrapping CodecContext with mismatched codec.')
162        self.codec = wrap_codec(codec if codec != NULL else self.ptr.codec)
163
164        # Set reasonable threading defaults.
165        # count == 0 -> use as many threads as there are CPUs.
166        # type == 2 -> thread within a frame. This does not change the API.
167        self.ptr.thread_count = 0
168        self.ptr.thread_type = 2
169
170    def _get_flags(self):
171        return self.ptr.flags
172
173    def _set_flags(self, value):
174        self.ptr.flags = value
175
176    flags = Flags.property(
177        _get_flags,
178        _set_flags,
179        """Flag property of :class:`.Flags`."""
180    )
181
182    unaligned = flags.flag_property('UNALIGNED')
183    qscale = flags.flag_property('QSCALE')
184    four_mv = flags.flag_property('4MV')
185    output_corrupt = flags.flag_property('OUTPUT_CORRUPT')
186    qpel = flags.flag_property('QPEL')
187    drop_changed = flags.flag_property('DROPCHANGED')
188    pass1 = flags.flag_property('PASS1')
189    pass2 = flags.flag_property('PASS2')
190    loop_filter = flags.flag_property('LOOP_FILTER')
191    gray = flags.flag_property('GRAY')
192    psnr = flags.flag_property('PSNR')
193    truncated = flags.flag_property('TRUNCATED')
194    interlaced_dct = flags.flag_property('INTERLACED_DCT')
195    low_delay = flags.flag_property('LOW_DELAY')
196    global_header = flags.flag_property('GLOBAL_HEADER')
197    bitexact = flags.flag_property('BITEXACT')
198    ac_pred = flags.flag_property('AC_PRED')
199    interlaced_me = flags.flag_property('INTERLACED_ME')
200    closed_gop = flags.flag_property('CLOSED_GOP')
201
202    def _get_flags2(self):
203        return self.ptr.flags2
204
205    def _set_flags2(self, value):
206        self.ptr.flags2 = value
207
208    flags2 = Flags2.property(
209        _get_flags2,
210        _set_flags2,
211        """Flag property of :class:`.Flags2`."""
212    )
213
214    fast = flags2.flag_property('FAST')
215    no_output = flags2.flag_property('NO_OUTPUT')
216    local_header = flags2.flag_property('LOCAL_HEADER')
217    drop_frame_timecode = flags2.flag_property('DROP_FRAME_TIMECODE')
218    chunks = flags2.flag_property('CHUNKS')
219    ignore_crop = flags2.flag_property('IGNORE_CROP')
220    show_all = flags2.flag_property('SHOW_ALL')
221    export_mvs = flags2.flag_property('EXPORT_MVS')
222    skip_manual = flags2.flag_property('SKIP_MANUAL')
223    ro_flush_noop = flags2.flag_property('RO_FLUSH_NOOP')
224
225    property extradata:
226        def __get__(self):
227            if self.ptr.extradata_size > 0:
228                return <bytes>(<uint8_t*>self.ptr.extradata)[:self.ptr.extradata_size]
229            else:
230                return None
231
232        def __set__(self, data):
233            self.extradata_source = bytesource(data)
234            self.ptr.extradata = self.extradata_source.ptr
235            self.ptr.extradata_size = self.extradata_source.length
236
237    property extradata_size:
238        def __get__(self):
239            return self.ptr.extradata_size
240
241    property is_open:
242        def __get__(self):
243            return lib.avcodec_is_open(self.ptr)
244
245    property is_encoder:
246        def __get__(self):
247            return lib.av_codec_is_encoder(self.ptr.codec)
248
249    property is_decoder:
250        def __get__(self):
251            return lib.av_codec_is_decoder(self.ptr.codec)
252
253    cpdef open(self, bint strict=True):
254
255        if lib.avcodec_is_open(self.ptr):
256            if strict:
257                raise ValueError('CodecContext is already open.')
258            return
259
260        # We might pass partial frames.
261        # TODO: What is this for?! This is causing problems with raw decoding
262        # as the internal parser doesn't seem to see a frame until it sees
263        # the next one.
264        # if self.codec.ptr.capabilities & lib.CODEC_CAP_TRUNCATED:
265        #     self.ptr.flags |= lib.CODEC_FLAG_TRUNCATED
266
267        # TODO: Do this better.
268        cdef _Dictionary options = Dictionary()
269        options.update(self.options or {})
270
271        # Assert we have a time_base.
272        if not self.ptr.time_base.num:
273            self._set_default_time_base()
274
275        err_check(lib.avcodec_open2(self.ptr, self.codec.ptr, &options.ptr))
276
277        self.options = dict(options)
278
279    cdef _set_default_time_base(self):
280        self.ptr.time_base.num = 1
281        self.ptr.time_base.den = lib.AV_TIME_BASE
282
283    cpdef close(self, bint strict=True):
284        if not lib.avcodec_is_open(self.ptr):
285            if strict:
286                raise ValueError('CodecContext is already closed.')
287            return
288        err_check(lib.avcodec_close(self.ptr))
289
290    def __dealloc__(self):
291        if self.ptr and self.allocated:
292            lib.avcodec_close(self.ptr)
293            lib.avcodec_free_context(&self.ptr)
294        if self.parser:
295            lib.av_parser_close(self.parser)
296
297    def __repr__(self):
298        return '<av.%s %s/%s at 0x%x>' % (
299            self.__class__.__name__,
300            self.type or '<notype>',
301            self.name or '<nocodec>',
302            id(self),
303        )
304
305    def parse(self, raw_input=None):
306        """Split up a byte stream into list of :class:`.Packet`.
307
308        This is only effectively splitting up a byte stream, and does no
309        actual interpretation of the data.
310
311        It will return all packets that are fully contained within the given
312        input, and will buffer partial packets until they are complete.
313
314        :param ByteSource raw_input: A chunk of a byte-stream to process.
315            Anything that can be turned into a :class:`.ByteSource` is fine.
316            ``None`` or empty inputs will flush the parser's buffers.
317
318        :return: ``list`` of :class:`.Packet` newly availible.
319
320        """
321
322        if not self.parser:
323            self.parser = lib.av_parser_init(self.codec.ptr.id)
324            if not self.parser:
325                raise ValueError('No parser for %s' % self.codec.name)
326
327        cdef ByteSource source = bytesource(raw_input, allow_none=True)
328
329        cdef unsigned char *in_data = source.ptr if source is not None else NULL
330        cdef int in_size = source.length if source is not None else 0
331
332        cdef unsigned char *out_data
333        cdef int out_size
334        cdef int consumed
335        cdef Packet packet = None
336
337        packets = []
338
339        while True:
340
341            with nogil:
342                consumed = lib.av_parser_parse2(
343                    self.parser,
344                    self.ptr,
345                    &out_data, &out_size,
346                    in_data, in_size,
347                    lib.AV_NOPTS_VALUE, lib.AV_NOPTS_VALUE,
348                    0
349                )
350            err_check(consumed)
351
352            if out_size:
353
354                # We copy the data immediately, as we have yet to figure out
355                # the expected lifetime of the buffer we get back. All of the
356                # examples decode it immediately.
357                #
358                # We've also tried:
359                #   packet = Packet()
360                #   packet.data = out_data
361                #   packet.size = out_size
362                #   packet.source = source
363                #
364                # ... but this results in corruption.
365
366                packet = Packet(out_size)
367                memcpy(packet.struct.data, out_data, out_size)
368
369                packets.append(packet)
370
371            if not in_size:
372                # This was a flush. Only one packet should ever be returned.
373                break
374
375            in_data += consumed
376            in_size -= consumed
377
378            if not in_size:
379                # Aaaand now we're done.
380                break
381
382        return packets
383
384    cdef _send_frame_and_recv(self, Frame frame):
385
386        cdef Packet packet
387
388        cdef int res
389        with nogil:
390            res = lib.avcodec_send_frame(self.ptr, frame.ptr if frame is not None else NULL)
391        err_check(res)
392
393        out = []
394        while True:
395            packet = self._recv_packet()
396            if packet:
397                out.append(packet)
398            else:
399                break
400        return out
401
402    cdef _send_packet_and_recv(self, Packet packet):
403
404        cdef Frame frame
405
406        cdef int res
407        with nogil:
408            res = lib.avcodec_send_packet(self.ptr, &packet.struct if packet is not None else NULL)
409        err_check(res)
410
411        out = []
412        while True:
413            frame = self._recv_frame()
414            if frame:
415                out.append(frame)
416            else:
417                break
418        return out
419
420    cdef _prepare_frames_for_encode(self, Frame frame):
421        return [frame]
422
423    cdef Frame _alloc_next_frame(self):
424        raise NotImplementedError('Base CodecContext cannot decode.')
425
426    cdef _recv_frame(self):
427
428        if not self._next_frame:
429            self._next_frame = self._alloc_next_frame()
430        cdef Frame frame = self._next_frame
431
432        cdef int res
433        with nogil:
434            res = lib.avcodec_receive_frame(self.ptr, frame.ptr)
435
436        if res == -EAGAIN or res == lib.AVERROR_EOF:
437            return
438        err_check(res)
439
440        if not res:
441            self._next_frame = None
442            return frame
443
444    cdef _recv_packet(self):
445
446        cdef Packet packet = Packet()
447
448        cdef int res
449        with nogil:
450            res = lib.avcodec_receive_packet(self.ptr, &packet.struct)
451        if res == -EAGAIN or res == lib.AVERROR_EOF:
452            return
453        err_check(res)
454
455        if not res:
456            return packet
457
458    cpdef encode(self, Frame frame=None):
459        """Encode a list of :class:`.Packet` from the given :class:`.Frame`."""
460
461        if self.ptr.codec_type not in [lib.AVMEDIA_TYPE_VIDEO, lib.AVMEDIA_TYPE_AUDIO]:
462            raise NotImplementedError('Encoding is only supported for audio and video.')
463
464        self.open(strict=False)
465
466        frames = self._prepare_frames_for_encode(frame)
467
468        # Assert the frames are in our time base.
469        # TODO: Don't mutate time.
470        for frame in frames:
471            if frame is not None:
472                frame._rebase_time(self.ptr.time_base)
473
474        res = []
475        for frame in frames:
476            for packet in self._send_frame_and_recv(frame):
477                self._setup_encoded_packet(packet)
478                res.append(packet)
479        return res
480
481    cdef _setup_encoded_packet(self, Packet packet):
482        # We coerced the frame's time_base into the CodecContext's during encoding,
483        # and FFmpeg copied the frame's pts/dts to the packet, so keep track of
484        # this time_base in case the frame needs to be muxed to a container with
485        # a different time_base.
486        #
487        # NOTE: if the CodecContext's time_base is altered during encoding, all bets
488        # are off!
489        packet._time_base = self.ptr.time_base
490
491    cpdef decode(self, Packet packet=None):
492        """Decode a list of :class:`.Frame` from the given :class:`.Packet`.
493
494        If the packet is None, the buffers will be flushed. This is useful if
495        you do not want the library to automatically re-order frames for you
496        (if they are encoded with a codec that has B-frames).
497
498        """
499
500        if not self.codec.ptr:
501            raise ValueError('cannot decode unknown codec')
502
503        self.open(strict=False)
504
505        res = []
506        for frame in self._send_packet_and_recv(packet):
507            if isinstance(frame, Frame):
508                self._setup_decoded_frame(frame, packet)
509            res.append(frame)
510        return res
511
512    cdef _setup_decoded_frame(self, Frame frame, Packet packet):
513
514        # Propagate our manual times.
515        # While decoding, frame times are in stream time_base, which PyAV
516        # is carrying around.
517        # TODO: Somehow get this from the stream so we can not pass the
518        # packet here (because flushing packets are bogus).
519        frame._time_base = packet._time_base
520
521        frame.index = self.ptr.frame_number - 1
522
523    property name:
524        def __get__(self):
525            return self.codec.name
526
527    property type:
528        def __get__(self):
529            return self.codec.type
530
531    property profile:
532        def __get__(self):
533            if self.ptr.codec and lib.av_get_profile_name(self.ptr.codec, self.ptr.profile):
534                return lib.av_get_profile_name(self.ptr.codec, self.ptr.profile)
535
536    property time_base:
537        def __get__(self):
538            return avrational_to_fraction(&self.ptr.time_base)
539
540        def __set__(self, value):
541            to_avrational(value, &self.ptr.time_base)
542
543    property ticks_per_frame:
544        def __get__(self):
545            return self.ptr.ticks_per_frame
546
547    property bit_rate:
548        def __get__(self):
549            return self.ptr.bit_rate if self.ptr.bit_rate > 0 else None
550
551        def __set__(self, int value):
552            self.ptr.bit_rate = value
553
554    property max_bit_rate:
555        def __get__(self):
556            if self.ptr.rc_max_rate > 0:
557                return self.ptr.rc_max_rate
558            else:
559                return None
560
561    property bit_rate_tolerance:
562        def __get__(self):
563            self.ptr.bit_rate_tolerance
564
565        def __set__(self, int value):
566            self.ptr.bit_rate_tolerance = value
567
568    property thread_count:
569        """How many threads to use; 0 means auto.
570
571        Wraps :ffmpeg:`AVCodecContext.thread_count`.
572
573        """
574
575        def __get__(self):
576            return self.ptr.thread_count
577
578        def __set__(self, int value):
579            if lib.avcodec_is_open(self.ptr):
580                raise RuntimeError("Cannot change thread_count after codec is open.")
581            self.ptr.thread_count = value
582
583    property thread_type:
584        """One of :class:`.ThreadType`.
585
586        Wraps :ffmpeg:`AVCodecContext.thread_type`.
587
588        """
589
590        def __get__(self):
591            return ThreadType.get(self.ptr.thread_type, create=True)
592
593        def __set__(self, value):
594            if lib.avcodec_is_open(self.ptr):
595                raise RuntimeError("Cannot change thread_type after codec is open.")
596            self.ptr.thread_type = ThreadType[value].value
597
598    property skip_frame:
599        """One of :class:`.SkipType`.
600
601        Wraps ffmpeg:`AVCodecContext.skip_frame`.
602
603        """
604
605        def __get__(self):
606            return SkipType._get(self.ptr.skip_frame, create=True)
607
608        def __set__(self, value):
609            self.ptr.skip_frame = SkipType[value].value
610