purpur/
transforms.rs

1use std::{path::Path, fs::File, io::Read, collections::hash_map::DefaultHasher, hash::{Hash, Hasher}};
2
3use image::{ColorType, GenericImageView, imageops::FilterType};
4use rand::prelude::SliceRandom;
5
6use crate::{ImageReturn, utils, create_path_with_hash};
7
8#[derive(Debug)]
9pub enum Apply<'a> {
10    ///Flip vertically
11    FlipV,
12    FlipH,
13    Contrast(f32),
14    Resize(u32, u32),
15    CenterCrop(u32, u32),
16    Encode,
17    GetVec(&'a mut Vec<u8>),
18    GetImgRet(&'a mut ImageReturn),
19    SaveTo(&'a str),
20}
21
22pub struct Transforms<'a> {
23    applies: Vec<Apply<'a>>,
24    shuffle: bool,
25}
26
27impl <'a>Transforms<'a> {
28    pub fn new(applies: Vec<Apply<'a>>) -> Transforms<'a> {
29        Transforms {
30            applies,
31            shuffle: false,
32        }
33    }
34    pub fn shuffle(mut self) -> Self {
35        self.shuffle = !self.shuffle;
36        self
37    }
38
39    pub fn apply<P: AsRef<Path>>(&mut self, path: P) -> Result<(), std::io::Error> {
40        let (found_in, mut exact_paths) = utils::get_paths(path)?;
41
42        let mut rng = rand::thread_rng();
43        
44        if self.shuffle {
45            exact_paths.shuffle(&mut rng);
46        }
47        
48        let mut channels = Vec::new();
49        let mut dims = Vec::new();
50
51        let image_count = exact_paths.len();
52
53        let mut img_encode = false;
54
55        for (idx, image_path) in exact_paths.iter().enumerate() {
56            match image::open(image_path) {
57                Ok(mut img) => {
58                    match img.color() {
59                        ColorType::L8 => channels.push(1),
60                        ColorType::La8 => channels.push(1),
61                        ColorType::Rgb8 => channels.push(3),
62                        ColorType::Rgba8 => channels.push(3),
63                        _ => continue,
64                    }
65                    for apply in self.applies.iter_mut() {
66                        match apply {
67                            Apply::FlipV => img = img.flipv(),
68                            Apply::FlipH => img = img.flipv(),
69                            Apply::Contrast(c) => img = img.adjust_contrast(*c),
70                            Apply::Resize(w, h) => img = img.resize_exact(*w, *h, FilterType::CatmullRom),
71                            Apply::GetVec(vec) => {      
72                                vec.extend_from_slice(img.as_bytes());
73                            },
74                            Apply::SaveTo(p) => {
75                                let mut hasher = DefaultHasher::default();                                
76                                img.as_bytes().hash(&mut hasher);
77                                let hash = hasher.finish();
78                                
79                                let path = create_path_with_hash(image_path, p, hash)?;
80                                //let path = create_new_path_from_old(image_path.path(), p, &mut rng)?;
81                                std::fs::create_dir_all(path.parent().unwrap())?;
82                                img.save(&path).unwrap();
83                            },
84                            Apply::GetImgRet(img_ret) => {
85                                
86                                if img_encode {
87                                    let mut from_encoder = Vec::new();
88                                    let mut file = File::open(image_path).unwrap();
89                                    file.read_to_end(&mut from_encoder).unwrap();
90                                    img_ret.data.append(&mut from_encoder);
91                                } else {
92                                    img_ret.data.extend_from_slice(img.as_bytes());
93                                };
94                                
95                                if idx == image_count-1 {
96                                    **img_ret = ImageReturn {
97                                        found_in: found_in.clone(),
98                                        exact_paths: exact_paths.clone(),
99                                        data: img_ret.data.clone(),
100                                        channels: channels.clone(),
101                                        dims: dims.clone(),
102                                    }    
103                                }
104                            },
105                            Apply::Encode => img_encode = true,
106                            Apply::CenterCrop(nw, nh) => {
107                                let (width, height) = img.dimensions();
108                                
109                                let subtract_w = width - *nw;
110                                let subtract_h = height - *nh;
111                                
112                                img = img.crop_imm(subtract_w/2, subtract_h/2, width-subtract_w, height-subtract_h);
113                            },                            
114                        }
115                    }
116                    let img_dims = img.dimensions();
117                    dims.push((img_dims.0 as usize, img_dims.1 as usize));
118                },
119                Err(_) => continue,
120            }
121        }
122        
123        Ok(())
124    }
125}
purpur/transforms.rs

purpur/
transforms.rs
