sarek 0.1.0

A work-in-progress, experimental neural network library utilizing TensorFlow Keras
use {
    byteorder::{
        BigEndian,
        ReadBytesExt
    },
    flate2::{
        read::{
            GzDecoder
        }
    },
    log::{
        info
    },
    sarek::{
        DataSet,
        SliceSource
    },
    std::{
        io::{
            Cursor,
            Read
        }
    }
};

static MNIST_TRAINING_IMAGES: &[u8] = include_bytes!( "data/train-images-idx3-ubyte.gz" );
static MNIST_TRAINING_LABELS: &[u8] = include_bytes!( "data/train-labels-idx1-ubyte.gz" );
static MNIST_TEST_IMAGES: &[u8] = include_bytes!( "data/t10k-images-idx3-ubyte.gz" );
static MNIST_TEST_LABELS: &[u8] = include_bytes!( "data/t10k-labels-idx1-ubyte.gz" );

fn create_reader( data: &'static [u8] ) -> Cursor< Vec< u8 > > {
    let mut decoder = GzDecoder::new( data );
    let mut data = Vec::new();
    decoder.read_to_end( &mut data ).unwrap();

    Cursor::new( data )
}

fn load_images( images: &'static [u8] ) -> (u32, u32, u32, Vec< u8 >) {
    let mut reader = create_reader( images );
    let magic = reader.read_u32::< BigEndian >().unwrap();
    assert_eq!( magic, 0x00000803 );

    let count = reader.read_u32::< BigEndian >().unwrap();
    let height = reader.read_u32::< BigEndian >().unwrap();
    let width = reader.read_u32::< BigEndian >().unwrap();

    let position = reader.position() as usize;
    let data = reader.into_inner()[ position.. ].to_vec();

    (count, width, height, data)
}

fn load_labels( labels: &'static [u8] ) -> (u32, Vec< u8 >) {
    let mut reader = create_reader( labels );
    let magic = reader.read_u32::< BigEndian >().unwrap();
    assert_eq!( magic, 0x00000801 );

    let count = reader.read_u32::< BigEndian >().unwrap();
    let position = reader.position() as usize;
    let data = reader.into_inner()[ position.. ].to_vec();

    (count, data)
}

type VecSource< T > = SliceSource< T, Vec< T > >;

fn load( raw_labels: &'static [u8], raw_images: &'static [u8] ) -> DataSet< VecSource< f32 >, VecSource< u32 > > {
    let (label_count, labels) = load_labels( raw_labels );
    let (image_count, width, height, image_data) = load_images( raw_images );

    assert_eq!( label_count, image_count );

    let images: Vec< f32 > = image_data.into_iter().map( |value| value as f32 / 255.0 ).collect();
    let categories: Vec< u32 > = labels.iter().map( |&category| category as u32 ).collect();

    let images = SliceSource::from( (width as usize, height as usize).into(), images );
    let categories = SliceSource::from( 1.into(), categories );

    DataSet::new( images, categories )
}

pub fn load_training_data_set() -> DataSet< VecSource< f32 >, VecSource< u32 > > {
    info!( "Starting to load the training data set..." );
    let data = load( MNIST_TRAINING_LABELS, MNIST_TRAINING_IMAGES );
    info!( "Finished loading the training data set!" );
    data
}

pub fn load_test_data_set() -> DataSet< VecSource< f32 >, VecSource< u32 > > {
    info!( "Starting to load the test data set..." );
    let data = load( MNIST_TEST_LABELS, MNIST_TEST_IMAGES );
    info!( "Finished loading the test data set!" );
    data
}