[+] feat: add downloader crate

Author: heng30

Cargo.lock

@@ -143,6 +143,7 @@ pmacro = { path = "lib/pmacro" }
 camera = { path = "lib/camera" }
 bytesio = { path = "lib/bytesio" }
 recorder = { path = "lib/recorder" }
+downloader = { path = "lib/downloader" }
 mp4-player = { path = "lib/mp4-player" }
 image-effect = { path = "lib/image-effect" }
 video-encoder = { path = "lib/video-encoder" }

Cargo.toml

@@ -2,51 +2,31 @@
 pub enum Model {
     Modnet,
     Rmbg14,
-    Rmbg14Fp16,
-    Rmbg14Quantized,
-    U2NET,
-    U2NETP,
 }
 
 impl Model {
     pub fn all_models() -> Vec<Self> {
-        vec![
-            Self::Modnet,
-            Self::Rmbg14,
-            Self::Rmbg14Fp16,
-            Self::Rmbg14Quantized,
-            Self::U2NET,
-            Self::U2NETP,
-        ]
+        vec![Self::Modnet, Self::Rmbg14]
     }
 
     pub fn to_input_size(&self) -> (u32, u32) {
         match self {
             Model::Modnet => (512, 512),
-            Model::Rmbg14 | Model::Rmbg14Fp16 | Model::Rmbg14Quantized => (1024, 1024),
-            Model::U2NET | Model::U2NETP => (320, 320),
+            Model::Rmbg14 => (1024, 1024),
         }
     }
 
     pub fn to_str(&self) -> &'static str {
         match self {
             Self::Modnet => "modnet_photographic_portrait_matting.onnx",
             Self::Rmbg14 => "rmbg-1.4.onnx",
-            Self::Rmbg14Fp16 => "rmbg-1.4_fp16.onnx",
-            Self::Rmbg14Quantized => "rmbg-1.4_quantized.onnx",
-            Self::U2NET => "u2net.onnx",
-            Self::U2NETP => "u2netp.onnx",
         }
     }
 
     pub fn try_from(model: &str) -> Option<Self> {
         match model {
             "modnet_photographic_portrait_matting.onnx" => Some(Model::Modnet),
             "rmbg-1.4.onnx" => Some(Model::Rmbg14),
-            "rmbg-1.4_fp16.onnx" => Some(Model::Rmbg14Fp16),
-            "rmbg-1.4_quantized.onnx" => Some(Model::Rmbg14Quantized),
-            "u2net.onnx" => Some(Model::U2NET),
-            "u2netp.onnx" => Some(Model::U2NETP),
             _ => None,
         }
     }
@@ -57,16 +37,6 @@ impl Model {
                 "https://huggingface.co/TheEeeeLin/HivisionIDPhotos_matting/resolve/034769305faf641ad94edfac654aba13be06e816/modnet_photographic_portrait_matting.onnx"
             }
             Self::Rmbg14 => "https://huggingface.co/briaai/RMBG-1.4/resolve/main/onnx/model.onnx",
-            Self::Rmbg14Fp16 => {
-                "https://huggingface.co/briaai/RMBG-1.4/resolve/main/onnx/model_fp16.onnx"
-            }
-            Self::Rmbg14Quantized => {
-                "https://huggingface.co/briaai/RMBG-1.4/resolve/main/onnx/model_quantized.onnx"
-            }
-            Self::U2NET => "https://huggingface.co/AlenZeng/u2netonnxmodel/resolve/main/u2net.onnx",
-            Self::U2NETP => {
-                "https://huggingface.co/martintomov/comfy/resolve/1b0c3477e152d8a2dea8e4e418a6dba32de56fda/rembg/u2netp.onnx"
-            }
         }
     }
 }

lib/background-remover/src/model.rs

@@ -0,0 +1 @@
+*.mp4

lib/downloader/.gitignore

@@ -0,0 +1,20 @@
+[package]
+name = "downloader"
+license.workspace = true
+edition.workspace = true
+version.workspace = true
+readme.workspace = true
+authors.workspace = true
+keywords.workspace = true
+homepage.workspace = true
+repository.workspace = true
+description.workspace = true
+
+[dependencies]
+thiserror.workspace = true
+futures.workspace = true
+reqwest = { workspace = true, features = ["stream"] }
+
+[dev-dependencies]
+anyhow.workspace = true
+tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

lib/downloader/Cargo.toml

