blender.rs

// TODO: use SIMD
//#![feature(portable_simd)]

// [dependencies]
// eyre = "*"
// num-traits = "*"
// rayon = "*"
// lazy_static = "*"

pub mod isqrt;

use eyre::Context;
use eyre::ContextCompat;
use eyre::Result;
//use core_simd::*;
use isqrt::IntegerSquareRoot;

use rayon::iter::ParallelIterator;
use rayon::prelude::*;
use std::collections::VecDeque;
use std::fs;
use std::fs::File;
use std::io;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Read;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::process::Child;
use std::process::ChildStderr;
use std::process::ExitStatus;
use std::str::FromStr;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::mpsc::Receiver;

use std::sync::mpsc::sync_channel;
use std::sync::mpsc::SyncSender;

use std::thread;
use std::time::SystemTime;

#[macro_use]
extern crate lazy_static;

pub enum Frame8or32 {
    Byte8(Vec<u8>),
    Int32(Vec<u32>),
}

fn list_dir_files(dir: &Path) -> io::Result<Vec<(PathBuf, SystemTime)>> {
    let mut out = Vec::new();
    let meta = fs::metadata(dir)?;
    if dir.is_dir() && meta.is_dir() {
        for entry in fs::read_dir(dir)? {
            let entry = entry?;
            let path = entry.path();
            let meta = entry.metadata()?;
            let datetime = meta.modified()?;
            if path.is_dir() {
                //list_dir_files(&path, cb)?;
            } else {
                out.push((path.into(), datetime));
            }
        }
    }else{
        Err(io::Error::from(io::ErrorKind::NotFound))?
    }
    Ok(out)
}

fn get_video_time_length(filepath: &str) -> Result<f64> {
    let (_, mut a_out, last_duration) =
        spawn_ffmpeg("ffmpeg_probe", &["-i", filepath, "-f", "null", "-"])?;

    let mut buf = Vec::with_capacity(4 * 1024);
    a_out.read_to_end(&mut buf)?;

    let dur = last_duration.value.lock().unwrap().clone();
    println!("vid len: {} == {:.2} sec", dur.raw_str, dur.seconds);

    Ok(dur.seconds)
}

fn print_yt_url_list(input_files: &Vec<PathBuf>) -> Result<()> {
    let multiplier = 1.0 / 8.0;
    let mut out = File::create("yt_urls.txt")?;
    let mut total_sec: f64 = 0.0;
    for file_path in input_files {
        //let pathbuf = &PathBuf::from_str(&file)?;
        let file_name = file_path
            .file_name()
            .map(|s| s.to_string_lossy().to_string())
            .unwrap_or("".to_owned());

        let video_id = file_name.split("[").nth(1).unwrap_or("video");
        let video_id = video_id.split("]").nth(0).unwrap_or("VID");

        let vid_name = file_name.split("[").nth(0).unwrap_or(video_id);

        let hour = (total_sec / 3600.0).floor() as i32;
        let mins = (total_sec / 60.0).floor() as i32 % 60;
        let secs = (total_sec).floor() as i32 % 60;
        writeln!(
            out,
            "{:02}:{:02}:{:02} {}\nhttps://www.youtube.com/watch?v={}",
            hour, mins, secs, vid_name, video_id
        )?;

        let seconds = get_video_time_length(&file_path.to_string_lossy())?;
        total_sec += seconds * multiplier;
    }
    out.flush()?;
    Ok(())
}

fn read_input(reuse_arx8: Receiver<Vec<u8>>, reader_tx: SyncSender<Vec<u8>>, input_name: &str) -> Result<()> {
    let dir = &PathBuf::from_str(&format!("./vid/input_{input_name}"))?;
    println!("reading dir: {:?}", dir);
    let mut flist = list_dir_files(dir)?;
    flist.sort_by(|a, b| a.1.cmp(&b.1));
    let input_files: Vec<PathBuf> = flist.iter().map(|f| f.0.clone()).collect();
    //print_yt_url_list(&input_files)?;

    for file_path in &input_files {
        //let pathbuf = &PathBuf::from_str(&file)?;
        println!("reading file: {:?}", &file_path);
        let result = read_input_file(
            &file_path.to_string_lossy().to_owned(),
            &reuse_arx8,
            &reader_tx,
        );

        println!("read_input_file: {:?}", result);
    }

    Ok(())
}

fn read_input_file(
    filepath: &str,
    reuse_arx8: &Receiver<Vec<u8>>,
    reader_tx: &SyncSender<Vec<u8>>,
) -> Result<()> {
    let (_, mut a_out, _) = spawn_ffmpeg(
        "ffmpeg_a",
        &["-i", filepath, "-f", "rawvideo", "-pix_fmt", "rgb24", "-"],
    )?;

    loop {
        let mut frame = reuse_arx8.recv().unwrap();

        let result = a_out.read_exact(&mut frame);
        if !result.is_ok() {
            frame.fill(0);
        }
        reader_tx.send(frame).unwrap();
        result?;
    }

    Ok(())
}

