#!/usr/bin/env python
# coding: utf-8

# In[1]:


get_ipython().run_line_magic('reload_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')
get_ipython().run_line_magic('matplotlib', 'inline')
import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID";
os.environ["CUDA_VISIBLE_DEVICES"]="0"; 


# ## Collect MNIST Dataset as Arrays

# In[2]:


from tensorflow.keras.datasets import mnist
import numpy as np
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train = np.expand_dims(x_train, axis=3)
x_test = np.expand_dims(x_test, axis=3)


# ## STEP 1:  Preprocess Dataset

# In[3]:


import ktrain
from ktrain import vision as vis


# In[4]:


data_aug = vis.get_data_aug(  rotation_range=15,
                              zoom_range=0.1,
                              width_shift_range=0.1,
                              height_shift_range=0.1)
classes = ['zero', 'one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine']


# In[5]:


(trn, val, preproc) = vis.images_from_array(x_train, y_train, 
                                            validation_data=None,
                                            val_pct=0.1,
                                            random_state=42,
                                            data_aug=data_aug,
                                            class_names=classes)


# ## STEP 2: Load Model and Wrap in `Learner`

# In[6]:


# Using a LeNet-style classifier
model = vis.image_classifier('default_cnn', trn, val)


# In[7]:


learner = ktrain.get_learner(model, train_data=trn, val_data=val, batch_size=128)


# ## STEP 3: Find Learning Rate

# In[8]:


learner.lr_find(show_plot=True, max_epochs=3)


# ## STEP 4: Train Model
# 
# We only train for three epochs for demonstration purposes.

# In[9]:


learner.fit_onecycle(1e-3, 3)


# In[10]:


learner.validate(class_names=preproc.get_classes())


# In[11]:


learner.view_top_losses(n=1)


# ## Make Predictions

# In[12]:


predictor = ktrain.get_predictor(learner.model, preproc)


# In[13]:


predictor.predict(x_test[0:1])[0]


# In[14]:


np.argmax(predictor.predict(x_test[0:1], return_proba=True)[0])


# In[15]:


predictor.save('/tmp/my_mnist')


# In[16]:


p = ktrain.load_predictor('/tmp/my_mnist')


# In[17]:


p.predict(x_test[0:1])[0]


# In[18]:


predictions = p.predict(x_test)


# In[19]:


import pandas as pd
df = pd.DataFrame(zip(predictions, y_test), columns=['Predicted', 'Actual'])
df.head()


# In[ ]: