pizzicato 0.1.0

[WIP] A library for reading and writing frames to video.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

use crate::{ Image, PixelFormat };

use anyhow::Result;

use rsmpeg::avcodec::{ AVCodec, AVCodecContext };
use rsmpeg::avformat::AVFormatContextOutput;
use rsmpeg::avutil::{ AVFrame, ra };
use rsmpeg::error::RsmpegError;
use rsmpeg::ffi;
use rsmpeg::swscale::SwsContext;

const DEFAULT_BIT_RATE: i64 = 400_000;
const DEFAULT_GOP_SIZE: i32 = 10;
const DEFAULT_MAX_B_FRAMES: i32 = 1;

#[derive(Clone, Copy, Debug, Hash, PartialEq)]
pub struct Options {
	pub size: (u32, u32),
	pub frame_rate: i32,
	pub pix_fmt: PixelFormat,

	pub bit_rate: Option<i64>,
	pub gop_size: Option<i32>,
	pub max_b_frames: Option<i32>,
}

pub fn encode(
	output_ctx: &mut AVFormatContextOutput,
	frames: impl Iterator<Item = impl Image>,
	codec: ffi::AVCodecID,
	options: Options,
) -> Result<()> {
	let encode_codec = AVCodec::find_encoder(codec)
		.ok_or(anyhow!("Failed to find encoder"))?;
	
	let mut encode_ctx = AVCodecContext::new(&encode_codec);

	encode_ctx.set_width(options.size.0 as i32);
	encode_ctx.set_height(options.size.1 as i32);
	encode_ctx.set_framerate(ra(options.frame_rate, 1));
	encode_ctx.set_time_base(ra(1, options.frame_rate));
	encode_ctx.set_pix_fmt(encode_codec.pix_fmts().unwrap()[0]);
	encode_ctx.set_bit_rate(options.bit_rate.unwrap_or(DEFAULT_BIT_RATE));
	encode_ctx.set_gop_size(options.gop_size.unwrap_or(DEFAULT_GOP_SIZE));
	encode_ctx.set_max_b_frames(options.max_b_frames.unwrap_or(DEFAULT_MAX_B_FRAMES));

	encode_ctx.open(None)?;

	let mut stream = output_ctx.new_stream();
	stream.set_codecpar(encode_ctx.extract_codecpar());
	stream.set_time_base(encode_ctx.time_base);
	drop(stream);

	let mut rescale_ctx = SwsContext::get_context(
		encode_ctx.width,
		encode_ctx.height,
		options.pix_fmt.into(),
		encode_ctx.width,
		encode_ctx.height,
		encode_ctx.pix_fmt,
		ffi::SWS_FAST_BILINEAR,
	).ok_or(anyhow!("Failed to init sws"))?;

	output_ctx.write_header()?;
	
	for (i, image) in frames.enumerate() {
		let raw_frame = image.to_avframe()
			.ok_or(anyhow!("Unable to create frame"))?;

		let mut frame = AVFrame::new();
		frame.set_format(encode_ctx.pix_fmt);
		frame.set_width(encode_ctx.width);
		frame.set_height(encode_ctx.height);
		frame.alloc_buffer()?;

		rescale_ctx.scale_frame(&raw_frame, 0, raw_frame.height, &mut frame)?;

		frame.set_pts(i as i64);

		encode_video_frame(Some(frame), output_ctx, &mut encode_ctx)?;
	}

	encode_video_frame(None, output_ctx, &mut encode_ctx)?;

	output_ctx.write_trailer()?;

	Ok(())
}

fn encode_video_frame(
	frame: Option<AVFrame>,
	output_ctx: &mut AVFormatContextOutput,
	encode_ctx: &mut AVCodecContext,
) -> Result<()> {
	encode_ctx.send_frame(frame.as_ref())?;

	let streams = output_ctx.streams();
	let stream_time_base = streams.get(streams.num() - 1)
		.unwrap().time_base;

	loop {
		let mut packet = match encode_ctx.receive_packet() {
			Ok(packet) => packet,
			Err(
				RsmpegError::EncoderDrainError |
				RsmpegError::EncoderFlushedError
			) => break,
			Err(err) => bail!(err),
		};

		packet.rescale_ts(
			encode_ctx.time_base,
			stream_time_base,
		);

		output_ctx.write_frame(&mut packet)?;
	}

	Ok(())
}