export: H.264 color range toggle (fix dark/oversaturated output)

The zero-copy VAAPI encoder emitted full-range BT.709 NV12 but wrote no color
tags, so players assumed limited range and stretched it — the H.264 output
looked dark and oversaturated (preview and VP9/software were fine).

- Rgba2Nv12 takes `full_range` and applies the matching Y/chroma scale+offset
  (limited 16-235 / full 0-255) via a uniform; the encoder sets color_range +
  bt709 colorspace/primaries/transfer tags to match. ffprobe-confirmed.
- New ColorRange { Limited, Full } on VideoExportSettings (default Limited, the
  universally-correct choice; serde(default) so saved settings still load),
  surfaced as a "Color range" dropdown in advanced export settings for H.264.

The swscale software fallback still emits Limited regardless of the toggle
(Full only affects the VAAPI zero-copy path).
This commit is contained in:
Skyler Lehmkuhl 2026-06-26 00:36:44 -04:00
parent f1fba186c1
commit 9411145ce9
7 changed files with 122 additions and 22 deletions

View File

@ -42,15 +42,17 @@ unsafe impl Send for ZeroCopyEncoder {}
impl ZeroCopyEncoder { impl ZeroCopyEncoder {
/// Build a zero-copy `h264_vaapi` encoder writing to `output_path` (container inferred /// Build a zero-copy `h264_vaapi` encoder writing to `output_path` (container inferred
/// from the extension, e.g. `.mp4`). `Err` if VAAPI/the device is unavailable. /// from the extension, e.g. `.mp4`). `Err` if VAAPI/the device is unavailable.
#[allow(clippy::too_many_arguments)]
pub fn new( pub fn new(
width: u32, width: u32,
height: u32, height: u32,
framerate: i32, framerate: i32,
bitrate_kbps: u32, bitrate_kbps: u32,
output_path: &Path, output_path: &Path,
full_range: bool,
) -> Result<Self, String> { ) -> Result<Self, String> {
let drm = vk_device::create()?; let drm = vk_device::create()?;
let renderer = Rgba2Nv12::new(&drm.device); let renderer = Rgba2Nv12::new(&drm.device, full_range);
unsafe { unsafe {
let mut hw_device = crate::vaapi::create_device()?; let mut hw_device = crate::vaapi::create_device()?;
let name = CString::new("h264_vaapi").unwrap(); let name = CString::new("h264_vaapi").unwrap();
@ -66,6 +68,17 @@ impl ZeroCopyEncoder {
(*enc).framerate = ff::AVRational { num: framerate, den: 1 }; (*enc).framerate = ff::AVRational { num: framerate, den: 1 };
(*enc).pix_fmt = ff::AVPixelFormat::AV_PIX_FMT_VAAPI; (*enc).pix_fmt = ff::AVPixelFormat::AV_PIX_FMT_VAAPI;
(*enc).bit_rate = (bitrate_kbps as i64) * 1000; (*enc).bit_rate = (bitrate_kbps as i64) * 1000;
// Color signalling for the H.264 VUI. The Rgba2Nv12 shader produces BT.709 luma/chroma
// in the matching range; without these tags players assume limited range and a
// full-range stream looks dark + oversaturated.
(*enc).color_range = if full_range {
ff::AVColorRange::AVCOL_RANGE_JPEG
} else {
ff::AVColorRange::AVCOL_RANGE_MPEG
};
(*enc).colorspace = ff::AVColorSpace::AVCOL_SPC_BT709;
(*enc).color_primaries = ff::AVColorPrimaries::AVCOL_PRI_BT709;
(*enc).color_trc = ff::AVColorTransferCharacteristic::AVCOL_TRC_BT709;
let frames_ref = ff::av_hwframe_ctx_alloc(hw_device); let frames_ref = ff::av_hwframe_ctx_alloc(hw_device);
{ {

View File

@ -2,20 +2,41 @@
//! surface's plane textures (R8 Y, RG8 UV). Render targets (not compute storage) so it //! surface's plane textures (R8 Y, RG8 UV). Render targets (not compute storage) so it
//! works with the DMA-BUF-imported plane images, which aren't storage-writable. //! works with the DMA-BUF-imported plane images, which aren't storage-writable.
//! //!
//! BT.709 full-range, matching `nv12::cpu_reference` and the encoder's color tags. //! BT.709 matrix; the Y/chroma scale+offset are chosen by `full_range` — limited (TV, 16235)
//! is the H.264 convention most players assume, full (PC, 0255) needs the matching color tag.
/// Converts a bound RGBA texture into a Y plane (R8) and a UV plane (RG8) via two passes. /// Converts a bound RGBA texture into a Y plane (R8) and a UV plane (RG8) via two passes.
pub struct Rgba2Nv12 { pub struct Rgba2Nv12 {
y_pipeline: wgpu::RenderPipeline, y_pipeline: wgpu::RenderPipeline,
uv_pipeline: wgpu::RenderPipeline, uv_pipeline: wgpu::RenderPipeline,
bgl: wgpu::BindGroupLayout, bgl: wgpu::BindGroupLayout,
params_buf: wgpu::Buffer,
} }
impl Rgba2Nv12 { impl Rgba2Nv12 {
pub fn new(device: &wgpu::Device) -> Self { /// `full_range`: true → full/PC range (Y 0255, no scaling); false → limited/TV range
/// (Y 16235, chroma 16240), the H.264 default. The encoder must tag the stream to match.
pub fn new(device: &wgpu::Device, full_range: bool) -> Self {
// (y_offset, y_scale, chroma_offset, chroma_scale). The shader applies these to the
// BT.709 luma/chroma dot products.
let params: [f32; 4] = if full_range {
[0.0, 1.0, 128.0 / 255.0, 1.0]
} else {
[16.0 / 255.0, 219.0 / 255.0, 128.0 / 255.0, 224.0 / 255.0]
};
let params_bytes: Vec<u8> = params.iter().flat_map(|f| f.to_le_bytes()).collect();
let params_buf = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("rgba2nv12_params"),
size: params_bytes.len() as u64,
usage: wgpu::BufferUsages::UNIFORM,
mapped_at_creation: true,
});
params_buf.slice(..).get_mapped_range_mut().copy_from_slice(&params_bytes);
params_buf.unmap();
let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("rgba2nv12_bgl"), label: Some("rgba2nv12_bgl"),
entries: &[wgpu::BindGroupLayoutEntry { entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0, binding: 0,
visibility: wgpu::ShaderStages::FRAGMENT, visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Texture { ty: wgpu::BindingType::Texture {
@ -24,7 +45,18 @@ impl Rgba2Nv12 {
multisampled: false, multisampled: false,
}, },
count: None, count: None,
}], },
wgpu::BindGroupLayoutEntry {
binding: 1,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
],
}); });
let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("rgba2nv12_pl"), label: Some("rgba2nv12_pl"),
@ -65,6 +97,7 @@ impl Rgba2Nv12 {
y_pipeline: mk("y_fs", wgpu::TextureFormat::R8Unorm), y_pipeline: mk("y_fs", wgpu::TextureFormat::R8Unorm),
uv_pipeline: mk("uv_fs", wgpu::TextureFormat::Rg8Unorm), uv_pipeline: mk("uv_fs", wgpu::TextureFormat::Rg8Unorm),
bgl, bgl,
params_buf,
} }
} }
@ -80,10 +113,16 @@ impl Rgba2Nv12 {
let bg = device.create_bind_group(&wgpu::BindGroupDescriptor { let bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("rgba2nv12_bg"), label: Some("rgba2nv12_bg"),
layout: &self.bgl, layout: &self.bgl,
entries: &[wgpu::BindGroupEntry { entries: &[
wgpu::BindGroupEntry {
binding: 0, binding: 0,
resource: wgpu::BindingResource::TextureView(rgba_view), resource: wgpu::BindingResource::TextureView(rgba_view),
}], },
wgpu::BindGroupEntry {
binding: 1,
resource: self.params_buf.as_entire_binding(),
},
],
}); });
for (pipeline, view) in [(&self.y_pipeline, y_view), (&self.uv_pipeline, uv_view)] { for (pipeline, view) in [(&self.y_pipeline, y_view), (&self.uv_pipeline, uv_view)] {
let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
@ -110,6 +149,8 @@ impl Rgba2Nv12 {
const SHADER: &str = r#" const SHADER: &str = r#"
@group(0) @binding(0) var input_rgba: texture_2d<f32>; @group(0) @binding(0) var input_rgba: texture_2d<f32>;
// (y_offset, y_scale, chroma_offset, chroma_scale) — selects limited vs full range.
@group(0) @binding(1) var<uniform> params: vec4<f32>;
// Fullscreen triangle. // Fullscreen triangle.
@vertex @vertex
@ -127,7 +168,8 @@ fn load(p: vec2<i32>) -> vec3<f32> {
@fragment @fragment
fn y_fs(@builtin(position) pos: vec4<f32>) -> @location(0) vec4<f32> { fn y_fs(@builtin(position) pos: vec4<f32>) -> @location(0) vec4<f32> {
let c = load(vec2<i32>(i32(pos.x), i32(pos.y))); let c = load(vec2<i32>(i32(pos.x), i32(pos.y)));
let y = 0.2126 * c.r + 0.7152 * c.g + 0.0722 * c.b; let luma = 0.2126 * c.r + 0.7152 * c.g + 0.0722 * c.b;
let y = params.x + params.y * luma;
return vec4<f32>(y, 0.0, 0.0, 1.0); return vec4<f32>(y, 0.0, 0.0, 1.0);
} }
@ -138,8 +180,10 @@ fn uv_fs(@builtin(position) pos: vec4<f32>) -> @location(0) vec4<f32> {
let sy = 2 * i32(pos.y); let sy = 2 * i32(pos.y);
let a = (load(vec2<i32>(sx, sy)) + load(vec2<i32>(sx + 1, sy)) let a = (load(vec2<i32>(sx, sy)) + load(vec2<i32>(sx + 1, sy))
+ load(vec2<i32>(sx, sy + 1)) + load(vec2<i32>(sx + 1, sy + 1))) * 0.25; + load(vec2<i32>(sx, sy + 1)) + load(vec2<i32>(sx + 1, sy + 1))) * 0.25;
let u = -0.1146 * a.r - 0.3854 * a.g + 0.5000 * a.b + 0.5; let cb = -0.1146 * a.r - 0.3854 * a.g + 0.5000 * a.b;
let v = 0.5000 * a.r - 0.4542 * a.g - 0.0458 * a.b + 0.5; let cr = 0.5000 * a.r - 0.4542 * a.g - 0.0458 * a.b;
let u = params.z + params.w * cb;
let v = params.z + params.w * cr;
return vec4<f32>(u, v, 0.0, 1.0); return vec4<f32>(u, v, 0.0, 1.0);
} }
"#; "#;

View File

@ -54,7 +54,7 @@ fn zerocopy_real_frame_render() {
wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 }, wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 },
); );
let conv = render_nv12::Rgba2Nv12::new(&drm.device); let conv = render_nv12::Rgba2Nv12::new(&drm.device, true);
let src_view = src.create_view(&Default::default()); let src_view = src.create_view(&Default::default());
let y_view = imported.y().create_view(&Default::default()); let y_view = imported.y().create_view(&Default::default());
let uv_view = imported.uv().create_view(&Default::default()); let uv_view = imported.uv().create_view(&Default::default());

View File

@ -10,7 +10,7 @@ fn zerocopy_encode_h264() {
let (w, h) = (640u32, 480u32); let (w, h) = (640u32, 480u32);
let out = std::env::temp_dir().join("gpu_video_encoder_zerocopy.mp4"); let out = std::env::temp_dir().join("gpu_video_encoder_zerocopy.mp4");
let _ = std::fs::remove_file(&out); let _ = std::fs::remove_file(&out);
let mut enc = match ZeroCopyEncoder::new(w, h, 30, 4000, &out) { let mut enc = match ZeroCopyEncoder::new(w, h, 30, 4000, &out, false) {
Ok(e) => e, Ok(e) => e,
Err(e) => { Err(e) => {
eprintln!("[zc-encode] unavailable, skipping: {e}"); eprintln!("[zc-encode] unavailable, skipping: {e}");

View File

@ -265,6 +265,29 @@ impl VideoCodec {
} }
} }
/// YUV color range for the encoded video (currently H.264). Limited/TV (16235) is what nearly
/// every player assumes; Full/PC (0255) keeps a bit more precision but only looks right in players
/// that honor the full-range tag — so Limited is the safe default.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ColorRange {
Limited,
Full,
}
impl Default for ColorRange {
fn default() -> Self { ColorRange::Limited }
}
impl ColorRange {
pub fn is_full(&self) -> bool { matches!(self, ColorRange::Full) }
pub fn name(&self) -> &'static str {
match self {
ColorRange::Limited => "Limited (TV, 16235)",
ColorRange::Full => "Full (PC, 0255)",
}
}
}
/// Video quality presets /// Video quality presets
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum VideoQuality { pub enum VideoQuality {
@ -322,6 +345,10 @@ pub struct VideoExportSettings {
/// Video quality /// Video quality
pub quality: VideoQuality, pub quality: VideoQuality,
/// YUV color range (H.264 only; ignored by other codecs).
#[serde(default)]
pub color_range: ColorRange,
/// Audio settings (None = no audio) /// Audio settings (None = no audio)
pub audio: Option<AudioExportSettings>, pub audio: Option<AudioExportSettings>,
@ -340,6 +367,7 @@ impl Default for VideoExportSettings {
height: None, height: None,
framerate: 60.0, framerate: 60.0,
quality: VideoQuality::High, quality: VideoQuality::High,
color_range: ColorRange::Limited,
audio: Some(AudioExportSettings::high_quality_aac()), audio: Some(AudioExportSettings::high_quality_aac()),
start_time: 0.0, start_time: 0.0,
end_time: 60.0, end_time: 60.0,

View File

@ -6,7 +6,7 @@ use eframe::egui;
use lightningbeam_core::export::{ use lightningbeam_core::export::{
AudioExportSettings, AudioFormat, AudioExportSettings, AudioFormat,
ImageExportSettings, ImageFormat, ImageExportSettings, ImageFormat,
VideoExportSettings, VideoCodec, VideoQuality, VideoExportSettings, VideoCodec, VideoQuality, ColorRange,
}; };
use std::path::PathBuf; use std::path::PathBuf;
@ -527,6 +527,20 @@ impl ExportDialog {
}); });
}); });
// Color range applies to H.264 (the VAAPI zero-copy encoder honors it). Limited/TV is the
// compatible default; Full/PC only looks right in players that read the full-range tag.
if matches!(self.video_settings.codec, VideoCodec::H264) {
ui.horizontal(|ui| {
ui.label("Color range:");
egui::ComboBox::from_id_salt("video_color_range")
.selected_text(self.video_settings.color_range.name())
.show_ui(ui, |ui| {
ui.selectable_value(&mut self.video_settings.color_range, ColorRange::Limited, ColorRange::Limited.name());
ui.selectable_value(&mut self.video_settings.color_range, ColorRange::Full, ColorRange::Full.name());
});
});
}
ui.checkbox(&mut self.include_audio, "Include Audio"); ui.checkbox(&mut self.include_audio, "Include Audio");
ui.add_space(8.0); ui.add_space(8.0);

View File

@ -899,6 +899,7 @@ impl ExportOrchestrator {
framerate.round() as i32, framerate.round() as i32,
settings.quality.bitrate_kbps(), settings.quality.bitrate_kbps(),
output_path, output_path,
settings.color_range.is_full(),
) { ) {
Ok(e) => e, Ok(e) => e,
Err(e) => { Err(e) => {