use anyhow::anyhow;
use ndarray::{s, Array1, Array3};
use std::fs::OpenOptions;
use std::io::prelude::*;
use std::io::BufReader;
use std::path::PathBuf;
pub struct MnistData {
_image_filename: String,
_label_filename: String,
pub(crate) images: Array3<u8>,
pub(crate) labels: Array1<u8>,
}
impl MnistData {
pub fn new(image_filename: String, label_filename: String) -> std::io::Result<MnistData> {
let image_path = PathBuf::from(image_filename.clone());
let image_file = OpenOptions::new().read(true).open(image_path)?;
let mut image_io = BufReader::new(image_file);
let images = read_image_file(&mut image_io);
let label_path = PathBuf::from(label_filename.clone());
let labels_file = OpenOptions::new().read(true).open(label_path)?;
let mut labels_io = BufReader::new(labels_file);
let labels = read_label_file(&mut labels_io);
Ok(MnistData {
_image_filename: image_filename,
_label_filename: label_filename,
images,
labels,
})
}
pub fn get_labels(&self) -> &Array1<u8> {
&self.labels
}
pub fn get_images(&self) -> &Array3<u8> {
&self.images
}
}
pub fn read_image_file(io_in: &mut dyn Read) -> Array3<u8> {
let toread: u32 = 0;
let it_slice = unsafe {
::std::slice::from_raw_parts_mut(
(&toread as *const u32) as *mut u8,
::std::mem::size_of::<u32>(),
)
};
io_in.read_exact(it_slice).unwrap();
let magic = u32::from_be(toread);
assert_eq!(magic, 2051);
let it_slice = unsafe {
::std::slice::from_raw_parts_mut(
(&toread as *const u32) as *mut u8,
::std::mem::size_of::<u32>(),
)
};
io_in.read_exact(it_slice).unwrap();
let nbitem = u32::from_be(toread);
assert!(nbitem == 60000 || nbitem == 10000);
let it_slice = unsafe {
::std::slice::from_raw_parts_mut(
(&toread as *const u32) as *mut u8,
::std::mem::size_of::<u32>(),
)
};
io_in.read_exact(it_slice).unwrap();
let nbrow = u32::from_be(toread);
assert_eq!(nbrow, 28);
let it_slice = unsafe {
::std::slice::from_raw_parts_mut(
(&toread as *const u32) as *mut u8,
::std::mem::size_of::<u32>(),
)
};
io_in.read_exact(it_slice).unwrap();
let nbcolumn = u32::from_be(toread);
assert_eq!(nbcolumn, 28);
let mut images = Array3::<u8>::zeros((nbrow as usize, nbcolumn as usize, nbitem as usize));
let mut datarow = vec![0u8; nbcolumn as usize];
for k in 0..nbitem as usize {
for i in 0..nbrow as usize {
let it_slice = datarow.as_mut_slice();
io_in.read_exact(it_slice).unwrap();
let mut smut_ik = images.slice_mut(s![i, .., k]);
assert_eq!(nbcolumn as usize, it_slice.len());
assert_eq!(nbcolumn as usize, smut_ik.len());
for j in 0..smut_ik.len() {
smut_ik[j] = it_slice[j];
}
}
}
images
}
pub fn read_label_file(io_in: &mut dyn Read) -> Array1<u8> {
let toread: u32 = 0;
let it_slice = unsafe {
::std::slice::from_raw_parts_mut(
(&toread as *const u32) as *mut u8,
::std::mem::size_of::<u32>(),
)
};
io_in.read_exact(it_slice).unwrap();
let magic = u32::from_be(toread);
assert_eq!(magic, 2049);
let it_slice = unsafe {
::std::slice::from_raw_parts_mut(
(&toread as *const u32) as *mut u8,
::std::mem::size_of::<u32>(),
)
};
io_in.read_exact(it_slice).unwrap();
let nbitem = u32::from_be(toread);
assert!(nbitem == 60000 || nbitem == 10000);
let mut labels_vec = vec![0u8; nbitem as usize];
io_in.read_exact(&mut labels_vec).unwrap();
Array1::from(labels_vec)
}
pub fn load_mnist_train_data(dname: &str) -> anyhow::Result<MnistData> {
let mut image_fname = String::from(dname);
image_fname.push_str("train-images-idx3-ubyte");
let image_path = PathBuf::from(image_fname.clone());
let image_file_res = OpenOptions::new().read(true).open(&image_path);
if image_file_res.is_err() {
log::error!("could not open image file : {:?}", image_fname);
return Err(anyhow!("io error"));
}
let mut label_fname = String::from(dname);
label_fname.push_str("train-labels-idx1-ubyte");
let label_path = PathBuf::from(label_fname.clone());
let label_file_res = OpenOptions::new().read(true).open(&label_path);
if label_file_res.is_err() {
log::error!("could not open label file : {:?}", label_fname);
return Err(anyhow!("io error"));
}
Ok(MnistData::new(image_fname, label_fname).unwrap())
}
pub fn load_mnist_test_data(dname: &str) -> anyhow::Result<MnistData> {
let mut image_fname = String::from(dname);
image_fname.push_str("t10k-images-idx3-ubyte");
let image_path = PathBuf::from(image_fname.clone());
let image_file_res = OpenOptions::new().read(true).open(&image_path);
if image_file_res.is_err() {
log::error!("could not open image file : {:?}", image_fname);
return Err(anyhow!("io error"));
}
let mut label_fname = String::from(dname);
label_fname.push_str("t10k-labels-idx1-ubyte");
let label_path = PathBuf::from(label_fname.clone());
let label_file_res = OpenOptions::new().read(true).open(&label_path);
if label_file_res.is_err() {
log::error!("could not open label file : {:?}", label_fname);
return Err(anyhow!("io error"));
}
Ok(MnistData::new(image_fname, label_fname).unwrap())
}