From f84e65bd20177472db23d48ae258659d03fa50d1 Mon Sep 17 00:00:00 2001
From: Sludge <96552222+SludgePhD@users.noreply.github.com>
Date: Mon, 4 Mar 2024 02:37:24 +0100
Subject: [PATCH] improve save-jpeg example

---
 examples/save-jpeg.rs | 78 +++++++++++++++++++++++++++++++++++--------
 1 file changed, 64 insertions(+), 14 deletions(-)

diff --git a/examples/save-jpeg.rs b/examples/save-jpeg.rs
index 96c4a9e..2603e98 100644
--- a/examples/save-jpeg.rs
+++ b/examples/save-jpeg.rs
@@ -3,12 +3,18 @@
 //! Uses the [`linuxvideo::stream::ReadStream`] returned by [`linuxvideo::VideoCaptureDevice::into_stream`]
 //! to read image data.
 
-use std::{env, fs::File, io::Write, path::Path};
+use std::{
+    env,
+    ffi::{OsStr, OsString},
+    fs::File,
+    io::{self, Write},
+    path::{Path, PathBuf},
+};
 
-use anyhow::anyhow;
+use anyhow::{anyhow, bail};
 use linuxvideo::{
     format::{PixFormat, PixelFormat},
-    Device,
+    BufType, Device,
 };
 
 fn main() -> anyhow::Result<()> {
@@ -18,31 +24,75 @@ fn main() -> anyhow::Result<()> {
 
     let device = args
         .next()
-        .ok_or_else(|| anyhow!("usage: save-stream <device> <file>"))?;
+        .ok_or_else(|| anyhow!("usage: save-stream <device> <file> [<count>]"))?;
 
     let file_path = args
         .next()
-        .ok_or_else(|| anyhow!("usage: save-stream <device> <file>"))?;
-    let mut file = File::create(file_path)?;
+        .ok_or_else(|| anyhow!("usage: save-stream <device> <file> [<count>]"))?;
 
-    let device = Device::open(Path::new(&device))?;
+    let count: u32 = args.next().map_or(1, |osstr| {
+        osstr
+            .to_str()
+            .expect("invalid UTF-8")
+            .parse()
+            .expect("invalid valud for <count>")
+    });
 
+    let device = Device::open(Path::new(&device))?;
     println!(
         "capabilities: {:?}",
         device.capabilities()?.device_capabilities()
     );
 
-    let capture = device.video_capture(PixFormat::new(u32::MAX, u32::MAX, PixelFormat::JPEG))?;
+    let formats = device
+        .formats(BufType::VIDEO_CAPTURE)
+        .map(|res| res.map(|f| f.pixel_format()))
+        .collect::<io::Result<Vec<_>>>()?;
+    let format = if formats.contains(&PixelFormat::MJPG) {
+        PixelFormat::MJPG
+    } else if formats.contains(&PixelFormat::JPEG) {
+        PixelFormat::JPEG
+    } else {
+        bail!(
+            "webcam does not support JPEG (supported formats are {:?})",
+            formats
+        );
+    };
+
+    let capture = device.video_capture(PixFormat::new(u32::MAX, u32::MAX, format))?;
     println!("negotiated format: {:?}", capture.format());
 
     let mut stream = capture.into_stream()?;
-
     println!("stream started, waiting for data");
-    stream.dequeue(|buf| {
-        file.write_all(&*buf)?;
-        println!("wrote file");
-        Ok(())
-    })?;
+    for i in 0..count {
+        let mut path = PathBuf::from(&file_path);
+        let stem = path.file_stem().unwrap_or(OsStr::new("image"));
+        let ext = path.extension().unwrap_or(OsStr::new("jpg"));
+        let number = if count == 1 {
+            String::new()
+        } else {
+            format!("-{i}")
+        };
+        let filename = [stem, OsStr::new(&number), OsStr::new("."), ext]
+            .into_iter()
+            .collect::<OsString>();
+        path.set_file_name(filename);
+
+        let mut file = File::create(&path)?;
+        stream.dequeue(|buf| {
+            if buf.is_error() {
+                eprintln!("WARNING: error flag is set on buffer");
+            }
+            file.write_all(&*buf)?;
+            println!(
+                "wrote {} bytes to {} (raw buffer size: {} bytes)",
+                buf.len(),
+                path.display(),
+                buf.raw_buffer().len(),
+            );
+            Ok(())
+        })?;
+    }
 
     Ok(())
 }