fn main() -> Result<()> {

    // main config

    let input_name = "podlasie";
    let img_dx = 1920;
    let img_dy = 1080;

    // how many frames to join into one
    let blend_target = 32;
    // how many frames in running average
    let second_blend_num = 4;

    let linear_gain = 4;
    // end config

    let suffix = format!(
        "sum{}_running{}_gain{}",
        blend_target, second_blend_num, linear_gain
    );

    let mut args: Vec<&str> = Vec::new();

    args.extend_from_slice(&["-f", "rawvideo", "-pixel_format", "rgb48", "-video_size"]);

    let size_str = format!("{}x{}", img_dx, img_dy);
    args.push(&size_str);

    args.extend_from_slice(&["-framerate", "24", "-i", "-"]);

    if img_dy < 1500 {
        // upscale for youtube
        args.extend_from_slice(&["-vf", "scale=iw*2:ih*2"]);
    }

    args.extend_from_slice(&[
        "-c:v",
        "libx265",
        "-crf",
        "12",
        // "-x265-params",
        // "profile=main10",
        // "-color_primaries",
        // "bt2020",
        // "-color_trc",
        // "smpte2084",
        // "-colorspace",
        // "bt2020nc",
        "-pix_fmt",
        "yuv420p10le",
        // "-b:v",
        // "28M",
        // // "-pass",
        // // "1",
        // "-pix_fmt",
        // "yuv420p10le",
        // "-color_primaries",
        // "9",
        // "-color_trc",
        // "16",
        // "-colorspace",
        // "9",
        // "-color_range",
        // "1",

        // "-maxrate",
        // "36M",
        // "-minrate",
        // "12M",
        // "-profile:v",
        // "2",
        // "-vcodec",
        // "libvpx-vp9",
    ]);

    let out_name_str = format!("out_{input_name}_{suffix}.mkv");
    args.extend_from_slice(&["-y", &out_name_str]);

    let (mut b_in, _, _) = spawn_ffmpeg("ffmpeg_b", &args)?;
    //let blend_targetf = blend_target as f32;
    
    
    let colors_mat_size = img_dx * img_dy * 3;
    let mut blend_counter = 0;
    let mut rgb_u8 = vec![0u8; colors_mat_size];
    let mut rgb_u16 = vec![0u16; colors_mat_size];

    let mut rgb_sum_u32 = vec![0u32; colors_mat_size];
    let mut rgb_sum_of_sum_u32 = vec![0u32; colors_mat_size];

    //let mut sum = vec![0u32; 1080 * 1920 * 3];

    let mut frame = 0;
    let gamma: f32 = 2.2;
    let use_integer_gamma = (gamma - 2.0).abs() < 0.0001;
    let igamma = gamma.floor();
    let gamma_inv = 1.0 / gamma;

    let mut ring_buf: VecDeque<Vec<u32>> = Default::default();

    for _i in 0..second_blend_num {
        ring_buf.push_back(vec![0u32; colors_mat_size]);
    }

    let (reuse_tx8, reuse_rx8) = sync_channel(1024);

    for _ in 0..((blend_target / 2) + 1) {
        reuse_tx8.send(vec![0u8; colors_mat_size]).unwrap();
    }

    let reader_tx: SyncSender<Vec<u8>>;
    let reader_rx;
    (reader_tx, reader_rx) = sync_channel(128);
    println!("buffer size={}", colors_mat_size);

    let use16bit = true;

    let input_thread = thread::spawn(move || -> Result<()> {
        let result = read_input(reuse_rx8, reader_tx, input_name);
        if let Err(err) = result {
            println!("read_input err: {}", err);
            return Err(err)
        }
        result
    });

    loop {
        if false {
            // let mut frames = Vec::new();
            // for _ in 0..blend_target {
            //     frames.push(reader_rx.recv().unwrap());
            // }

            // let out_u32 = frames.into_par_iter().reduce(
            //     || arx32.lock().unwrap().recv().unwrap(),
            //     |a, b| {
            //         add_frames(
            //             tx8.clone(),
            //             tx32.clone(),
            //             arx32.clone(),
            //             a,
            //             b,
            //             colors_mat_size,
            //         )
            //     },
            // );
            // if let Frame8or32::Int32(ref out_u32) = out_u32 {
            //     rgb_sum_u32.clone_from_slice(&out_u32);
            // }
            // tx32.send(out_u32).unwrap();

            // blend_counter += blend_target;
        } else {
            //a_out.read_exact(&mut rgb_u8)?;

            let frame_rgb_u8 = reader_rx.recv();
            if frame_rgb_u8.is_err() {
                break;
            }
            let frame_rgb_u8 = frame_rgb_u8.unwrap();
            // Stage 1: rgb^2
            if use_integer_gamma {
                add_rgb_gamma2_int(&frame_rgb_u8, &mut rgb_sum_u32);
            } else {
                // rgb_u8
                //     .par_iter()
                //     .zip(rgb_sum_u32.par_iter_mut())
                //     .for_each(|(a, b)| {
                //         let a = *a as f32;
                //         *b += a.powf(2.2) as u32;
                //     });
                add_rgb_gamma2_2_table(&frame_rgb_u8, &mut rgb_sum_u32);
            }
            let _ = reuse_tx8.send(frame_rgb_u8);
            blend_counter += 1;
        }

        if blend_counter >= blend_target {
            println!("frame={} blend_counter={}", frame, blend_counter);

            let mut ring_u32 = ring_buf.pop_front().unwrap();
            rgb_sum_of_sum_u32
                .par_iter_mut()
                .zip(ring_u32.par_iter())
                .for_each(|(r, &s)| {
                    *r -= s;
                });

            ring_u32.clone_from_slice(&rgb_sum_u32);
            ring_buf.push_back(ring_u32);

            rgb_sum_of_sum_u32
                .par_iter_mut()
                .zip(rgb_sum_u32.par_iter())
                .for_each(|(r, &s)| {
                    *r += s;
                });

            let out_now_u32 = &rgb_sum_of_sum_u32;

            let pixel_div: u32 = second_blend_num * blend_target / linear_gain;
            let pixel_divf = pixel_div as f32;

            let u8_to_u16: u32 = 256;
            let u8_to_u16f: f32 = u8_to_u16 as f32;
            //let pixel_mulf16 = u8_to_u16f.powf(gamma);
            //let pixel_mulu16 = pixel_mulf16.floor() as u32;

            // Stage 2: sqrt(rgb)
            if use_integer_gamma {
                if use16bit {
                    rgb_u16
                        .par_iter_mut()
                        .zip(out_now_u32.par_iter())
                        .for_each(|(a, &b)| {
                            let b = b * u8_to_u16;
                            let b = b / pixel_div;
                            let b = b * u8_to_u16;

                            // if igamma == 1 {
                            //     *a = b as u8;
                            // } else {
                            *a = b.integer_sqrt() as u16;
                            // }
                        });
                } else {
                    rgb_u8
                        .par_iter_mut()
                        .zip(out_now_u32.par_iter())
                        .for_each(|(a, &b)| {
                            let b = b / pixel_div;
                            // if igamma == 1 {
                            //     *a = b as u8;
                            // } else {
                            *a = b.integer_sqrt() as u8;
                            // }
                        });
                }
            } else {
                if use16bit {
                    let mul = (1.0 / pixel_divf) * u8_to_u16f.powf(gamma);
                    rgb_u16
                        .par_iter_mut()
                        .zip(out_now_u32.par_iter())
                        .for_each(|(a, &b)| {
                            let b = b as f32 * mul;
                            *a = b.powf(gamma_inv) as u16;
                        });
                } else {
                    rgb_u8
                        .par_iter_mut()
                        .zip(out_now_u32.par_iter())
                        .for_each(|(a, &b)| {
                            let b = b as f32 / pixel_divf;
                            *a = b.powf(gamma_inv) as u8;
                        });
                }
            }
            if use16bit {
                let u16_as_u8s = unsafe {
                    let (prefix, ints, suffix) = rgb_u16.align_to::<u8>();
                    assert!(prefix.len() == 0);
                    assert!(suffix.len() == 0);
                    ints
                };
                b_in.write_all(&u16_as_u8s[..])?;
            } else {
                b_in.write_all(&rgb_u8[..])?;
            }

            blend_counter = 0;

            // Reset vec
            {
                let len = rgb_sum_u32.len();
                rgb_sum_u32.clear();
                rgb_sum_u32.resize(len, 0);
            }
        }

        frame += 1;
    }
    input_thread.join().unwrap()?;
    Ok(())
    // });

    // let res = t1.join();
    // res.unwrap()?;
}

