Add --output_ivf flag to remux VP9 and AV1 to ivf

Move check to startup_check

av1an-core/src/concat.rs

@@ -1,70 +1,123 @@
-use av_format::buffer::AccReader;
 use av_format::demuxer::Context as DemuxerContext;
 use av_format::demuxer::Event;
 use av_format::muxer::Context as MuxerContext;
+use av_format::{buffer::AccReader, muxer::Muxer};
 use av_ivf::demuxer::*;
 use av_ivf::muxer::*;
-use std::fs::File;
-use std::path::Path;
+use std::fs::{self, DirEntry, File};
+use std::io;
+use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
 use anyhow::Context;
 
-// pub fn concatenate_ivf(out: &Path) -> anyhow::Result<()> {
-//     let files = [];
+pub fn concat_ivf(input: &Path, out: &Path) -> anyhow::Result<()> {
+  let mut files: Vec<PathBuf> = std::fs::read_dir(input)?
+    .into_iter()
+    .filter_map(|d| d.ok())
+    .filter_map(|d| {
+      if let Ok(file_type) = d.file_type() {
+        if file_type.is_file() {
+          Some(d.path())
+        } else {
+          None
+        }
+      } else {
+        None
+      }
+    })
+    .collect();
 
-//     let output = File::create(out)?;
+  assert!(!files.is_empty());
 
-//     let mux = Box::new(IvfMuxer::new());
-//     let mut muxer = MuxerContext::new(mux, Box::new(output));
-//     // muxer.set_global_info(demuxer.info.clone())?;
-//     muxer.configure()?;
-//     muxer.write_header()?;
+  files.sort_unstable_by_key(|x| {
+    // If the temp directory follows the expected format of 00000.ivf, 00001.ivf, etc.,
+    // then these unwraps will not fail
+    x.file_stem()
+      .unwrap()
+      .to_str()
+      .unwrap()
+      .parse::<u32>()
+      .unwrap()
+  });
 
-//     let mut pos_offset: usize = 0;
-//     for file in files.iter() {
-//         let mut last_pos: usize = 0;
-//         let input = std::fs::File::open(file)
-//             .with_context(|| format!("Input file {:?} does not exist.", file))?;
+  let output = File::create(out)?;
 
-//         let acc = AccReader::new(input);
+  let mut muxer = MuxerContext::new(Box::new(IvfMuxer::new()), Box::new(output));
 
-//         let mut demuxer = DemuxerContext::new(Box::new(IvfDemuxer::new()), Box::new(acc));
-//         demuxer.read_headers()?;
+  let global_info = {
+    let acc = AccReader::new(std::fs::File::open(&files[0]).unwrap());
+    let mut demuxer = DemuxerContext::new(Box::new(IvfDemuxer::new()), Box::new(acc));
 
-//         trace!("global info: {:#?}", demuxer.info);
+    demuxer.read_headers().unwrap();
 
-//         loop {
-//             match demuxer.read_event() {
-//                 Ok(event) => match event {
-//                     Event::MoreDataNeeded(sz) => panic!("we needed more data: {} bytes", sz),
-//                     Event::NewStream(s) => panic!("new stream: {:?}", s),
-//                     Event::NewPacket(mut packet) => {
-//                         if let Some(p) = packet.pos.as_mut() {
-//                             last_pos = *p;
-//                             *p += pos_offset;
-//                         }
+    // attempt to set the duration correctly
+    let duration = demuxer.info.duration.unwrap_or(0)
+      + files
+        .iter()
+        .skip(1)
+        .filter_map(|file| {
+          let acc = AccReader::new(std::fs::File::open(file).unwrap());
+          let mut demuxer = DemuxerContext::new(Box::new(IvfDemuxer::new()), Box::new(acc));
 
-//                         debug!("received packet with pos: {:?}", packet.pos);
-//                         muxer.write_packet(Arc::new(packet))?;
-//                     }
-//                     Event::Continue => continue,
-//                     Event::Eof => {
-//                         debug!("EOF received.");
-//                         break;
-//                     }
-//                     _ => unimplemented!(),
-//                 },
-//                 Err(e) => {
-//                     debug!("error: {:?}", e);
-//                     break;
-//                 }
-//             }
-//         }
-//         pos_offset += last_pos + 1;
-//     }
+          demuxer.read_headers().unwrap();
+          demuxer.info.duration
+        })
+        .sum::<u64>();
 
-//     muxer.write_trailer()?;
+    let mut info = demuxer.info;
+    info.duration = Some(duration);
+    info
+  };
 
-//     Ok(())
-// }
+  muxer.set_global_info(global_info)?;
+
+  muxer.configure()?;
+  muxer.write_header()?;
+
+  let mut pos_offset: usize = 0;
+  for file in files.iter() {
+    let mut last_pos: usize = 0;
+    let input = std::fs::File::open(file)?;
+
+    let acc = AccReader::new(input);
+
+    let mut demuxer = DemuxerContext::new(Box::new(IvfDemuxer::new()), Box::new(acc));
+    demuxer.read_headers()?;
+
+    trace!("global info: {:#?}", demuxer.info);
+
+    loop {
+      match demuxer.read_event() {
+        Ok(event) => match event {
+          Event::MoreDataNeeded(sz) => panic!("needed more data: {} bytes", sz),
+          Event::NewStream(s) => panic!("new stream: {:?}", s),
+          Event::NewPacket(mut packet) => {
+            if let Some(p) = packet.pos.as_mut() {
+              last_pos = *p;
+              *p += pos_offset;
+            }
+
+            debug!("received packet with pos: {:?}", packet.pos);
+            muxer.write_packet(Arc::new(packet))?;
+          }
+          Event::Continue => continue,
+          Event::Eof => {
+            debug!("EOF received.");
+            break;
+          }
+          _ => unimplemented!(),
+        },
+        Err(e) => {
+          debug!("error: {:?}", e);
+          break;
+        }
+      }
+    }
+    pos_offset += last_pos + 1;
+  }
+
+  muxer.write_trailer()?;
+
+  Ok(())
+}

