1 use std::fs::File;
2 use std::io::{self, Write};
3 use std::path::Path;
4 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
5 use std::sync::Arc;
6 use std::{cell::RefCell, collections::VecDeque, time::Instant};
7 
8 use v4l2r::{
9     device::{
10         poller::PollError,
11         queue::{
12             direction::Capture,
13             dqbuf::DqBuffer,
14             generic::{GenericBufferHandles, GenericQBuffer, GenericSupportedMemoryType},
15             handles_provider::MmapProvider,
16             qbuf::OutputQueueable,
17         },
18     },
19     encoder::*,
20     memory::{MmapHandle, UserPtrHandle},
21     Format,
22 };
23 use v4l2r_utils::framegen::FrameGenerator;
24 
25 use anyhow::ensure;
26 use clap::{App, Arg};
27 
main()28 fn main() {
29     env_logger::init();
30 
31     let matches = App::new("V4L2 stateful encoder")
32         .arg(
33             Arg::with_name("num_frames")
34                 .long("stop_after")
35                 .takes_value(true)
36                 .help("Stop after encoding a given number of buffers"),
37         )
38         .arg(
39             Arg::with_name("device")
40                 .required(true)
41                 .help("Path to the vicodec device file"),
42         )
43         .arg(
44             Arg::with_name("frame_size")
45                 .long("frame_size")
46                 .required(false)
47                 .takes_value(true)
48                 .default_value("640x480")
49                 .help("Size of the frames to encode (e.g. \"640x480\")"),
50         )
51         .arg(
52             Arg::with_name("output_file")
53                 .long("save")
54                 .required(false)
55                 .takes_value(true)
56                 .help("Save the encoded stream to a file"),
57         )
58         .arg(
59             Arg::with_name("output_mem")
60                 .long("output_mem")
61                 .required(false)
62                 .takes_value(true)
63                 .default_value("mmap")
64                 .help("Type of memory to use for the OUTPUT queue (mmap, user or dmabuf)"),
65         )
66         .get_matches();
67 
68     let device_path = matches.value_of("device").unwrap_or("/dev/video0");
69 
70     let mut stop_after = match clap::value_t!(matches.value_of("num_frames"), usize) {
71         Ok(v) => Some(v),
72         Err(e) if e.kind == clap::ErrorKind::ArgumentNotFound => None,
73         Err(e) => panic!("Invalid value for stop_after: {}", e),
74     };
75 
76     let frame_size = matches
77         .value_of("frame_size")
78         .map(|s| {
79             const ERROR_MSG: &str = "Invalid parameter for frame_size";
80             let split: Vec<&str> = s.split('x').collect();
81             if split.len() != 2 {
82                 panic!("{}", ERROR_MSG);
83             }
84             let width: usize = split[0].parse().expect(ERROR_MSG);
85             let height: usize = split[1].parse().expect(ERROR_MSG);
86 
87             (width, height)
88         })
89         .unwrap();
90 
91     let mut output_file = matches
92         .value_of("output_file")
93         .map(|s| File::create(s).expect("Invalid output file specified."));
94 
95     let output_mem = match matches.value_of("output_mem") {
96         Some("mmap") => GenericSupportedMemoryType::Mmap,
97         Some("user") => GenericSupportedMemoryType::UserPtr,
98         Some("dmabuf") => GenericSupportedMemoryType::DmaBuf,
99         _ => panic!("Invalid value for output_mem"),
100     };
101 
102     let lets_quit = Arc::new(AtomicBool::new(false));
103     // Setup the Ctrl+c handler.
104     {
105         let lets_quit_handler = lets_quit.clone();
106         ctrlc::set_handler(move || {
107             lets_quit_handler.store(true, Ordering::SeqCst);
108         })
109         .expect("Failed to set Ctrl-C handler.");
110     }
111 
112     let encoder = Encoder::open(Path::new(&device_path))
113         .expect("Failed to open device")
114         .set_capture_format(|f| {
115             let format: Format = f.set_pixelformat(b"FWHT").apply()?;
116 
117             ensure!(
118                 format.pixelformat == b"FWHT".into(),
119                 "FWHT format not supported"
120             );
121 
122             Ok(())
123         })
124         .expect("Failed to set capture format")
125         .set_output_format(|f| {
126             let format: Format = f
127                 .set_pixelformat(b"RGB3")
128                 .set_size(frame_size.0, frame_size.1)
129                 .apply()?;
130 
131             ensure!(
132                 format.pixelformat == b"RGB3".into(),
133                 "RGB3 format not supported"
134             );
135             ensure!(
136                 format.width as usize == frame_size.0 && format.height as usize == frame_size.1,
137                 "Output frame resolution not supported"
138             );
139 
140             Ok(())
141         })
142         .expect("Failed to set output format");
143 
144     let output_format = encoder
145         .get_output_format()
146         .expect("Failed to get output format");
147     println!("Adjusted output format: {:?}", output_format);
148 
149     let capture_format = encoder
150         .get_capture_format()
151         .expect("Failed to get capture format");
152     println!("Adjusted capture format: {:?}", capture_format);
153 
154     println!(
155         "Configured encoder for {}x{} ({} bytes per line)",
156         output_format.width, output_format.height, output_format.plane_fmt[0].bytesperline
157     );
158 
159     let mut frame_gen = FrameGenerator::new(
160         output_format.width as usize,
161         output_format.height as usize,
162         output_format.plane_fmt[0].bytesperline as usize,
163     )
164     .expect("Failed to create frame generator");
165 
166     const NUM_BUFFERS: usize = 2;
167 
168     let free_buffers: Option<VecDeque<_>> = match output_mem {
169         GenericSupportedMemoryType::Mmap | GenericSupportedMemoryType::DmaBuf => None,
170         GenericSupportedMemoryType::UserPtr => Some(
171             std::iter::repeat(vec![0u8; output_format.plane_fmt[0].sizeimage as usize])
172                 .take(NUM_BUFFERS)
173                 .collect(),
174         ),
175     };
176     let free_buffers = RefCell::new(free_buffers);
177 
178     let dmabufs: Option<VecDeque<_>> = match output_mem {
179         GenericSupportedMemoryType::Mmap | GenericSupportedMemoryType::UserPtr => None,
180         GenericSupportedMemoryType::DmaBuf => Some(
181             v4l2r_utils::dmabuf_exporter::export_dmabufs(&output_format, NUM_BUFFERS)
182                 .unwrap()
183                 .into_iter()
184                 .collect(),
185         ),
186     };
187     let dmabufs = RefCell::new(dmabufs);
188 
189     let input_done_cb = |buffer: CompletedOutputBuffer<GenericBufferHandles>| {
190         let handles = match buffer {
191             CompletedOutputBuffer::Dequeued(mut buf) => buf.take_handles().unwrap(),
192             CompletedOutputBuffer::Canceled(buf) => buf.plane_handles,
193         };
194         match handles {
195             // We have nothing to do for MMAP buffers.
196             GenericBufferHandles::Mmap(_) => {}
197             // For user-allocated memory, return the buffer to the free list.
198             GenericBufferHandles::User(mut u) => {
199                 free_buffers
200                     .borrow_mut()
201                     .as_mut()
202                     .unwrap()
203                     .push_back(u.remove(0).0);
204             }
205             GenericBufferHandles::DmaBuf(d) => {
206                 dmabufs.borrow_mut().as_mut().unwrap().push_back(d);
207             }
208         };
209     };
210 
211     let mut total_size = 0usize;
212     let start_time = Instant::now();
213     let poll_count_reader = Arc::new(AtomicUsize::new(0));
214     let poll_count_writer = Arc::clone(&poll_count_reader);
215     let mut frame_counter = 0usize;
216     let output_ready_cb = move |cap_dqbuf: DqBuffer<Capture, Vec<MmapHandle>>| {
217         let bytes_used = *cap_dqbuf.data.get_first_plane().bytesused as usize;
218         // Ignore zero-sized buffers.
219         if bytes_used == 0 {
220             return;
221         }
222 
223         total_size = total_size.wrapping_add(bytes_used);
224         let elapsed = start_time.elapsed();
225         frame_counter += 1;
226         let fps = frame_counter as f32 / elapsed.as_millis() as f32 * 1000.0;
227         let ppf = poll_count_reader.load(Ordering::SeqCst) as f32 / frame_counter as f32;
228         print!(
229             "\rEncoded buffer {:#5}, index: {:#2}), bytes used:{:#6} total encoded size:{:#8} fps: {:#5.2} ppf: {:#4.2}" ,
230             cap_dqbuf.data.sequence(),
231             cap_dqbuf.data.index(),
232             bytes_used,
233             total_size,
234             fps,
235             ppf,
236         );
237         io::stdout().flush().unwrap();
238 
239         if let Some(ref mut output) = output_file {
240             let mapping = cap_dqbuf
241                 .get_plane_mapping(0)
242                 .expect("Failed to map capture buffer");
243             output
244                 .write_all(mapping.as_ref())
245                 .expect("Error while writing output data");
246         }
247     };
248 
249     let mut encoder = encoder
250         .allocate_output_buffers_generic::<GenericBufferHandles>(output_mem, NUM_BUFFERS)
251         .expect("Failed to allocate OUTPUT buffers")
252         .allocate_capture_buffers(NUM_BUFFERS, MmapProvider::new(&capture_format))
253         .expect("Failed to allocate CAPTURE buffers")
254         .set_poll_counter(poll_count_writer)
255         .start(input_done_cb, output_ready_cb)
256         .expect("Failed to start encoder");
257 
258     while !lets_quit.load(Ordering::SeqCst) {
259         if let Some(max_cpt) = &mut stop_after {
260             if *max_cpt == 0 {
261                 break;
262             }
263             *max_cpt -= 1;
264         }
265 
266         let v4l2_buffer = match encoder.get_buffer() {
267             Ok(buffer) => buffer,
268             // If we got interrupted while waiting for a buffer, just exit normally.
269             Err(GetBufferError::PollError(PollError::EPollWait(nix::errno::Errno::EINTR))) => break,
270             Err(e) => panic!("{}", e),
271         };
272         let bytes_used = frame_gen.frame_size();
273         match v4l2_buffer {
274             GenericQBuffer::Mmap(buf) => {
275                 let mut mapping = buf
276                     .get_plane_mapping(0)
277                     .expect("Failed to get MMAP mapping");
278                 frame_gen
279                     .next_frame(&mut mapping)
280                     .expect("Failed to generate frame");
281                 buf.queue(&[bytes_used])
282                     .expect("Failed to queue input frame");
283             }
284             GenericQBuffer::User(buf) => {
285                 let mut buffer = free_buffers
286                     .borrow_mut()
287                     .as_mut()
288                     .unwrap()
289                     .pop_front()
290                     .expect("No backing buffer to bind");
291                 frame_gen
292                     .next_frame(&mut buffer)
293                     .expect("Failed to generate frame");
294                 buf.queue_with_handles(
295                     GenericBufferHandles::from(vec![UserPtrHandle::from(buffer)]),
296                     &[bytes_used],
297                 )
298                 .expect("Failed to queue input frame");
299             }
300             GenericQBuffer::DmaBuf(buf) => {
301                 let buffer = dmabufs
302                     .borrow_mut()
303                     .as_mut()
304                     .unwrap()
305                     .pop_front()
306                     .expect("No backing dmabuf to bind");
307                 let mut mapping = buffer[0].map().unwrap();
308                 frame_gen
309                     .next_frame(&mut mapping)
310                     .expect("Failed to generate frame");
311                 buf.queue_with_handles(GenericBufferHandles::from(buffer), &[bytes_used])
312                     .expect("Failed to queue input frame");
313             }
314         }
315     }
316 
317     encoder.stop().unwrap();
318 
319     // Insert new line since we were overwriting the same one
320     println!();
321 
322     if output_mem == GenericSupportedMemoryType::UserPtr {
323         // All the OUTPUT buffers should have been returned
324         assert_eq!(free_buffers.borrow().as_ref().unwrap().len(), NUM_BUFFERS);
325     }
326 }
327