#[inline(never)]
fn add_rgb_gamma2_int(aa: &[u8], bb: &mut [u32]) {
    for (a, b) in aa.iter().zip(bb.iter_mut()) {
        let a = *a as u32;
        *b += a * a;
    }
}

#[inline(never)]
#[allow(unused)]
fn add_rgb_gamma2_int_parallel(aa: &[u8], bb: &mut [u32]) {
    aa.par_iter().zip(bb.par_iter_mut()).for_each(|(a, b)| {
        let a = *a as u32;
        *b += a * a;
    });
}

#[inline(never)]
#[allow(unused)]
fn add_rgb_gamma2_2(aa: &[u8], bb: &mut [u32]) {
    for (a, b) in aa.iter().zip(bb.iter_mut()) {
        let a = *a as f32;
        *b += (a * a) as u32;
    }
}

#[inline(never)]
fn add_rgb_gamma2_2_table(aa: &[u8], bb: &mut [u32]) {
    for (&a, b) in aa.iter().zip(bb.iter_mut()) {
        let powered_22 = unsafe { GAMMA_22_TABLE.get_unchecked(a as usize) };
        *b += powered_22;
    }
}

lazy_static! {
    static ref GAMMA_22_TABLE: Vec<u32> = {
        let mut v = Vec::with_capacity(256);

        for i in 0..256 {
            let ii = i as f64;
            let powered = ii.powf(2.2);

            v.push(powered.floor() as u32);
        }

        v
    };
}

