alumina/data/
image_folder.rs

1
2use walkdir::WalkDir;
3use image;
4use image::{GenericImage, DynamicImage, Pixel};
5use ndarray::{ArrayD, ArrayViewD, IxDyn};
6use std::path::{PathBuf, Path};
7use data::DataSet;
8use std::usize;
9use std::io::*;
10
11pub const CHANNELS: usize = 3;
12
13pub struct ImageFolder {
14	paths: Vec<PathBuf>,
15}
16
17impl ImageFolder {
18	
19	pub fn new<P: AsRef<Path>>(root_path: P, subfolders: bool) -> ImageFolder {
20		let root_path = root_path.as_ref();
21
22		print!("Loading paths for {} ... ", root_path.to_string_lossy());
23		stdout().flush().ok();
24
25		let walker = WalkDir::new(root_path).max_depth(if subfolders {usize::MAX} else {1}).into_iter();
26		let mut paths = walker.filter_map(|e| e.ok()).filter_map(|e| {
27			let path = e.path();
28			
29			if path.is_file() {
30				if let Some(extension) = path.extension() {
31					let extension = extension.to_string_lossy().to_lowercase();
32					if ["jpg", "jpeg", "png", "bmp", "tiff"].iter().any(|&ext| ext == extension) {
33						Some(path.to_path_buf())
34					} else {
35						None
36					}
37				} else {
38					None
39				}
40			} else {
41				None
42			}
43		}).collect::<Vec<_>>();
44		println!("loaded {} paths.", paths.len()); //TODO change to info!() logging
45		paths.sort();
46		ImageFolder{
47			paths: paths,
48		}
49	}
50}
51
52impl DataSet for ImageFolder {
53	fn get(&mut self, i: usize) -> Vec<ArrayD<f32>> {
54
55		let image = match image::open(&self.paths[i]) {
56			Ok(ref dyn_image) => image_to_data(dyn_image),
57			Err(err) => {
58					eprintln!("Image load error '{}' {}", self.paths[i].to_string_lossy(), err);
59					ArrayD::zeros(IxDyn(&[1, 1, CHANNELS][..]))
60				},
61		};
62
63		vec![image]
64	}
65
66	fn length(&self) -> usize {
67		self.paths.len()
68	}
69
70	fn width(&self) -> usize {
71		1
72	}
73
74	fn components(&self) -> Vec<String> {
75		vec!["Images".to_string()]
76	}
77}
78
79
80
81
82pub fn data_to_image(image_data: ArrayViewD<f32>) -> DynamicImage {
83		assert_eq!(image_data.shape()[2], CHANNELS);
84		let data = &image_data.as_slice().unwrap();
85		let width = image_data.shape()[1] as u32;
86		let height = image_data.shape()[0] as u32;
87
88		let mut img = DynamicImage::new_rgba8(width, height);
89
90		for y in 0..height {
91			for x in 0..width {
92				let data = &data[(x + y*width) as usize*CHANNELS..][..CHANNELS];
93				img.put_pixel(x, y, image::Rgba::from_channels((data[0]*255.0 + 0.5).min(255.0).max(0.0) as u8, (data[1]*255.0 + 0.5).min(255.0).max(0.0) as u8, (data[2]*255.0 + 0.5).min(255.0).max(0.0) as u8, 255u8));	
94			}
95		}
96		// for (x, y, pixel) in imgbuf.enumerate_pixels_mut() {
97		// 		let data = &mut data[(x + y*width) as usize*CHANNELS..][..CHANNELS];
98		// 		*pixel = image::Rgba::from_channels((data[0]*255.0).min(255.0) as u8, (data[1]*255.0).min(255.0) as u8, (data[2]*255.0).min(255.0) as u8, 0u8);		
99		// }
100		img
101
102}
103
104pub fn image_to_data(image: &DynamicImage) -> ArrayD<f32> {
105		let width = image.dimensions().0;
106		let height = image.dimensions().1;
107
108		let mut data = unsafe{
109				ArrayD::uninitialized(IxDyn(&[height as usize, width as usize, CHANNELS][..]))
110			};
111		{
112			let data_slice = data.as_slice_mut().unwrap();
113			for y in 0..height {
114				for x in 0..width {
115					let pixel = image.get_pixel(x, y);
116					let channels = pixel.channels();
117					let data_slice = &mut data_slice[(x + y*width) as usize*CHANNELS..][..CHANNELS];
118					for i in 0..CHANNELS {
119						data_slice[i] = channels[i] as f32/255.0;
120					}
121				}
122			}
123		}
124
125		data
126}
127
128
129#[test]
130fn image_crop_test() {
131	_image_crop_test()
132}
133
134fn _image_crop_test() {
135	use data::{Cropping, DataStream};
136
137	let mut d = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
138    d.push("res");
139
140	let mut images = ImageFolder::new(d, true)
141		.crop(0, &[25, 25, 3], Cropping::Random)
142		.sequential();
143
144	assert_eq!(&[25, 25, 3], images.next()[0].shape());
145}
alumina/data/image_folder.rs

alumina/data/
image_folder.rs
