mirror of
https://github.com/storytold/effectcraft.git
synced 2026-11-07 23:00:51 +00:00
JonRC
d0e7e09ecb
ProRes renders had no colr atom, so the stream reported no primaries, transfer or matrix. Write an nclc atom, as FilmCraft does for MOV, with the codes of what the render path makes: the project's output space when it has a working space, BT.709 otherwise. sRGB output stays at transfer 1, which is still open on #340 and #464. The frames don't change, so the pinned ProRes bytes move by exactly the atom's 18 bytes.
570 lines
25 KiB
Rust
570 lines
25 KiB
Rust
//! Movie export: FilmCraft H.264, HEVC (`effectcraft-hevcenc`) and AV1 (`effectcraft-av1enc`)
|
||
//! → MP4 (+ AAC), ProRes and uncompressed `raw ` → MOV (+ PCM), muxed by
|
||
//! FilmCraft's ISO BMFF / QuickTime writer.
|
||
|
||
use std::io::Write;
|
||
|
||
use effectcraft_project::render_queue::{AudioFormat, Channels, OutputFormat, ProResProfile};
|
||
use effectcraft_project::{Comp, ProjectSettings};
|
||
use effectcraft_time::{FrameRate, Tick};
|
||
use filmcraft_isobmff::{Brand, ColorInfo, FourCc, Mp4Writer, PcmConfig, SampleEntry, TrackConfig, WriteSample, WriterOptions};
|
||
use rayon::prelude::*;
|
||
|
||
use crate::{Cx, ExportError, Report, Result, State, batch_size, wants_audio};
|
||
|
||
fn enc(e: impl std::fmt::Display) -> ExportError {
|
||
ExportError::Encode(e.to_string())
|
||
}
|
||
fn mux_err(e: impl std::fmt::Display) -> ExportError {
|
||
ExportError::Io(e.to_string())
|
||
}
|
||
|
||
/// Tag a video [`SampleEntry`] with the BT.709 limited-range `colr` (nclx) box every video
|
||
/// encoder here outputs (`rgba_to_yuv420`, `rgba_to_yuv_10` and `rgba_to_yuv420_bits` all use
|
||
/// the 0.2126/0.7152/0.0722 luma and a 16–235 luma range). Without it a re-import falls back to
|
||
/// the container's resolution-based `default_matrix` — BT.601 below 1024×577 — and shifts the
|
||
/// colours (pure red came back as [232,0,1]).
|
||
pub(crate) fn tag_bt709_limited(entry: &mut SampleEntry) {
|
||
if let Some(v) = entry.video.as_mut() {
|
||
v.color = Some(ColorInfo::Nclx { primaries: 1, transfer: 1, matrix: 1, full_range: false });
|
||
}
|
||
}
|
||
|
||
pub(crate) struct Packet {
|
||
pub(crate) data: Vec<u8>,
|
||
pub(crate) key: bool,
|
||
/// pts − dts in the track timescale.
|
||
pub(crate) cto: i32,
|
||
}
|
||
|
||
/// A frame for the movie encoder: 8-bit RGBA, or 8-bit 4:2:0 planes already converted on the
|
||
/// render threads (for encoders whose [`VideoEncoder::takes_yuv420`] is true).
|
||
pub(crate) enum Pixels {
|
||
Rgba(Vec<u8>),
|
||
Yuv420(Yuv420),
|
||
}
|
||
|
||
/// BT.709 limited-range 8-bit 4:2:0 planes ([`rgba_to_yuv420`], or Y then U and V from the
|
||
/// accelerator).
|
||
pub(crate) enum Yuv420 {
|
||
Planes { y: Vec<u8>, u: Vec<u8>, v: Vec<u8> },
|
||
Packed(Vec<u8>),
|
||
}
|
||
|
||
impl Yuv420 {
|
||
fn from_rgba(rgba: &[u8], w: usize, h: usize) -> Yuv420 {
|
||
let (mut y, mut u, mut v) = (vec![], vec![], vec![]);
|
||
rgba_to_yuv420(rgba, w, h, &mut y, &mut u, &mut v);
|
||
Yuv420::Planes { y, u, v }
|
||
}
|
||
fn slices(&self, w: usize, h: usize) -> (&[u8], &[u8], &[u8]) {
|
||
match self {
|
||
Yuv420::Planes { y, u, v } => (y, u, v),
|
||
Yuv420::Packed(p) => {
|
||
let (cw, ch) = (w.div_ceil(2), h.div_ceil(2));
|
||
let (y, uv) = p.split_at(w * h);
|
||
let (u, v) = uv.split_at(cw * ch);
|
||
(y, u, v)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
pub(crate) trait VideoEncoder: Send {
|
||
fn sample_entry(&self) -> SampleEntry;
|
||
fn encode(&mut self, rgba: &[u8], index: u64) -> Result<Vec<Packet>>;
|
||
/// Whether frames should come as [`Pixels::Yuv420`] (converted while other frames render).
|
||
fn takes_yuv420(&self) -> bool {
|
||
false
|
||
}
|
||
fn encode_yuv420(&mut self, _frame: &Yuv420, _index: u64) -> Result<Vec<Packet>> {
|
||
Err(ExportError::Unsupported("this encoder takes RGBA frames".into()))
|
||
}
|
||
fn flush(&mut self) -> Result<Vec<Packet>>;
|
||
/// Edit-list media start (B-frame delay) in the track timescale.
|
||
fn media_start(&self) -> Option<i64> {
|
||
None
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------- H.264
|
||
|
||
struct H264 {
|
||
enc: filmcraft_h264enc::Encoder,
|
||
w: u32,
|
||
h: u32,
|
||
rate: FrameRate,
|
||
y: Vec<u8>,
|
||
u: Vec<u8>,
|
||
v: Vec<u8>,
|
||
}
|
||
|
||
impl H264 {
|
||
fn new(w: u32, h: u32, rate: FrameRate, kbps: u32) -> Result<H264> {
|
||
let mut cfg = filmcraft_h264enc::EncoderConfig::new(w, h, rate.num as u32, rate.den as u32);
|
||
cfg.format = filmcraft_h264enc::PacketFormat::LengthPrefixed;
|
||
cfg.aud = false;
|
||
cfg.keyint = (rate.as_f64() * 2.0).round().max(1.0) as u32;
|
||
let kbps = kbps.max(100);
|
||
cfg.rate = filmcraft_h264enc::RateControl::Vbr { target_kbps: kbps, max_kbps: kbps * 3 / 2 };
|
||
let enc = filmcraft_h264enc::Encoder::new(cfg).map_err(enc)?;
|
||
Ok(H264 { enc, w, h, rate, y: vec![], u: vec![], v: vec![] })
|
||
}
|
||
fn packets(&self, ps: Vec<filmcraft_h264enc::Packet>) -> Vec<Packet> {
|
||
ps.into_iter().map(|p| Packet { data: p.data, key: p.keyframe, cto: (p.pts - p.dts) as i32 }).collect()
|
||
}
|
||
}
|
||
|
||
impl VideoEncoder for H264 {
|
||
fn sample_entry(&self) -> SampleEntry {
|
||
let cfg = filmcraft_isobmff::AvcConfig::parse(&self.enc.avcc()).unwrap_or_else(|_| {
|
||
let (sps, pps) = self.enc.sps_pps();
|
||
filmcraft_isobmff::AvcConfig::new(vec![sps], vec![pps], 4)
|
||
});
|
||
let mut entry = SampleEntry::avc(cfg, self.w as u16, self.h as u16);
|
||
tag_bt709_limited(&mut entry);
|
||
entry
|
||
}
|
||
fn encode(&mut self, rgba: &[u8], index: u64) -> Result<Vec<Packet>> {
|
||
let (mut y, mut u, mut v) = (std::mem::take(&mut self.y), std::mem::take(&mut self.u), std::mem::take(&mut self.v));
|
||
rgba_to_yuv420(rgba, self.w as usize, self.h as usize, &mut y, &mut u, &mut v);
|
||
let planes = Yuv420::Planes { y, u, v };
|
||
let r = self.encode_yuv420(&planes, index);
|
||
if let Yuv420::Planes { y, u, v } = planes {
|
||
(self.y, self.u, self.v) = (y, u, v);
|
||
}
|
||
r
|
||
}
|
||
fn takes_yuv420(&self) -> bool {
|
||
true
|
||
}
|
||
fn encode_yuv420(&mut self, p: &Yuv420, index: u64) -> Result<Vec<Packet>> {
|
||
let (w, cw) = (self.w as usize, (self.w as usize).div_ceil(2));
|
||
let (y, u, v) = p.slices(w, self.h as usize);
|
||
let frame = filmcraft_h264enc::YuvFrame { y, u, v, y_stride: w, uv_stride: cw };
|
||
let ps = self.enc.try_encode(&frame, index as i64 * self.rate.den).map_err(enc)?;
|
||
Ok(self.packets(ps))
|
||
}
|
||
fn flush(&mut self) -> Result<Vec<Packet>> {
|
||
let ps = self.enc.flush();
|
||
Ok(self.packets(ps))
|
||
}
|
||
fn media_start(&self) -> Option<i64> {
|
||
(self.enc.delay() > 0).then_some(self.rate.den)
|
||
}
|
||
}
|
||
|
||
pub(crate) use effectcraft_raster::yuv::rgba_to_yuv420;
|
||
|
||
// ---------------------------------------------------------------- ProRes
|
||
|
||
struct ProRes {
|
||
enc: filmcraft_prores::Encoder,
|
||
profile: filmcraft_prores::Profile,
|
||
w: u32,
|
||
h: u32,
|
||
alpha: bool,
|
||
colr: ColorInfo,
|
||
}
|
||
|
||
/// The `colr` atom of a ProRes render: what its frames are, in H.273 code points. MOV uses
|
||
/// QuickTime's `nclc`, which has no range flag; ProRes frames are limited range. The matrix is the
|
||
/// one [`rgba_to_yuv_10`] always uses. Primaries and transfer follow the render path
|
||
/// (`Pipe::output`): only a working space converts to the output space, so unmanaged frames pass
|
||
/// through. sRGB output keeps transfer 1 rather than 13 (#340), and an alpha matte is not in the
|
||
/// output space.
|
||
fn prores_colr(settings: &ProjectSettings, channels: Channels) -> ColorInfo {
|
||
use effectcraft_project::ColorSpace as Space;
|
||
let output = settings.working_space.and(settings.output_space).filter(|_| channels != Channels::Alpha);
|
||
let (primaries, transfer) = match output {
|
||
None | Some(Space::Srgb | Space::Rec709) => (1, 1),
|
||
Some(Space::Rec2020) => (9, 1),
|
||
Some(Space::DisplayP3) => (12, 13),
|
||
Some(Space::Rec2100Pq) => (9, 16),
|
||
Some(Space::Rec2100Hlg) => (9, 18),
|
||
// H.273 has no code for the ACES primaries: unspecified, in linear light.
|
||
Some(Space::AcesCg | Space::Aces2065) => (2, 8),
|
||
};
|
||
ColorInfo::Nclc { primaries, transfer, matrix: 1 }
|
||
}
|
||
|
||
fn prores_profile(p: ProResProfile) -> filmcraft_prores::Profile {
|
||
use filmcraft_prores::Profile;
|
||
match p {
|
||
ProResProfile::Proxy => Profile::Proxy,
|
||
ProResProfile::Lt => Profile::Lt,
|
||
ProResProfile::Standard => Profile::Standard,
|
||
ProResProfile::Hq => Profile::Hq,
|
||
ProResProfile::P4444 => Profile::P4444,
|
||
ProResProfile::P4444Xq => Profile::P4444Xq,
|
||
}
|
||
}
|
||
|
||
impl ProRes {
|
||
fn new(w: u32, h: u32, profile: ProResProfile, channels: Channels, colr: ColorInfo) -> ProRes {
|
||
// Alpha needs a 4444 flavour.
|
||
let profile = if channels == Channels::Rgba && !profile.is_4444() { ProResProfile::P4444 } else { profile };
|
||
let profile = prores_profile(profile);
|
||
let alpha = channels == Channels::Rgba;
|
||
let mut cfg = filmcraft_prores::EncoderConfig::new(profile, w, h);
|
||
cfg.encode_alpha = alpha;
|
||
ProRes { enc: filmcraft_prores::Encoder::with_config(cfg), profile, w, h, alpha, colr }
|
||
}
|
||
fn is_444(&self) -> bool {
|
||
matches!(self.profile, filmcraft_prores::Profile::P4444 | filmcraft_prores::Profile::P4444Xq)
|
||
}
|
||
}
|
||
|
||
impl VideoEncoder for ProRes {
|
||
fn sample_entry(&self) -> SampleEntry {
|
||
let mut entry = SampleEntry::prores(FourCc(self.profile.fourcc()), self.w as u16, self.h as u16);
|
||
if let Some(v) = entry.video.as_mut() {
|
||
v.color = Some(self.colr.clone());
|
||
}
|
||
entry
|
||
}
|
||
fn encode(&mut self, rgba: &[u8], _index: u64) -> Result<Vec<Packet>> {
|
||
let chroma = if self.is_444() { filmcraft_prores::ChromaFormat::Yuv444 } else { filmcraft_prores::ChromaFormat::Yuv422 };
|
||
let mut fr = filmcraft_prores::Frame::new(self.w, self.h, chroma, 10, self.alpha);
|
||
rgba_to_yuv_10(rgba, self.w as usize, self.h as usize, self.is_444(), &mut fr.y, &mut fr.cb, &mut fr.cr);
|
||
if let Some(a) = fr.alpha.as_mut() {
|
||
for (d, s) in a.iter_mut().zip(rgba.as_chunks::<4>().0.iter()) {
|
||
*d = ((s[3] as u32 * 1023 + 127) / 255) as u16;
|
||
}
|
||
}
|
||
let data = self.enc.encode(&fr).map_err(enc)?;
|
||
Ok(vec![Packet { data, key: true, cto: 0 }])
|
||
}
|
||
fn flush(&mut self) -> Result<Vec<Packet>> {
|
||
Ok(vec![])
|
||
}
|
||
}
|
||
|
||
/// BT.709 limited-range 10-bit 4:2:2 (horizontal chroma average) or 4:4:4 from RGBA8.
|
||
fn rgba_to_yuv_10(rgba: &[u8], w: usize, h: usize, full_chroma: bool, y: &mut [u16], cb: &mut [u16], cr: &mut [u16]) {
|
||
let cw = if full_chroma { w } else { w.div_ceil(2) };
|
||
debug_assert!(y.len() >= w * h && rgba.len() >= w * h * 4);
|
||
y.par_chunks_mut(w).zip(cb.par_chunks_mut(cw).zip(cr.par_chunks_mut(cw))).enumerate().for_each(|(row, (yr, (cbr, crr)))| {
|
||
let src = &rgba[row * w * 4..(row + 1) * w * 4];
|
||
let mut us = vec![0f32; w];
|
||
let mut vs = vec![0f32; w];
|
||
for x in 0..w {
|
||
let (r, g, b) = (src[x * 4] as f32 / 255.0, src[x * 4 + 1] as f32 / 255.0, src[x * 4 + 2] as f32 / 255.0);
|
||
let l = 0.2126 * r + 0.7152 * g + 0.0722 * b;
|
||
yr[x] = (64.0 + 876.0 * l).round().clamp(4.0, 1019.0) as u16;
|
||
us[x] = (b - l) / 1.8556;
|
||
vs[x] = (r - l) / 1.5748;
|
||
}
|
||
for cx in 0..cw {
|
||
let (u, v) = if full_chroma {
|
||
(us[cx], vs[cx])
|
||
} else {
|
||
let (a, b2) = (cx * 2, (cx * 2 + 1).min(w - 1));
|
||
((us[a] + us[b2]) * 0.5, (vs[a] + vs[b2]) * 0.5)
|
||
};
|
||
cbr[cx] = (512.0 + 896.0 * u).round().clamp(4.0, 1019.0) as u16;
|
||
crr[cx] = (512.0 + 896.0 * v).round().clamp(4.0, 1019.0) as u16;
|
||
}
|
||
});
|
||
}
|
||
|
||
// ---------------------------------------------------------------- audio
|
||
|
||
enum Audio {
|
||
Aac(Box<filmcraft_aac::Encoder>),
|
||
Pcm,
|
||
}
|
||
|
||
fn deinterleave(buf: &[f32], channels: usize) -> Vec<Vec<f32>> {
|
||
(0..channels).map(|c| buf.chunks_exact(channels).map(|p| p[c]).collect()).collect()
|
||
}
|
||
|
||
/// The comp's audio over `n` samples from `start`, interleaved with the module's channel count
|
||
/// (mono = the average of left and right).
|
||
pub(crate) fn mix(cx: &Cx, start: effectcraft_time::Tick, n: usize, sr: u32) -> Vec<f32> {
|
||
let st = effectcraft_render::audio::mix_comp(&cx.project, cx.footage, cx.expr, cx.comp, start, n, sr);
|
||
if cx.output.audio_channels == 1 { st.as_chunks::<2>().0.iter().map(|p| (p[0] + p[1]) * 0.5).collect() } else { st }
|
||
}
|
||
|
||
/// PCM bytes of samples in `fmt` (little or big endian).
|
||
pub(crate) fn pcm_bytes(buf: &[f32], fmt: AudioFormat, big_endian: bool) -> Vec<u8> {
|
||
let mut out = Vec::with_capacity(buf.len() * 4);
|
||
for &s in buf {
|
||
let s = s.clamp(-1.0, 1.0);
|
||
match fmt {
|
||
AudioFormat::S16 => {
|
||
let v = (s * 32767.0).round() as i16;
|
||
out.extend_from_slice(&if big_endian { v.to_be_bytes() } else { v.to_le_bytes() });
|
||
}
|
||
AudioFormat::S24 => {
|
||
let v = (s * 8_388_607.0).round() as i32;
|
||
let b = v.to_le_bytes();
|
||
if big_endian {
|
||
out.extend_from_slice(&[b[2], b[1], b[0]]);
|
||
} else {
|
||
out.extend_from_slice(&b[..3]);
|
||
}
|
||
}
|
||
AudioFormat::F32 => out.extend_from_slice(&if big_endian { s.to_be_bytes() } else { s.to_le_bytes() }),
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
// ---------------------------------------------------------------- the movie
|
||
|
||
type Writer<'f, 's> = Mp4Writer<&'f mut crate::out::Out<'s>>;
|
||
|
||
pub(crate) async fn movie(job: &Cx<'_>, comp: &Comp, w: u32, h: u32, st: &mut State<'_>) -> Result<Report> {
|
||
let rate = job.settings.rate(comp);
|
||
let fmt = job.output.format;
|
||
let channels = match job.output.channels {
|
||
Channels::Rgba if !matches!(fmt, OutputFormat::ProRes | OutputFormat::Uncompressed) => Channels::Rgb,
|
||
c => c,
|
||
};
|
||
let mut venc: Box<dyn VideoEncoder> = match fmt {
|
||
OutputFormat::H264 => Box::new(H264::new(w, h, rate, job.output.bitrate_kbps)?),
|
||
OutputFormat::Hevc => Box::new(crate::hevc_av1::Hevc::new(w, h, rate, job.output)?),
|
||
OutputFormat::Av1 => Box::new(crate::hevc_av1::Av1::new(w, h, rate, job.output)?),
|
||
OutputFormat::Uncompressed => Box::new(crate::uncompressed::Raw::new(w, h, channels == Channels::Rgba)?),
|
||
_ => {
|
||
let colr = prores_colr(&job.project.settings, channels);
|
||
Box::new(ProRes::new(w, h, job.output.prores_profile, if channels == Channels::Alpha { Channels::Rgb } else { channels }, colr))
|
||
}
|
||
};
|
||
let brand = if matches!(fmt, OutputFormat::ProRes | OutputFormat::Uncompressed) { Brand::Mov } else { Brand::Mp4 };
|
||
let batch = batch_size();
|
||
let yuv = venc.takes_yuv420();
|
||
let render = |_: u64, img: effectcraft_raster::Image| -> Pixels {
|
||
let px = job.pixels(&img, comp, channels, w, h);
|
||
if yuv { Pixels::Yuv420(Yuv420::from_rgba(&px, w as usize, h as usize)) } else { Pixels::Rgba(px) }
|
||
};
|
||
// H.264 frames the accelerator can convert itself come back as planes (12 MB instead of
|
||
// 133 MB of RGBA f32 read back at 4K); everything else renders and converts as before.
|
||
// Deferred readbacks (a browser worker) keep their pass-based path: no blocking there.
|
||
let gpu_yuv = yuv && channels == Channels::Rgb && !job.deferred();
|
||
let batch_frames = async |ks: Vec<u64>| -> Vec<Pixels> {
|
||
if gpu_yuv { job.h264_frames(comp, ks, w, h, |p| Pixels::Yuv420(Yuv420::Packed(p)), render) } else { job.frames(comp, ks, render).await }
|
||
};
|
||
let encode = |venc: &mut Box<dyn VideoEncoder>, px: &Pixels, index: u64| match px {
|
||
Pixels::Rgba(rgba) => venc.encode(rgba, index),
|
||
Pixels::Yuv420(planes) => venc.encode_yuv420(planes, index),
|
||
};
|
||
|
||
let beside = crate::pipeline::Beside::shared();
|
||
let total = st.total;
|
||
|
||
// Encode the first batch before creating tracks: encoders finalise their config on frame 1.
|
||
// Each batch encodes while the next one renders.
|
||
let first_end = batch.min(total);
|
||
let first: Vec<Pixels> = batch_frames((0..first_end).collect()).await;
|
||
let next_end = (first_end + batch).min(total);
|
||
let (pending, mut ready) = job
|
||
.beside(beside, batch_frames((first_end..next_end).collect()), || -> Result<Vec<Packet>> {
|
||
let mut pending = Vec::new();
|
||
for (k, px) in first.iter().enumerate() {
|
||
pending.extend(encode(&mut venc, px, k as u64)?);
|
||
}
|
||
Ok(pending)
|
||
})
|
||
.await?;
|
||
let pending = pending?;
|
||
drop(first);
|
||
st.advance(first_end)?;
|
||
|
||
let path = job.place(job.path, 1);
|
||
let mut file = job.create(&path)?;
|
||
let opts = WriterOptions::new(brand);
|
||
let movie_ts = opts.movie_timescale.max(1) as i128;
|
||
let mut mux: Writer = Mp4Writer::new(&mut file, opts).map_err(mux_err)?;
|
||
let mut vcfg = TrackConfig::new(venc.sample_entry(), rate.num as u32);
|
||
if let Some(start) = venc.media_start() {
|
||
// An explicit edit: the convenience `media_start` edit would subtract the B-frame delay
|
||
// from the media duration and drop the last frame.
|
||
let frames_ts = st.total as i128 * rate.den as i128;
|
||
let seg = (frames_ts * movie_ts + rate.num as i128 / 2) / rate.num as i128;
|
||
vcfg.edits = vec![filmcraft_isobmff::Edit { segment_duration: seg as u64, media_time: start, media_rate: 0x10000 }];
|
||
}
|
||
let vt = mux.add_track(vcfg).map_err(mux_err)?;
|
||
let sr = job.output.audio_sample_rate.clamp(8_000, 192_000);
|
||
let with_audio = wants_audio(job);
|
||
let chans = if job.output.audio_channels == 1 { 1u32 } else { 2 };
|
||
let afmt = job.output.audio_format;
|
||
let mut audio = None;
|
||
if with_audio {
|
||
let (a, entry, start) = if brand == Brand::Mp4 {
|
||
let e = filmcraft_aac::Encoder::new(filmcraft_aac::EncoderConfig::cbr(sr, chans as usize, 160_000 * chans)).map_err(enc)?;
|
||
let entry = SampleEntry::aac(e.audio_specific_config(), chans, sr);
|
||
let priming = e.priming_samples() as i64;
|
||
(Audio::Aac(Box::new(e)), entry, Some(priming))
|
||
} else {
|
||
let (bits, float) = match afmt {
|
||
AudioFormat::S16 => (16, false),
|
||
AudioFormat::S24 => (24, false),
|
||
AudioFormat::F32 => (32, true),
|
||
};
|
||
let pcm = PcmConfig { bits, float, big_endian: false, signed: true, channels: chans, sample_rate: sr as f64 };
|
||
(Audio::Pcm, SampleEntry::pcm(pcm), None)
|
||
};
|
||
let mut c = TrackConfig::new(entry, sr);
|
||
c.media_start = start;
|
||
let at = mux.add_track(c).map_err(mux_err)?;
|
||
audio = Some((a, at));
|
||
}
|
||
let (span_start, span_end) = job.settings.span(comp);
|
||
let mut cursor = span_start.to_units_floor(sr as i64);
|
||
|
||
let write_video = |mux: &mut Writer, ps: Vec<Packet>| -> Result<()> {
|
||
for p in ps {
|
||
mux.write_sample(vt, WriteSample { data: &p.data, duration: rate.den as u32, composition_offset: p.cto, is_sync: p.key }).map_err(mux_err)?;
|
||
}
|
||
Ok(())
|
||
};
|
||
let mut write_audio_until = |mux: &mut Writer, until: Tick, cursor: &mut i64| -> Result<()> {
|
||
let Some((a, at)) = audio.as_mut() else { return Ok(()) };
|
||
let end = until.to_units_floor(sr as i64);
|
||
if end <= *cursor {
|
||
return Ok(());
|
||
}
|
||
let n = (end - *cursor) as usize;
|
||
let start = Tick::from_units(*cursor, i64::from(sr));
|
||
let start = if start.to_units_floor(i64::from(sr)) < *cursor { Tick(start.0.saturating_add(1)) } else { start };
|
||
let buf = mix(job, start, n, sr);
|
||
*cursor = end;
|
||
match a {
|
||
Audio::Aac(e) => {
|
||
let planar = deinterleave(&buf, chans as usize);
|
||
let refs: Vec<&[f32]> = planar.iter().map(Vec::as_slice).collect();
|
||
for au in e.encode(&refs) {
|
||
mux.write_sample(*at, WriteSample { data: &au, duration: 1024, composition_offset: 0, is_sync: true }).map_err(mux_err)?;
|
||
}
|
||
}
|
||
Audio::Pcm => {
|
||
let pcm = pcm_bytes(&buf, afmt, false);
|
||
mux.write_sample(*at, WriteSample { data: &pcm, duration: n as u32, composition_offset: 0, is_sync: true }).map_err(mux_err)?;
|
||
}
|
||
}
|
||
Ok(())
|
||
};
|
||
let frame_end_time = |k: u64| -> Tick { if k >= total { span_end } else { job.settings.frame_time(comp, k) } };
|
||
write_video(&mut mux, pending)?;
|
||
write_audio_until(&mut mux, frame_end_time(first_end), &mut cursor)?;
|
||
let mut i = first_end;
|
||
while i < total {
|
||
let end = (i + batch).min(total);
|
||
let next_end = (end + batch).min(total);
|
||
let frames = std::mem::take(&mut ready);
|
||
let (done, next) = job
|
||
.beside(beside, batch_frames((end..next_end).collect()), || -> Result<()> {
|
||
for (k, px) in frames.iter().enumerate() {
|
||
let ps = encode(&mut venc, px, i + k as u64)?;
|
||
write_video(&mut mux, ps)?;
|
||
}
|
||
write_audio_until(&mut mux, frame_end_time(end), &mut cursor)
|
||
})
|
||
.await?;
|
||
done?;
|
||
ready = next;
|
||
st.advance(end - i)?;
|
||
i = end;
|
||
}
|
||
write_video(&mut mux, venc.flush()?)?;
|
||
write_audio_until(&mut mux, span_end, &mut cursor)?;
|
||
if let Some((Audio::Aac(e), at)) = audio.as_mut() {
|
||
for au in e.flush() {
|
||
mux.write_sample(*at, WriteSample { data: &au, duration: 1024, composition_offset: 0, is_sync: true }).map_err(mux_err)?;
|
||
}
|
||
}
|
||
let out = mux.finish().map_err(mux_err)?;
|
||
out.flush().map_err(mux_err)?;
|
||
let bytes = file.finish()?;
|
||
Ok(Report { path, frames: 0, width: w, height: h, seconds: 0.0, bytes, audio: with_audio, log: None, overflow: vec![] })
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use effectcraft_raster::yuv::{UNIT8, round_pos};
|
||
|
||
use super::*;
|
||
|
||
/// The conversion as first written (per-pixel division and `f32::round`).
|
||
fn yuv420_reference(rgba: &[u8], w: usize, h: usize) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
|
||
let (cw, ch) = (w.div_ceil(2), h.div_ceil(2));
|
||
let (mut y, mut u, mut v) = (vec![0u8; w * h], vec![0u8; cw * ch], vec![0u8; cw * ch]);
|
||
for cy in 0..ch {
|
||
let mut acc = vec![(0f32, 0f32, 0f32); cw];
|
||
for row in cy * 2..(cy * 2 + 2).min(h) {
|
||
for x in 0..w {
|
||
let s = &rgba[(row * w + x) * 4..];
|
||
let (r, g, b) = (s[0] as f32 / 255.0, s[1] as f32 / 255.0, s[2] as f32 / 255.0);
|
||
let l = 0.2126 * r + 0.7152 * g + 0.0722 * b;
|
||
y[row * w + x] = (16.0 + 219.0 * l).round().clamp(1.0, 254.0) as u8;
|
||
let a = &mut acc[x / 2];
|
||
a.0 += (b - l) / 1.8556;
|
||
a.1 += (r - l) / 1.5748;
|
||
a.2 += 1.0;
|
||
}
|
||
}
|
||
for (cx, a) in acc.iter().enumerate() {
|
||
u[cy * cw + cx] = (128.0 + 224.0 * a.0 / a.2).round().clamp(1.0, 254.0) as u8;
|
||
v[cy * cw + cx] = (128.0 + 224.0 * a.1 / a.2).round().clamp(1.0, 254.0) as u8;
|
||
}
|
||
}
|
||
(y, u, v)
|
||
}
|
||
|
||
#[test]
|
||
fn round_pos_is_round() {
|
||
for i in 0..2_000_000u32 {
|
||
let v = i as f32 * 0.000_25;
|
||
assert_eq!(round_pos(v).to_bits(), v.round().to_bits(), "{v}");
|
||
let n = f32::from_bits(v.to_bits().wrapping_sub(1)).max(0.0);
|
||
assert_eq!(round_pos(n).to_bits(), n.round().to_bits(), "{n}");
|
||
}
|
||
for k in 0..1024u32 {
|
||
for v in [k as f32 + 0.5, f32::from_bits((k as f32 + 0.5).to_bits() - 1), f32::from_bits((k as f32 + 0.5).to_bits() + 1)] {
|
||
assert_eq!(round_pos(v), v.round(), "{v}");
|
||
}
|
||
}
|
||
for (i, u) in UNIT8.iter().enumerate() {
|
||
assert_eq!(u.to_bits(), (i as f32 / 255.0).to_bits());
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn yuv420_matches_the_reference() {
|
||
// Every 8-bit level in every channel, odd sizes, and pseudo-random pixels.
|
||
for (w, h) in [(1, 1), (2, 2), (3, 5), (17, 9), (64, 33)] {
|
||
let mut rgba = vec![0u8; w * h * 4];
|
||
let mut seed = (w * 31 + h) as u32;
|
||
for (i, b) in rgba.iter_mut().enumerate() {
|
||
seed = seed.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
|
||
*b = if i % 7 == 0 { (i / 4 % 256) as u8 } else { (seed >> 24) as u8 };
|
||
}
|
||
let (mut y, mut u, mut v) = (vec![], vec![], vec![]);
|
||
rgba_to_yuv420(&rgba, w, h, &mut y, &mut u, &mut v);
|
||
assert_eq!((y, u, v), yuv420_reference(&rgba, w, h), "{w}x{h}");
|
||
}
|
||
let grey: Vec<u8> = (0..=255u8).flat_map(|g| [g, g, g, 255, g, 255 - g, g / 2, 255]).collect();
|
||
let (mut y, mut u, mut v) = (vec![], vec![], vec![]);
|
||
rgba_to_yuv420(&grey, 512, 1, &mut y, &mut u, &mut v);
|
||
assert_eq!((y, u, v), yuv420_reference(&grey, 512, 1));
|
||
}
|
||
|
||
#[test]
|
||
fn yuv_of_primaries() {
|
||
let rgba = [255u8, 255, 255, 255, 0, 0, 0, 255, 255, 0, 0, 255, 0, 255, 0, 255];
|
||
let (mut y, mut u, mut v) = (vec![], vec![], vec![]);
|
||
rgba_to_yuv420(&rgba, 2, 2, &mut y, &mut u, &mut v);
|
||
assert_eq!(y[0], 235);
|
||
assert_eq!(y[1], 16);
|
||
assert_eq!((u.len(), v.len()), (1, 1));
|
||
let (mut y10, mut cb, mut cr) = (vec![0u16; 4], vec![0u16; 4], vec![0u16; 4]);
|
||
rgba_to_yuv_10(&rgba, 2, 2, true, &mut y10, &mut cb, &mut cr);
|
||
assert_eq!(y10[0], 940);
|
||
assert_eq!(y10[1], 64);
|
||
assert_eq!(cb[0], 512);
|
||
assert!(cr[2] > 900, "red has high Cr: {}", cr[2]);
|
||
}
|
||
}
|