twgpu_tools/
encoding.rs

1use std::path::PathBuf;
2use std::sync::mpsc;
3use std::thread;
4
5use clap::Subcommand;
6use rayon::iter::{ParallelBridge, ParallelIterator};
7
8#[cfg(feature = "ffmpeg")]
9use crate::ff;
10use crate::DownloadTexture;
11
12pub enum EncodingKind {
13    Images,
14    #[cfg(feature = "ffmpeg")]
15    Videos,
16}
17
18pub enum Encoding {
19    Images(ImageEncoding),
20    #[cfg(feature = "ffmpeg")]
21    Videos(ff::VideoEncoding),
22}
23
24#[derive(Debug, Clone, Subcommand)]
25pub enum EncodingCfg {
26    Images {
27        /// Don't render/overwrite existing images.
28        #[arg(long)]
29        skip_existing: bool,
30    },
31    #[cfg(feature = "ffmpeg")]
32    Ffmpeg(ff::Cfg),
33}
34
35pub enum Stream {
36    Images(ImageStream),
37    #[cfg(feature = "ffmpeg")]
38    Video(ff::Video),
39}
40
41impl Encoding {
42    pub fn init(cfg: EncodingCfg, _resolution: (u32, u32), _fps: u32, parallelized: bool) -> Self {
43        match cfg {
44            EncodingCfg::Images { .. } => {
45                let encoding = ImageEncoding::init(parallelized);
46                Self::Images(encoding)
47            }
48            #[cfg(feature = "ffmpeg")]
49            EncodingCfg::Ffmpeg(cfg) => {
50                let encoding = ff::VideoEncoding::init(_resolution, _fps, cfg);
51                Self::Videos(encoding)
52            }
53        }
54    }
55
56    pub fn new_stream(&mut self, _path: PathBuf) -> Stream {
57        match self {
58            Encoding::Images(encode) => Stream::Images(encode.new_stream()),
59            #[cfg(feature = "ffmpeg")]
60            Encoding::Videos(encode) => Stream::Video(encode.new_video(_path)),
61        }
62    }
63
64    pub fn finalize(self) {
65        match self {
66            Encoding::Images(encode) => encode.finish(),
67            #[cfg(feature = "ffmpeg")]
68            Encoding::Videos(encode) => encode.finish(),
69        }
70    }
71}
72
73impl Stream {
74    pub fn push_frame(&mut self, frame: DownloadTexture, path: PathBuf, _is_final_frame: bool) {
75        match self {
76            Self::Images(stream) => stream.push_frame(frame, path),
77            #[cfg(feature = "ffmpeg")]
78            Self::Video(video) => video.push_frame(frame, _is_final_frame),
79        }
80    }
81}
82
83pub struct ImageEncoding {
84    encoding_thread: thread::JoinHandle<()>,
85    frame_sender: mpsc::SyncSender<(DownloadTexture, PathBuf)>,
86}
87
88pub struct ImageStream {
89    frame_sender: mpsc::SyncSender<(DownloadTexture, PathBuf)>,
90}
91
92impl ImageEncoding {
93    fn init(parallelized: bool) -> Self {
94        let (frame_sender, frame_receiver) = mpsc::sync_channel(50);
95        let encoding_thread = thread::Builder::new()
96            .name("Image Encoder".to_string())
97            .spawn(move || {
98                if parallelized {
99                    // TODO: It's unfortunate to fall back to single threaded here.
100                    // However, multiple par-iters will pretty much deadlock.
101                    // One image encoder thread could be shared globally.
102                    frame_receiver
103                        .into_iter()
104                        .for_each(|frame: (DownloadTexture, PathBuf)| {
105                            frame.0.download_mapped_rgba().save(frame.1).unwrap();
106                        });
107                } else {
108                    frame_receiver.into_iter().par_bridge().for_each(
109                        |frame: (DownloadTexture, PathBuf)| {
110                            frame.0.download_mapped_rgba().save(frame.1).unwrap();
111                        },
112                    );
113                }
114            })
115            .unwrap();
116        Self {
117            encoding_thread,
118            frame_sender,
119        }
120    }
121
122    fn new_stream(&self) -> ImageStream {
123        ImageStream {
124            frame_sender: self.frame_sender.clone(),
125        }
126    }
127
128    fn finish(self) {
129        std::mem::drop(self.frame_sender);
130        self.encoding_thread.join().unwrap();
131    }
132}
133
134impl ImageStream {
135    fn push_frame(&self, frame: DownloadTexture, path: PathBuf) {
136        frame.map_then_send_sync(self.frame_sender.clone(), path)
137    }
138}
twgpu_tools/encoding.rs

twgpu_tools/
encoding.rs