av1an-core/src/ffmpeg.rs

@@ -32,8 +32,7 @@ pub fn ffmpeg_get_frame_count(source: &Path) -> usize {
 
   let cap = re.captures(&output).unwrap();
 
-  let frame_count = cap[cap.len() - 1].parse::<usize>().unwrap();
-  frame_count
+  cap[cap.len() - 1].parse::<usize>().unwrap()
 }
 
 /// Returns vec of all keyframes

av1an-core/src/lib.rs

@@ -13,7 +13,7 @@ use std::str::FromStr;
 use std::{fs::File, io::Write};
 use sysinfo::SystemExt;
 
-mod concat;
+pub mod concat;
 pub mod ffmpeg;
 pub mod vapoursynth;
 

av1an-core/src/target_quality.rs

@@ -0,0 +1 @@
+

av1an-pyo3/src/lib.rs

@@ -37,7 +37,7 @@ fn hash_path(path: &str) -> PyResult<String> {
 fn create_vs_file(temp: &str, source: &str, chunk_method: &str) -> PyResult<String> {
   // only for python code, remove if being called by rust
   let temp = Path::new(temp);
-  let source = Path::new(source);
+  let source = Path::new(source).canonicalize()?;
   let chunk_method = ChunkMethod::from_str(chunk_method)
     // TODO implement this in the FromStr implementation itself
     .map_err(|_| pyo3::exceptions::PyTypeError::new_err("Invalid chunk method"))?;
@@ -117,6 +117,12 @@ fn ffmpeg_get_frame_count(source: &str) -> usize {
   av1an_core::ffmpeg::ffmpeg_get_frame_count(Path::new(source))
 }
 
+#[pyfunction]
+fn concatenate_ivf(input: &str, output: &str) -> PyResult<()> {
+  av1an_core::concat::concat_ivf(Path::new(input), Path::new(output))
+    .map_err(|e| pyo3::exceptions::PyTypeError::new_err(format!("{}", e)))
+}
+
 /// A Python module implemented in Rust.
 #[pymodule]
 fn av1an_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
@@ -128,6 +134,7 @@ fn av1an_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
   m.add_function(wrap_pyfunction!(frame_probe_vspipe, m)?)?;
   m.add_function(wrap_pyfunction!(ffmpeg_get_frame_count, m)?)?;
   m.add_function(wrap_pyfunction!(get_keyframes, m)?)?;
+  m.add_function(wrap_pyfunction!(concatenate_ivf, m)?)?;
 
   Ok(())
 }

av1an/arg_parse.py

@@ -109,7 +109,11 @@ class Args:
             help="Use mkvmerge instead of ffmpeg to concatenate",
             action="store_true",
         )
-
+        io_group.add_argument(
+            "--output_ivf",
+            help="Output an .ivf file (will not have audio)",
+            action="store_true",
+        )
         io_group.add_argument(
             "--quiet",
             "-q",

av1an/project/Project.py

@@ -14,6 +14,7 @@ from av1an_pyo3 import (
     hash_path,
     create_vs_file,
     determine_workers as determine_workers_rust,
+    concatenate_ivf,
 )
 
 
@@ -30,6 +31,7 @@ class Project(object):
         self.temp: Path = None
         self.output_file: Path = None
         self.mkvmerge: bool = None
+        self.output_ivf: bool = None
         self.config = None
         self.webm = None
 
@@ -111,7 +113,9 @@ class Project(object):
                 self.input
                 if self.is_vs
                 else create_vs_file(
-                    self.temp.as_posix(), self.input.as_posix(), self.chunk_method
+                    str(self.temp.resolve()),
+                    str(self.input.resolve()),
+                    self.chunk_method,
                 )
             )
             fr = frame_probe_vspipe(vs)
@@ -240,6 +244,11 @@ class Project(object):
         try:
             if self.encoder == "vvc":
                 vvc_concat(self.temp, self.output_file.with_suffix(".h266"))
+            elif self.output_ivf:
+                concatenate_ivf(
+                    str((self.temp / "encode").resolve()),
+                    str(self.output_file.with_suffix(".ivf").resolve()),
+                )
             elif self.mkvmerge:
                 concatenate_mkvmerge(self.temp, self.output_file)
             else:

av1an/startup/setup.py

@@ -90,6 +90,10 @@ def startup_check(project: Project):
 
         atexit.register(restore_term)
 
+    if project.encoder not in ["rav1e", "aom", "svt_av1", "vpx"] and project.output_ivf:
+        print(".ivf only supports VP8, VP9, and AV1")
+        sys.exit(1)
+
     if not project.chunk_method:
         project.select_best_chunking_method()