Merge pull request #60 from nus-vv-streams/bc/fix-pred

Bennett's fix-pred

Cargo.toml

@@ -103,6 +103,9 @@ name = "vv"
 name = "vvplay_async"
 required-features = ["render", "dash"]
 
+[[bin]]
+name = "exporter"
+
 [features]
 default = ["render", "dash", "with-tmc2-rs-decoder"]
 with-tmc2-rs-decoder = ["dep:tmc2rs", "dep:ffmpeg-next"]

benches/abr.rs

@@ -1,5 +1,5 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
-use vivotk::abr::quetra::QuetraMultiview;
+use vivotk::abr::quetra::{Quetra, QuetraMultiview};
 use vivotk::abr::RateAdapter;
 
 fn bench_quetra_multiview(c: &mut Criterion) {
@@ -30,5 +30,33 @@ fn bench_quetra_multiview(c: &mut Criterion) {
     });
 }
 
-criterion_group!(benches, bench_quetra_multiview);
+fn bench_quetra(c: &mut Criterion) {
+    c.bench_function("quetra", |b| {
+        let abr = Quetra::new(4, 30.0);
+        b.iter(|| {
+            abr.select_quality(
+                black_box(2),
+                black_box(1728450.56),
+                black_box(&[
+                    vec![827392, 1089536, 1490944, 2646016, 4972544, 8110080],
+                    vec![335872, 368640, 368640, 532480, 786432, 729088],
+                    vec![729088, 999424, 1466368, 2596864, 4767744, 7340032],
+                    vec![794624, 1048576, 1466368, 2547712, 4685824, 7872512],
+                    vec![270336, 303104, 319488, 442368, 704512, 679936],
+                    vec![802816, 1105920, 1572864, 2842624, 5349376, 7626752],
+                ]),
+                black_box(&[
+                    -0.07874092,
+                    -0.054728724,
+                    0.9953917,
+                    0.07874092,
+                    0.054728724,
+                    -0.9953917,
+                ]),
+            );
+        })
+    });
+}
+
+criterion_group!(benches, bench_quetra_multiview, bench_quetra);
 criterion_main!(benches);

src/bin/exporter.rs

