@@ -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`.
11491149pub 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
12281305pub 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