Skip to content

Commit fc3f7c0

Browse files
Merge pull request #2813 from Theophilus2003/perf/thumbnail-mjpeg
perf(desktop-gpui): encode recording thumbnails with FFmpeg's MJPEG encoder
2 parents 72598be + 94ed75a commit fc3f7c0

1 file changed

Lines changed: 285 additions & 51 deletions

File tree

‎apps/desktop-gpui/src/library.rs‎

Lines changed: 285 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1143,14 +1143,33 @@ pub fn thumbnail_rgba(path: &Path) -> Option<image::RgbaImage> {
11431143
/// rather than deferred.
11441144
///
11451145
/// Faithful in the parts that matter to the output: decode packets until the
1146-
/// *first* video frame comes out, scale to RGB24 at the source's own size
1147-
/// (`size: None` at both call sites, so a native-resolution JPEG, not a
1148-
/// thumbnail), save as JPEG. Blocking; callers hand it to `spawn_blocking`.
1146+
/// *first* video frame comes out and save it as a JPEG at the source's own
1147+
/// size (`size: None` at every call site, so a native-resolution JPEG, not a
1148+
/// thumbnail). Blocking; callers hand it to `spawn_blocking`.
11491149
pub fn create_screenshot(
11501150
input: &Path,
11511151
output: &Path,
11521152
size: Option<(u32, u32)>,
11531153
) -> Result<(), String> {
1154+
let frame = first_video_frame(input)?;
1155+
let size = size.unwrap_or((frame.width(), frame.height()));
1156+
let jpeg = match encode_mjpeg(&frame, size) {
1157+
Some(jpeg) => jpeg,
1158+
None => encode_rgb_jpeg(&frame, size)?,
1159+
};
1160+
if let Some(parent) = output.parent() {
1161+
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
1162+
}
1163+
let staged = output.with_extension(format!("jpg.tmp-{}", crate::store::new_uuid_v4()));
1164+
std::fs::write(&staged, jpeg)
1165+
.and_then(|()| std::fs::rename(&staged, output))
1166+
.map_err(|e| {
1167+
let _ = std::fs::remove_file(&staged);
1168+
e.to_string()
1169+
})
1170+
}
1171+
1172+
fn first_video_frame(input: &Path) -> Result<ffmpeg::frame::Video, String> {
11541173
let mut ictx = ffmpeg::format::input(input).map_err(|e| e.to_string())?;
11551174
let input_stream = ictx
11561175
.streams()
@@ -1164,65 +1183,123 @@ pub fn create_screenshot(
11641183
.video()
11651184
.map_err(|e| e.to_string())?;
11661185

1167-
let mut scaler = ffmpeg::software::scaling::context::Context::get(
1168-
decoder.format(),
1169-
decoder.width(),
1170-
decoder.height(),
1171-
ffmpeg::format::Pixel::RGB24,
1172-
size.map_or(decoder.width(), |s| s.0),
1173-
size.map_or(decoder.height(), |s| s.1),
1174-
ffmpeg::software::scaling::flag::Flags::BILINEAR,
1175-
)
1176-
.map_err(|e| e.to_string())?;
1177-
11781186
let mut frame = ffmpeg::frame::Video::empty();
11791187
for (stream, packet) in ictx.packets() {
11801188
if stream.index() != video_stream_index {
11811189
continue;
11821190
}
11831191
decoder.send_packet(&packet).map_err(|e| e.to_string())?;
1184-
if decoder.receive_frame(&mut frame).is_err() {
1185-
continue;
1192+
if decoder.receive_frame(&mut frame).is_ok() {
1193+
return Ok(frame);
11861194
}
1195+
}
11871196

1188-
let mut rgb_frame = ffmpeg::frame::Video::empty();
1189-
scaler
1190-
.run(&frame, &mut rgb_frame)
1191-
.map_err(|e| e.to_string())?;
1192-
1193-
// The scaler's rows are padded to its own stride; the image buffer is
1194-
// tight, so the copy is row by row.
1195-
let width = rgb_frame.width() as usize;
1196-
let height = rgb_frame.height() as usize;
1197-
let src_stride = rgb_frame.stride(0);
1198-
let dst_stride = width * 3;
1199-
let buffer = cap_gpui_kernels::frame::copy_rows(
1200-
rgb_frame.data(0),
1201-
dst_stride,
1202-
src_stride,
1203-
height,
1204-
false,
1205-
)
1206-
.ok_or("Failed to create image from frame data")?;
1197+
Err("Failed to create screenshot".to_string())
1198+
}
12071199

1208-
let image = image::RgbImage::from_raw(width as u32, height as u32, buffer)
1209-
.ok_or("Failed to create image from frame data")?;
1210-
if let Some(parent) = output.parent() {
1211-
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
1212-
}
1213-
let staged = output.with_extension(format!("jpg.tmp-{}", crate::store::new_uuid_v4()));
1214-
cap_gpui_kernels::codec::save_rgb_image(&image, &staged, image::ImageFormat::Jpeg)
1215-
.map_err(|e| {
1216-
let _ = std::fs::remove_file(&staged);
1217-
e.to_string()
1218-
})?;
1219-
return std::fs::rename(&staged, output).map_err(|e| {
1220-
let _ = std::fs::remove_file(&staged);
1221-
e.to_string()
1222-
});
1200+
/// FFmpeg's MJPEG encoder takes the decoded YUV frame directly, so a 4K
1201+
/// thumbnail skips the RGB round trip and `image`'s much slower JPEG encoder.
1202+
/// 4:4:4 at qscale 4 keeps the unsubsampled layout of the quality-75 JPEG
1203+
/// `image` wrote, with at least its quality.
1204+
fn encode_mjpeg(frame: &ffmpeg::frame::Video, (width, height): (u32, u32)) -> Option<Vec<u8>> {
1205+
const QSCALE: i32 = 4;
1206+
let format = ffmpeg::format::Pixel::YUV444P;
1207+
let codec = ffmpeg::codec::encoder::find(ffmpeg::codec::Id::MJPEG)?;
1208+
let mut scaler = ffmpeg::software::scaling::context::Context::get(
1209+
frame.format(),
1210+
frame.width(),
1211+
frame.height(),
1212+
format,
1213+
width,
1214+
height,
1215+
ffmpeg::software::scaling::flag::Flags::BILINEAR,
1216+
)
1217+
.ok()?;
1218+
write_full_range(&mut scaler)?;
1219+
let mut yuv = ffmpeg::frame::Video::empty();
1220+
scaler.run(frame, &mut yuv).ok()?;
1221+
1222+
let mut encoder = ffmpeg::codec::context::Context::new_with_codec(codec)
1223+
.encoder()
1224+
.video()
1225+
.ok()?;
1226+
encoder.set_width(width);
1227+
encoder.set_height(height);
1228+
encoder.set_format(format);
1229+
encoder.set_color_range(ffmpeg::color::Range::JPEG);
1230+
encoder.set_time_base((1, 1));
1231+
encoder.set_flags(ffmpeg::codec::Flags::QSCALE);
1232+
encoder.set_global_quality(QSCALE * ffmpeg::ffi::FF_QP2LAMBDA);
1233+
encoder.set_qmin(QSCALE);
1234+
encoder.set_qmax(QSCALE);
1235+
let mut encoder = encoder.open().ok()?;
1236+
1237+
yuv.set_pts(Some(0));
1238+
encoder.send_frame(&yuv).ok()?;
1239+
encoder.send_eof().ok()?;
1240+
let mut packet = ffmpeg::Packet::empty();
1241+
encoder.receive_packet(&mut packet).ok()?;
1242+
packet.data().map(<[u8]>::to_vec)
1243+
}
1244+
1245+
/// JPEG stores full-range YUV, while swscale writes limited range unless told
1246+
/// otherwise; the source keeps whatever range swscale detected for it.
1247+
fn write_full_range(scaler: &mut ffmpeg::software::scaling::context::Context) -> Option<()> {
1248+
let mut inv_table = std::ptr::null_mut();
1249+
let mut table = std::ptr::null_mut();
1250+
let (mut src_range, mut dst_range) = (0, 0);
1251+
let (mut brightness, mut contrast, mut saturation) = (0, 0, 0);
1252+
unsafe {
1253+
let context = scaler.as_mut_ptr();
1254+
let read = ffmpeg::ffi::sws_getColorspaceDetails(
1255+
context,
1256+
&mut inv_table,
1257+
&mut src_range,
1258+
&mut table,
1259+
&mut dst_range,
1260+
&mut brightness,
1261+
&mut contrast,
1262+
&mut saturation,
1263+
);
1264+
(read >= 0
1265+
&& ffmpeg::ffi::sws_setColorspaceDetails(
1266+
context, inv_table, src_range, table, 1, brightness, contrast, saturation,
1267+
) >= 0)
1268+
.then_some(())
12231269
}
1270+
}
12241271

1225-
Err("Failed to create screenshot".to_string())
1272+
/// The original path, for an FFmpeg build without the MJPEG encoder.
1273+
fn encode_rgb_jpeg(
1274+
frame: &ffmpeg::frame::Video,
1275+
(width, height): (u32, u32),
1276+
) -> Result<Vec<u8>, String> {
1277+
let mut rgb_frame = ffmpeg::frame::Video::empty();
1278+
ffmpeg::software::scaling::context::Context::get(
1279+
frame.format(),
1280+
frame.width(),
1281+
frame.height(),
1282+
ffmpeg::format::Pixel::RGB24,
1283+
width,
1284+
height,
1285+
ffmpeg::software::scaling::flag::Flags::BILINEAR,
1286+
)
1287+
.map_err(|e| e.to_string())?
1288+
.run(frame, &mut rgb_frame)
1289+
.map_err(|e| e.to_string())?;
1290+
1291+
// The scaler's rows are padded to its own stride; the image buffer is
1292+
// tight, so the copy is row by row.
1293+
let row_bytes = width as usize * 3;
1294+
let buffer = cap_gpui_kernels::frame::copy_rows(
1295+
rgb_frame.data(0),
1296+
row_bytes,
1297+
rgb_frame.stride(0),
1298+
height as usize,
1299+
false,
1300+
)
1301+
.ok_or("Failed to create image from frame data")?;
1302+
cap_gpui_kernels::codec::encode_jpeg_rgb(&buffer, width, height, 75).map_err(|e| e.to_string())
12261303
}
12271304

12281305
pub fn bundle_thumbnail_path(project_dir: &Path) -> PathBuf {
@@ -1293,6 +1370,163 @@ mod tests {
12931370
bundle
12941371
}
12951372

1373+
fn psnr(jpeg: &[u8], screen: &image::RgbImage) -> f64 {
1374+
let decoded = image::load_from_memory(jpeg).unwrap().into_rgb8();
1375+
assert_eq!(decoded.dimensions(), screen.dimensions());
1376+
let squared: f64 = decoded
1377+
.as_raw()
1378+
.iter()
1379+
.zip(screen.as_raw())
1380+
.map(|(&a, &b)| (f64::from(a) - f64::from(b)).powi(2))
1381+
.sum();
1382+
10.0 * (255.0_f64.powi(2) * screen.as_raw().len() as f64 / squared).log10()
1383+
}
1384+
1385+
/// `screen` as a recording would hand it over: limited-range BT.709 4:2:0.
1386+
fn recorded_frame(screen: &image::RgbImage) -> ffmpeg::frame::Video {
1387+
let (width, height) = screen.dimensions();
1388+
let mut rgb = ffmpeg::frame::Video::new(ffmpeg::format::Pixel::RGB24, width, height);
1389+
let stride = rgb.stride(0);
1390+
for (row, pixels) in rgb
1391+
.data_mut(0)
1392+
.chunks_mut(stride)
1393+
.zip(screen.as_raw().chunks(width as usize * 3))
1394+
{
1395+
row[..pixels.len()].copy_from_slice(pixels);
1396+
}
1397+
let mut scaler = ffmpeg::software::scaling::context::Context::get(
1398+
ffmpeg::format::Pixel::RGB24,
1399+
width,
1400+
height,
1401+
ffmpeg::format::Pixel::YUV420P,
1402+
width,
1403+
height,
1404+
ffmpeg::software::scaling::flag::Flags::BILINEAR,
1405+
)
1406+
.unwrap();
1407+
unsafe {
1408+
let bt709 = ffmpeg::ffi::sws_getCoefficients(ffmpeg::ffi::SWS_CS_ITU709);
1409+
let set = ffmpeg::ffi::sws_setColorspaceDetails(
1410+
scaler.as_mut_ptr(),
1411+
bt709,
1412+
1,
1413+
bt709,
1414+
0,
1415+
0,
1416+
1 << 16,
1417+
1 << 16,
1418+
);
1419+
assert!(set >= 0);
1420+
}
1421+
let mut frame = ffmpeg::frame::Video::empty();
1422+
scaler.run(&rgb, &mut frame).unwrap();
1423+
frame.set_color_space(ffmpeg::color::Space::BT709);
1424+
frame.set_color_range(ffmpeg::color::Range::MPEG);
1425+
frame
1426+
}
1427+
1428+
#[test]
1429+
fn screenshot_is_written_with_the_mjpeg_encoder() {
1430+
ffmpeg::init().unwrap();
1431+
let root = temp_dir("mjpeg-screenshot");
1432+
let source = root.join("display.mp4");
1433+
std::fs::write(
1434+
&source,
1435+
include_bytes!("../../media-server/src/__tests__/fixtures/test-with-audio.mp4"),
1436+
)
1437+
.unwrap();
1438+
let output = bundle_thumbnail_path(&root);
1439+
1440+
create_screenshot(&source, &output, None).unwrap();
1441+
1442+
let frame = first_video_frame(&source).unwrap();
1443+
let written = std::fs::read(&output).unwrap();
1444+
let size = (frame.width(), frame.height());
1445+
assert_eq!(written, encode_mjpeg(&frame, size).unwrap());
1446+
let decoded = image::load_from_memory_with_format(&written, image::ImageFormat::Jpeg);
1447+
assert_eq!(decoded.unwrap().to_rgb8().dimensions(), size);
1448+
std::fs::remove_dir_all(root).unwrap();
1449+
}
1450+
1451+
#[test]
1452+
fn mjpeg_screenshot_is_at_least_as_close_to_the_screen() {
1453+
ffmpeg::init().unwrap();
1454+
let screen = image::RgbImage::from_fn(640, 360, |x, y| {
1455+
let text = (x / 2 + y / 3) % 7 == 0 && y % 24 < 14;
1456+
match (x / 160, text) {
1457+
(0, _) => image::Rgb([61, 106, 242]),
1458+
(_, true) => image::Rgb([212, 212, 212]),
1459+
(3, false) => image::Rgb([(x - 480) as u8, 53, (y / 2) as u8 + 60]),
1460+
_ => image::Rgb([30, 31, 34]),
1461+
}
1462+
});
1463+
let frame = recorded_frame(&screen);
1464+
let size = screen.dimensions();
1465+
1466+
let mjpeg = psnr(&encode_mjpeg(&frame, size).unwrap(), &screen);
1467+
let rgb = psnr(&encode_rgb_jpeg(&frame, size).unwrap(), &screen);
1468+
1469+
assert!(mjpeg >= rgb, "mjpeg {mjpeg:.2} dB, rgb {rgb:.2} dB");
1470+
}
1471+
1472+
/// Each source is a video; a `.png` beside it with the same name is the
1473+
/// screen it recorded, which both JPEGs are scored against.
1474+
#[test]
1475+
#[ignore = "needs CAP_SCREENSHOT_BENCH_SOURCES"]
1476+
fn benchmark_create_screenshot() {
1477+
ffmpeg::init().unwrap();
1478+
let sources = std::env::var("CAP_SCREENSHOT_BENCH_SOURCES").unwrap();
1479+
let median = |mut runs: Vec<f64>| {
1480+
runs.sort_by(f64::total_cmp);
1481+
runs[runs.len() / 2]
1482+
};
1483+
println!(
1484+
"source,width,height,decode_ms,rgb_ms,mjpeg_ms,rgb_bytes,mjpeg_bytes,rgb_psnr,mjpeg_psnr"
1485+
);
1486+
for source in sources.split(',').map(Path::new) {
1487+
let started = std::time::Instant::now();
1488+
let frame = first_video_frame(source).unwrap();
1489+
let decode = started.elapsed().as_secs_f64() * 1e3;
1490+
let size = (frame.width(), frame.height());
1491+
let rgb = || encode_rgb_jpeg(&frame, size).unwrap();
1492+
let mjpeg = || encode_mjpeg(&frame, size).unwrap();
1493+
let time = |encode: &dyn Fn() -> Vec<u8>| {
1494+
let started = std::time::Instant::now();
1495+
std::hint::black_box(encode());
1496+
started.elapsed().as_secs_f64() * 1e3
1497+
};
1498+
let (mut rgb_runs, mut mjpeg_runs) = (Vec::new(), Vec::new());
1499+
for run in 0..15 {
1500+
if run % 2 == 0 {
1501+
rgb_runs.push(time(&rgb));
1502+
mjpeg_runs.push(time(&mjpeg));
1503+
} else {
1504+
mjpeg_runs.push(time(&mjpeg));
1505+
rgb_runs.push(time(&rgb));
1506+
}
1507+
}
1508+
let (rgb, mjpeg) = (rgb(), mjpeg());
1509+
let screen = image::open(source.with_extension("png")).map(|png| png.into_rgb8());
1510+
let score = |jpeg: &[u8]| {
1511+
screen
1512+
.as_ref()
1513+
.map_or(String::new(), |screen| format!("{:.2}", psnr(jpeg, screen)))
1514+
};
1515+
println!(
1516+
"{},{},{},{decode:.1},{:.1},{:.1},{},{},{},{}",
1517+
source.display(),
1518+
size.0,
1519+
size.1,
1520+
median(rgb_runs),
1521+
median(mjpeg_runs),
1522+
rgb.len(),
1523+
mjpeg.len(),
1524+
score(&rgb),
1525+
score(&mjpeg),
1526+
);
1527+
}
1528+
}
1529+
12961530
fn temp_dir(tag: &str) -> PathBuf {
12971531
let dir = std::env::temp_dir().join(format!(
12981532
"cap-gpui-library-{tag}-{}-{:?}",

0 commit comments

Comments
 (0)