@@ -0,0 +1,47 @@
+use anyhow::Result;
+use downloader::{DownloadError, DownloadStatus, Downloader};
+use std::io::Write;
+
+#[tokio::main]
+async fn main() -> Result<(), DownloadError> {
+    let download_url =
+        "https://freetestdata.com/wp-content/uploads/2022/02/Free_Test_Data_1MB_MP4.mp4";
+    let save_path = "./test_video.mp4";
+
+    println!("Starting download from: {}", download_url);
+    println!("Saving to: {}", save_path);
+    println!("Press Ctrl+C to cancel...\n");
+
+    let downloader = Downloader::new(download_url.to_string(), save_path.to_string());
+
+    match downloader
+        .start(|downloaded: u64, total: u64, progress: f32| {
+            let percent = progress * 100.0;
+            let mb_downloaded = downloaded as f64 / 1024.0 / 1024.0;
+            let mb_total = total as f64 / 1024.0 / 1024.0;
+            print!(
+                "\rProgress: {:.2}% ({:.2} MB / {:.2} MB)",
+                percent, mb_downloaded, mb_total
+            );
+            std::io::stdout().flush().unwrap();
+        })
+        .await
+    {
+        Ok(DownloadStatus::Finsished) => {
+            println!("\n✓ Download completed successfully!");
+            println!("File saved to: {}", save_path);
+        }
+        Ok(DownloadStatus::Cancelled) => {
+            println!("\n✗ Download was cancelled!");
+        }
+        Ok(DownloadStatus::Downloading) => {
+            println!("\n⏳ Download in progress...");
+        }
+        Err(e) => {
+            println!("\n✗ Download failed: {}", e);
+            return Err(e);
+        }
+    }
+
+    Ok(())
+}

lib/downloader/examples/downloader_demo.rs

@@ -0,0 +1,98 @@
+use crate::{DownloadError, Result};
+use futures::StreamExt;
+use reqwest::Client;
+use std::{
+    fs,
+    io::Write,
+    sync::{
+        Arc,
+        atomic::{AtomicBool, Ordering},
+    },
+};
+
+pub enum DownloadStatus {
+    Finsished,
+    Cancelled,
+    Downloading,
+}
+
+#[derive(Debug, Clone)]
+#[non_exhaustive]
+pub struct Downloader {
+    url: String,
+    save_path: String,
+    cancel_sig: Arc<AtomicBool>,
+}
+
+impl Downloader {
+    pub fn new(url: String, save_path: String) -> Downloader {
+        Downloader {
+            url,
+            save_path,
+            cancel_sig: Arc::new(AtomicBool::new(false)),
+        }
+    }
+
+    pub async fn start(
+        &self,
+        mut progress_cb: impl FnMut(u64, u64, f32) + 'static,
+    ) -> Result<DownloadStatus> {
+        let tmp_filepath = format!("{}.tmp", self.save_path);
+
+        let mut save_file =
+            fs::File::create(&tmp_filepath).map_err(|e| DownloadError::FileCreateError {
+                error: e,
+                path: tmp_filepath.clone(),
+            })?;
+
+        let response =
+            Client::new()
+                .get(&self.url)
+                .send()
+                .await
+                .map_err(|e| DownloadError::RequestError {
+                    error: e,
+                    url: self.url.to_string(),
+                })?;
+
+        let total_size = response
+            .content_length()
+            .ok_or_else(|| DownloadError::ContentLengthError)?;
+
+        let mut downloaded: u64 = 0;
+        let mut stream = response.bytes_stream();
+
+        while let Some(chunk) = stream.next().await {
+            if self.cancel_sig.load(Ordering::Relaxed) {
+                return Ok(DownloadStatus::Cancelled);
+            }
+
+            let chunk = chunk.map_err(|e| DownloadError::IncompleteDownload {
+                error: e.to_string(),
+                downloaded,
+                total: total_size,
+            })?;
+            save_file.write_all(&chunk)?;
+
+            downloaded += chunk.len() as u64;
+
+            let progress = downloaded as f32 / total_size as f32;
+            progress_cb(downloaded, total_size, progress);
+        }
+
+        if total_size == downloaded {
+            _ = fs::rename(&tmp_filepath, &self.save_path);
+            Ok(DownloadStatus::Finsished)
+        } else {
+            Ok(DownloadStatus::Downloading)
+        }
+    }
+
+    pub fn cancel(&self) {
+        self.cancel_sig.store(true, Ordering::Relaxed);
+    }
+
+    pub fn cancel_sig(&self) -> Arc<AtomicBool> {
+        self.cancel_sig.clone()
+    }
+}

lib/downloader/src/downloader.rs

@@ -0,0 +1,30 @@
+pub mod downloader;
+
+pub use downloader::{DownloadStatus, Downloader};
+
+pub type Result<T> = std::result::Result<T, DownloadError>;
+
+#[derive(thiserror::Error, Debug)]
+pub enum DownloadError {
+    #[error("HTTP request {url} failed. Error: {error}")]
+    RequestError { error: reqwest::Error, url: String },
+
+    #[error("IO error: {0}")]
+    IoError(#[from] std::io::Error),
+
+    #[error("Failed to get content length from response")]
+    ContentLengthError,
+
+    #[error("Download operation cancelled")]
+    Cancelled,
+
+    #[error("Download incomplete: {downloaded}/{total} bytes. Error: {error}")]
+    IncompleteDownload {
+        error: String,
+        downloaded: u64,
+        total: u64,
+    },
+
+    #[error("Failed to create file: {path}. Error: {error}")]
+    FileCreateError { error: std::io::Error, path: String },
+}

lib/downloader/src/lib.rs

@@ -52,6 +52,7 @@ native-dialog.workspace = true
 platform-dirs.workspace = true
 screen-capture.workspace = true
 fast_image_resize.workspace = true
+background-remover.workspace = true
 rodio = { workspace = true, features = ["playback"] }
 
 [target.'cfg(target_os = "linux")'.dependencies]