use std::process::Command;
use phasm_core::codec::h264::encoder::encoder::Encoder;
fn main() {
let w: u32 = 1920;
let h: u32 = 1072;
let n_frames: usize = std::env::var("N").ok().and_then(|s| s.parse().ok()).unwrap_or(30);
let yuv_path = std::env::var("YUV").unwrap_or_else(|_| format!("/tmp/img4138_1080p_f{n_frames}.yuv"));
let frame_size = (w * h * 3 / 2) as usize;
let y_size = (w * h) as usize;
let _c_size = (w * h / 4) as usize;
let pixels = std::fs::read(&yuv_path).expect("yuv missing");
assert!(pixels.len() >= n_frames * frame_size, "yuv too short: {} vs need {}", pixels.len(), n_frames * frame_size);
let q: u8 = std::env::var("Q").ok().and_then(|s| s.parse().ok()).unwrap_or(80);
let mut enc = Encoder::new(w, h, Some(q)).unwrap();
let mut bytes = enc.encode_i_frame(&pixels[..frame_size]).unwrap();
let mut enc_recons = vec![(enc.recon.y.clone(), enc.recon.cb.clone(), enc.recon.cr.clone())];
for fi in 1..n_frames {
let src = &pixels[fi * frame_size..(fi + 1) * frame_size];
bytes.extend_from_slice(&enc.encode_p_frame(src).unwrap());
enc_recons.push((enc.recon.y.clone(), enc.recon.cb.clone(), enc.recon.cr.clone()));
}
let h264_path = std::env::var("PHASM_OUT_H264")
.unwrap_or_else(|_| "/tmp/nframe_drift.h264".to_string());
let dec_path = std::env::var("PHASM_OUT_DEC_YUV")
.unwrap_or_else(|_| "/tmp/nframe_drift_decoded.yuv".to_string());
std::fs::write(&h264_path, &bytes).unwrap();
Command::new("ffmpeg").args([
"-y", "-loglevel", "error", "-f", "h264", "-i", h264_path.as_str(),
"-f", "rawvideo", "-pix_fmt", "yuv420p", dec_path.as_str(),
]).output().unwrap();
let decoded = std::fs::read(&dec_path).unwrap();
let n_decoded = decoded.len() / frame_size;
if n_decoded != n_frames {
println!("WARN: decoded {n_decoded} frames, encoded {n_frames}");
}
let n_check = n_decoded.min(n_frames);
println!("frame | enc-vs-dec | src-vs-recon | src-vs-dec | notes");
println!("------+-------------+----------------+-------------+--------");
for fi in 0..n_check {
let (ey, _, _) = &enc_recons[fi];
let base = fi * frame_size;
let dy = &decoded[base..base + y_size];
let src_y = &pixels[base..base + y_size];
let mut ed_sqe = 0.0f64;
let mut diff_pixels = 0usize;
let mut max_diff = 0i32;
for (a, b) in ey.iter().zip(dy.iter()) {
let d = *a as i32 - *b as i32;
if d != 0 { diff_pixels += 1; }
if d.abs() > max_diff { max_diff = d.abs(); }
ed_sqe += (d as f64) * (d as f64);
}
let ed_mse = ed_sqe / y_size as f64;
let ed_psnr = if ed_mse > 0.0 { 10.0 * (255.0 * 255.0 / ed_mse).log10() } else { 99.99 };
let mut se_sqe = 0.0f64;
for i in 0..y_size { let d = src_y[i] as f64 - ey[i] as f64; se_sqe += d * d; }
let se_mse = se_sqe / y_size as f64;
let se_psnr = 10.0 * (255.0 * 255.0 / se_mse).log10();
let mut sd_sqe = 0.0f64;
for i in 0..y_size { let d = src_y[i] as f64 - dy[i] as f64; sd_sqe += d * d; }
let sd_mse = sd_sqe / y_size as f64;
let sd_psnr = 10.0 * (255.0 * 255.0 / sd_mse).log10();
let mb_w = (w / 16) as usize;
let mb_h = (h / 16) as usize;
let mut first_bad: Option<(usize, usize, usize, i32)> = None;
if diff_pixels > 0 {
'outer: for mby in 0..mb_h {
for mbx in 0..mb_w {
let mut md = 0usize;
let mut mm = 0i32;
for dy_ in 0..16 {
for dx_ in 0..16 {
let x = mbx * 16 + dx_;
let yy = mby * 16 + dy_;
let d = ey[yy * w as usize + x] as i32 - dy[yy * w as usize + x] as i32;
if d != 0 { md += 1; }
if d.abs() > mm { mm = d.abs(); }
}
}
if md > 0 {
first_bad = Some((mbx, mby, md, mm));
break 'outer;
}
}
}
}
let note = if diff_pixels == 0 {
String::from("in sync")
} else if let Some((mx, my, md, mm)) = first_bad {
format!("DESYNC {diff_pixels}/{y_size} ({:.2}%) max{max_diff} firstMB=({mx},{my}) {md}/256 px max{mm}",
diff_pixels as f64 * 100.0 / y_size as f64)
} else {
format!("DESYNC {diff_pixels}/{y_size} ({:.2}%) max{max_diff}", diff_pixels as f64 * 100.0 / y_size as f64)
};
println!("{fi:5} | {ed_psnr:8.2} dB | {se_psnr:10.2} dB | {sd_psnr:8.2} dB | {note}");
let dump_frame = std::env::var("DUMP_FRAME").ok().and_then(|s| s.parse::<usize>().ok());
if dump_frame == Some(fi) {
let override_mb = std::env::var("DUMP_MB").ok().and_then(|s| {
let mut p = s.split(',');
let x: usize = p.next()?.trim().parse().ok()?;
let y: usize = p.next()?.trim().parse().ok()?;
Some((x, y, 0usize, 0i32))
});
let target = override_mb.or(first_bad);
if let Some((mx, my, _, _)) = target {
println!("\n ENC recon at MB ({mx},{my}) frame {fi}:");
for dy_ in 0..16 {
let mut r = String::new();
for dx_ in 0..16 {
let x = mx * 16 + dx_;
let yy = my * 16 + dy_;
r.push_str(&format!("{:4}", ey[yy * w as usize + x]));
}
println!(" {r}");
}
println!("\n DEC recon at MB ({mx},{my}) frame {fi}:");
for dy_ in 0..16 {
let mut r = String::new();
for dx_ in 0..16 {
let x = mx * 16 + dx_;
let yy = my * 16 + dy_;
r.push_str(&format!("{:4}", dy[yy * w as usize + x]));
}
println!(" {r}");
}
println!("\n DIFF (ENC - DEC) at MB ({mx},{my}) frame {fi}:");
for dy_ in 0..16 {
let mut r = String::new();
for dx_ in 0..16 {
let x = mx * 16 + dx_;
let yy = my * 16 + dy_;
let d = ey[yy * w as usize + x] as i32 - dy[yy * w as usize + x] as i32;
r.push_str(&format!("{:4}", d));
}
println!(" {r}");
}
if fi >= 1 {
let ref_y = &enc_recons[fi - 1].0;
let ref_dec_base = (fi - 1) * frame_size;
let ref_dec = &decoded[ref_dec_base..ref_dec_base + y_size];
let mut ref_diff_cnt = 0;
for dy_ in 0..16 {
for dx_ in 0..16 {
let x = mx * 16 + dx_;
let yy = my * 16 + dy_;
if ref_y[yy * w as usize + x] != ref_dec[yy * w as usize + x] {
ref_diff_cnt += 1;
}
}
}
println!("\n Reference frame ({}) at same MB: {ref_diff_cnt}/256 pixels differ enc-vs-dec", fi - 1);
}
}
}
}
}