1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
use bytes::Bytes;
use image::RgbImage;
use std::error::Error;
use std::fs;
use std::io::Write;
use std::path::Path;

pub fn remove_non_alphanum(s: &str) -> String {
    let s_new: String = s
        .chars()
        .map(|x| match x {
            '0'..='9' => x,
            'A'..='Z' => x.to_ascii_lowercase(),
            'a'..='z' => x,
            ' ' => '_',
            _ => ' ',
        })
        .collect();
    s_new.replace(" ", "_")
}

pub fn model_filename(name: &str) -> String {
    let clean_name = remove_non_alphanum(name);
    format!("models/{}.onnx", clean_name)
}

pub fn save_file_get(url: &str, path: &str) -> Result<(), String> {
    let client = reqwest::blocking::Client::builder()
        .referer(false)
        .build()
        .map_err(|e| e.to_string())?;

    let response = client.get(url).send().map_err(|e| e.to_string())?;

    let status = response.status();
    if !status.is_success() {
        let text = response.text().unwrap();
        return Err(if text.is_empty() {
            status.to_string()
        } else {
            text
        });
    }

    if !Path::new("models/").exists() {
        fs::create_dir("models").map_err(|e| e.to_string())?;
    }
    let mut out = fs::File::create(path).map_err(|e| e.to_string())?;
    out.write_all(&response.bytes().expect("Failed to convert to bytes"));

    Ok(())
}

pub fn read_bytes_url(url: &str) -> reqwest::Result<Bytes> {
    let client = reqwest::blocking::Client::builder()
        .referer(false)
        .build()
        .map_err(|e| e)?;
    let response = client.get(url).send().map_err(|e| e)?;
    response.bytes()
}

pub fn image_from_bytes(bytes: &Bytes) -> Result<RgbImage, Box<dyn Error>> {
    let dynimg = image::load_from_memory(bytes).unwrap();
    let img = dynimg.to_rgb();
    Ok(img)
}