struct LastDuration {
    pub value: Mutex<LastDurationValue>,
}

impl LastDuration {
    pub fn new() -> Self {
        Self {
            value: Mutex::new(LastDurationValue::default()),
        }
    }
}

#[derive(Default, Clone)]
struct LastDurationValue {
    raw_str: String,
    seconds: f64,
}

fn spawn_ffmpeg(
    name: &'static str,
    args: &[&str],
) -> Result<(
    std::process::ChildStdin,
    std::process::ChildStdout,
    Arc<LastDuration>,
)> {
    use std::process::Command;
    use std::process::Stdio;

    let mut cmd = Box::new(Command::new("ffmpeg"));

    // ffmpeg -i - -c:v copy -c:a copy -f mpegts -
    // cmd.arg("-i")
    //     .arg("-")
    //     .arg("-c")
    //     .arg("copy")
    //     .arg("-f")
    //     .arg("mpegts")
    //     .arg("-");

    cmd.args(args);

    let write_stderr = true;

    cmd.stderr(Stdio::piped())
        .stdout(Stdio::piped())
        .stdin(Stdio::piped());

    let mut child = cmd.spawn().wrap_err("failed to spawn ffmpeg")?;

    let stdin = child
        .stdin
        .take()
        .wrap_err("child ffmpeg did not have a handle to stdin")?;

    let stdout = child
        .stdout
        .take()
        .wrap_err("child ffmpeg did not have a handle to stdout")?;

    let stderr = child
        .stderr
        .take()
        .wrap_err("child ffmpeg did not have a handle to stderr")?;

    let last_duration = Arc::new(LastDuration::new());

    let ld = last_duration.clone();
    std::thread::spawn(move || {
        let result = thread_watcher_stderr(child, stderr, write_stderr, name, ld);
        match result {
            Err(err) => {
                println!("child ffmpeg wait result was: {}\n", err);
            }
            Ok(status) => {
                println!("child ffmpeg status was: {}\n", status);
            }
        }
    });

    Ok((stdin, stdout, last_duration))
}

fn thread_watcher_stderr(
    mut child: Child,
    stderr: ChildStderr,
    write_stderr: bool,
    name: &'static str,
    last_duration: Arc<LastDuration>,
) -> Result<ExitStatus> {
    let mut stderr_file = if write_stderr {
        Some(File::create(format!("ffmpeg_{}_stderr.txt", name)).unwrap())
    } else {
        None
    };

    let br = BufReader::new(stderr);

    for line in br.lines() {
        let line = line?;

        let words = line.split(" ");

        for word in words {
            if let Some(time_str) = word.strip_prefix("time=") {
                // time=00:02:07.78
                let mut parts = time_str.split(":");
                let hour: f64 = parts
                    .next()
                    .map(|v| v.parse().ok())
                    .flatten()
                    .unwrap_or(0.0);
                let mins: f64 = parts
                    .next()
                    .map(|v| v.parse().ok())
                    .flatten()
                    .unwrap_or(0.0);
                let secs: f64 = parts
                    .next()
                    .map(|v| v.parse().ok())
                    .flatten()
                    .unwrap_or(0.0);

                let seconds = secs + mins * 60.0 + hour * 3600.0;

                *last_duration.value.lock().unwrap() = LastDurationValue {
                    seconds,
                    raw_str: time_str.to_owned(),
                };
            }
        }

        if let Some(out_file) = &mut stderr_file {
            writeln!(out_file, "{}", line)?;
        }
    }

    let status = child
        .wait()
        .wrap_err("child ffmpeg process encountered an error")?;

    Ok(status)
}