test_encode.py

from __future__ import division

from fractions import Fraction
from unittest import SkipTest
import math

from av import AudioFrame, VideoFrame
from av.audio.stream import AudioStream
from av.video.stream import VideoStream
import av

from .common import Image, TestCase, fate_suite


WIDTH = 320
HEIGHT = 240
DURATION = 48


def write_rgb_rotate(output):

    if not Image:
        raise SkipTest()

    output.metadata['title'] = 'container'
    output.metadata['key'] = 'value'

    stream = output.add_stream("mpeg4", 24)
    stream.width = WIDTH
    stream.height = HEIGHT
    stream.pix_fmt = "yuv420p"

    for frame_i in range(DURATION):

        frame = VideoFrame(WIDTH, HEIGHT, 'rgb24')
        image = Image.new('RGB', (WIDTH, HEIGHT), (
            int(255 * (0.5 + 0.5 * math.sin(frame_i / DURATION * 2 * math.pi))),
            int(255 * (0.5 + 0.5 * math.sin(frame_i / DURATION * 2 * math.pi + 2 / 3 * math.pi))),
            int(255 * (0.5 + 0.5 * math.sin(frame_i / DURATION * 2 * math.pi + 4 / 3 * math.pi))),
        ))
        frame.planes[0].update(image.tobytes())

        for packet in stream.encode(frame):
            output.mux(packet)

    for packet in stream.encode(None):
        output.mux(packet)

    # Done!
    output.close()


def assert_rgb_rotate(self, input_):

    # Now inspect it a little.
    self.assertEqual(len(input_.streams), 1)
    self.assertEqual(input_.metadata.get('title'), 'container', input_.metadata)
    self.assertEqual(input_.metadata.get('key'), None)
    stream = input_.streams[0]
    self.assertIsInstance(stream, VideoStream)
    self.assertEqual(stream.type, 'video')
    self.assertEqual(stream.name, 'mpeg4')
    self.assertEqual(stream.average_rate, 24)  # Only because we constructed is precisely.
    self.assertEqual(stream.rate, Fraction(24, 1))
    self.assertEqual(stream.time_base * stream.duration, 2)
    self.assertEqual(stream.format.name, 'yuv420p')
    self.assertEqual(stream.format.width, WIDTH)
    self.assertEqual(stream.format.height, HEIGHT)


class TestBasicVideoEncoding(TestCase):

    def test_rgb_rotate(self):

        path = self.sandboxed('rgb_rotate.mov')
        output = av.open(path, 'w')

        write_rgb_rotate(output)
        assert_rgb_rotate(self, av.open(path))

    def test_encoding_with_pts(self):

        path = self.sandboxed('video_with_pts.mov')
        output = av.open(path, 'w')

        stream = output.add_stream('libx264', 24)
        stream.width = WIDTH
        stream.height = HEIGHT
        stream.pix_fmt = "yuv420p"

        for i in range(DURATION):
            frame = VideoFrame(WIDTH, HEIGHT, 'rgb24')
            frame.pts = i * 2000
            frame.time_base = Fraction(1, 48000)

            for packet in stream.encode(frame):
                self.assertEqual(packet.time_base, Fraction(1, 24))
                output.mux(packet)

        for packet in stream.encode(None):
            self.assertEqual(packet.time_base, Fraction(1, 24))
            output.mux(packet)

        output.close()


class TestBasicAudioEncoding(TestCase):

    def test_audio_transcode(self):

        path = self.sandboxed('audio_transcode.mov')
        output = av.open(path, 'w')
        output.metadata['title'] = 'container'
        output.metadata['key'] = 'value'

        sample_rate = 48000
        channel_layout = 'stereo'
        channels = 2
        sample_fmt = 's16'

        stream = output.add_stream('mp2', sample_rate)

        ctx = stream.codec_context
        ctx.time_base = sample_rate
        ctx.sample_rate = sample_rate
        ctx.format = sample_fmt
        ctx.layout = channel_layout
        ctx.channels = channels

        src = av.open(fate_suite('audio-reference/chorusnoise_2ch_44kHz_s16.wav'))
        for frame in src.decode(audio=0):
            frame.pts = None
            for packet in stream.encode(frame):
                output.mux(packet)

        for packet in stream.encode(None):
            output.mux(packet)

        output.close()

        container = av.open(path)
        self.assertEqual(len(container.streams), 1)
        self.assertEqual(container.metadata.get('title'), 'container', container.metadata)
        self.assertEqual(container.metadata.get('key'), None)

        stream = container.streams[0]
        self.assertIsInstance(stream, AudioStream)
        self.assertEqual(stream.codec_context.sample_rate, sample_rate)
        self.assertEqual(stream.codec_context.format.name, 's16p')
        self.assertEqual(stream.codec_context.channels, channels)


class TestEncodeStreamSemantics(TestCase):

    def test_audio_default_options(self):
        output = av.open(self.sandboxed('output.mov'), 'w')

        stream = output.add_stream('mp2')
        self.assertEqual(stream.bit_rate, 128000)
        self.assertEqual(stream.format.name, 's16')
        self.assertEqual(stream.rate, 48000)
        self.assertEqual(stream.ticks_per_frame, 1)
        self.assertEqual(stream.time_base, None)

    def test_video_default_options(self):
        output = av.open(self.sandboxed('output.mov'), 'w')

        stream = output.add_stream('mpeg4')
        self.assertEqual(stream.bit_rate, 1024000)
        self.assertEqual(stream.format.height, 480)
        self.assertEqual(stream.format.name, 'yuv420p')
        self.assertEqual(stream.format.width, 640)
        self.assertEqual(stream.height, 480)
        self.assertEqual(stream.pix_fmt, 'yuv420p')
        self.assertEqual(stream.rate, Fraction(24, 1))
        self.assertEqual(stream.ticks_per_frame, 1)
        self.assertEqual(stream.time_base, None)
        self.assertEqual(stream.width, 640)

    def test_stream_index(self):
        output = av.open(self.sandboxed('output.mov'), 'w')

        vstream = output.add_stream('mpeg4', 24)
        vstream.pix_fmt = 'yuv420p'
        vstream.width = 320
        vstream.height = 240

        astream = output.add_stream('mp2', 48000)
        astream.channels = 2
        astream.format = 's16'

        self.assertEqual(vstream.index, 0)
        self.assertEqual(astream.index, 1)

        vframe = VideoFrame(320, 240, 'yuv420p')
        vpacket = vstream.encode(vframe)[0]

        self.assertIs(vpacket.stream, vstream)
        self.assertEqual(vpacket.stream_index, 0)

        for i in range(10):
            aframe = AudioFrame('s16', 'stereo', samples=astream.frame_size)
            aframe.rate = 48000
            apackets = astream.encode(aframe)
            if apackets:
                apacket = apackets[0]
                break

        self.assertIs(apacket.stream, astream)
        self.assertEqual(apacket.stream_index, 1)