@@ -0,0 +1,173 @@
+//! This binary is used to export pngs from a video playback session.
+//!
+//! The camera trace file is a csv file that contains the camera position and orientation for each frame.
+//! The format of the csv file is (x, y, z, pitch, yaw, roll):
+//! ```csv
+//! 0,0,1.5,0,-90,0.0
+//! 0,0,1.5,0,-90,0.0
+//! ```
+//!
+//! The quality trace file is a csv file that contains the quality for each frame.
+//! The format of the csv file is qualities for each views (view_0, view_1, view_2, view_3, view_4, view_5, view_6).
+//! You can also provide only 1 quality. In this case, this number will be the quality for all views.
+//! ```csv
+//! 0,0,0,5,5,5
+//! 3
+//! ```
+//!
+//! # Usage
+//!
+//! For remote files
+//! ```bash
+//! cargo run --bin exporter --release -- --url http://localhost:3000/longdress.mpd --quality quality_trace.csv [camera_trace.csv] [output_folder]
+//! ```
+//!
+//! For local files
+//! ```bash
+//! cargo run --bin exporter --release -- --ply_folder [ply_folder] [camera_trace.csv] [output_folder]
+//! ```
+//!
+
+use clap::Parser;
+use std::fs::File;
+use std::io::BufReader;
+use std::path::Path;
+use std::path::PathBuf;
+use vivotk::codec::{decoder::Tmc2rsDecoder, Decoder};
+use vivotk::dash::fetcher::Fetcher;
+use vivotk::render::wgpu::png::{PngWriter, RenderFormat};
+use vivotk::simulation::CameraTrace;
+use vivotk::utils::read_file_to_point_cloud;
+
+#[derive(Parser)]
+struct Args {
+    /// Camera trace file
+    camera_trace: PathBuf,
+    /// Output folder for pngs
+    output_folder: PathBuf,
+    /// Quality trace file
+    #[clap(long)]
+    quality: Option<PathBuf>,
+    /// Remote url that points to MPD file
+    #[clap(long)]
+    url: Option<String>,
+    /// Folder containing original ply files
+    #[clap(long)]
+    ply_folder: Option<PathBuf>,
+    /// Total number of frames in original video
+    #[clap(long, default_value_t = 300)]
+    frame_count: usize,
+    /// number of frames in each segment. By default, it is set to 30 frames (1 second)
+    #[clap(long, default_value_t = 30)]
+    segment_size: usize,
+    /// Set screen width.
+    ///
+    /// To enable rendering at full screen, compile with `--features fullscreen` (depends on device gpu support)
+    #[clap(short, long, default_value_t = 1600)]
+    width: u32,
+    /// Set screen height.
+    ///
+    /// To enable rendering at full screen, compile with `--features fullscreen` (depends on device gpu support)
+    #[clap(short, long, default_value_t = 900)]
+    height: u32,
+}
+
+struct QualityTrace {
+    data: Vec<Vec<usize>>,
+}
+
+impl QualityTrace {
+    fn new(path: &Path) -> Self {
+        use std::io::BufRead;
+
+        let file = File::open(path).unwrap();
+        let reader = BufReader::new(file);
+        let data = reader
+            .lines()
+            .map(|line| {
+                line.unwrap()
+                    .trim()
+                    .split(',')
+                    .map(|x| x.parse::<usize>().unwrap())
+                    .collect()
+            })
+            .collect();
+        QualityTrace { data }
+    }
+}
+
+#[tokio::main(flavor = "multi_thread", worker_threads = 2)]
+async fn main() {
+    let args: Args = Args::parse();
+    let mut png_writer = PngWriter::new(
+        args.output_folder.clone().into_os_string(),
+        0.0,
+        0.0,
+        0.0,
+        cgmath::Deg(0.0).into(),
+        cgmath::Deg(0.0).into(),
+        args.width,
+        args.height,
+        "#000000",
+        RenderFormat::Png,
+    );
+    png_writer.set_background_color(wgpu::Color::WHITE);
+
+    let camera_trace = CameraTrace::new(&args.camera_trace, false);
+
+    if let Some(quality) = args.quality {
+        let quality_trace = QualityTrace::new(&quality);
+        let mut fetcher = Fetcher::new(
+            &args.url.expect("url must be provided"),
+            &args.output_folder,
+            true,
+        )
+        .await;
+        let mut frame_number = 0;
+        for quality in quality_trace.data.iter() {
+            let res = fetcher
+                .download(
+                    0,
+                    frame_number % args.frame_count as u64,
+                    quality,
+                    true,
+                    None,
+                )
+                .await
+                .unwrap();
+            let paths = res.paths.into_iter().flatten().collect::<Vec<_>>();
+            let mut decoder = Tmc2rsDecoder::new(&paths);
+            decoder.start().unwrap();
+            for i in 0..args.segment_size {
+                let cam_pos = camera_trace.next();
+                png_writer.update_camera_pos(cam_pos);
+                let pc = decoder.poll().unwrap();
+                png_writer.write_to_png(&pc, &(i + frame_number as usize).to_string());
+            }
+            frame_number += 30;
+        }
+    } else if let Some(ply_folder) = args.ply_folder {
+        let mut ply_files: Vec<PathBuf> = vec![];
+
+        let mut dir = tokio::fs::read_dir(ply_folder).await.unwrap();
+        while let Some(entry) = dir.next_entry().await.unwrap() {
+            let f = entry.path();
+            if !f.extension().map(|f| "ply".eq(f)).unwrap_or(false) {
+                continue;
+            }
+            ply_files.push(f);
+        }
+
+        ply_files.sort();
+        for frame_number in 0..64 * args.segment_size {
+            let cam_pos = camera_trace.next();
+            png_writer.update_camera_pos(cam_pos);
+            let pc =
+                read_file_to_point_cloud(ply_files.get(frame_number % args.frame_count).unwrap())
+                    .unwrap();
+            png_writer.write_to_png(&pc, &frame_number.to_string());
+        }
+    } else {
+        unreachable!("Either quality or ply_folder must be specified")
+    }
+}