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