#!/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" 
import ktrain
from ktrain import vision as vis


# # Image Regression: Age Prediction
# 
# In this example, we will build a model that predicts the age of a person given the person's photo. 
# 
# ## Download the Dataset
# 
# From [this blog post](https://medium.com/analytics-vidhya/fastai-image-regression-age-prediction-based-on-image-68294d34f2ed) by Abhik Jha, we see that there are several face datasets with age annotations from which to choose:
# - [UTK Face Dataset](http://aicip.eecs.utk.edu/wiki/UTKFace)
# - [IMDb-Wiki Face Dataset](https://data.vision.ee.ethz.ch/cvl/rrothe/imdb-wiki/)
# - [Appa Real Face Dataset](http://chalearnlap.cvc.uab.es/dataset/26/description/)
# 
# In this notebook, we use the UTK Face Dataset.   Download the data from [http://aicip.eecs.utk.edu/wiki/UTKFace](http://aicip.eecs.utk.edu/wiki/UTKFace) and extrct each of the three zip files to the same folder.
# 
# 
# ## STEP 1: Load and Preprocess the Dataset
# 
# The target **age** attribute in this dataset is encoded in the filename.  More specifically, filenames are of the form:
# ```
#  [age]_[gender]_[race]_[date&time].jpg
# ```
# where
# - `[age]` is an integer from 0 to 116, indicating the age
# - `[gender]` is either 0 (male) or 1 (female)
# - `[race]` is an integer from 0 to 4, denoting White, Black, Asian, Indian, and Others (like Hispanic, Latino, Middle Eastern).
# - `[date&time]` is in the format of yyyymmddHHMMSSFFF, showing the date and time an image was collected to UTKFace
# 
# We are only interested in extracting the age for use as a numerical target.  Let us first construct a regular expression to extract the age from the filename.  Then, we can supply the pattern to `images_from_fname` to load and preprocess the dataset.  Supplying `is_regression=True` is important here, as it tells *ktrain* that the integer targets representing age should be treated as a numerical target, as oppposed to a class label.

# In[2]:


# build a regular expression that extracts the age from file name
PATTERN = r'([^/]+)_\d+_\d+_\d+.jpg$'
import re
p = re.compile(PATTERN)
r = p.search('/hello/world/40_1_0_20170117134417715.jpg')
print("Extracted Age:%s" % (int(r.group(1))))


# Set `DATADIR` to the folder where you extracted all the images.

# In[3]:


DATADIR='data/age_estimation/images'
data_aug = vis.get_data_aug(horizontal_flip=True)
(train_data, val_data, preproc) = vis.images_from_fname(DATADIR, pattern = PATTERN, data_aug = data_aug, 
                                                        is_regression=True, random_state=42)


# From the warnings above, we see that a few filenames in the dataset are constructed incorrectly. For instance, the first filename incorrectly has two consecutive underscore characters after the age attribute. Although the age attribute appears to be intact despite the errors and we could modify the regular expression to process these files, we will ignore them in this demonstration.
# 
# ## STEP 2: Create a  Model and Wrap in `Learner`
# 
# We use the `image_regression_model` function to create a `ResNet50` model.  By default, the model freezes all layers except the final randomly-initialized dense layer.

# In[4]:


vis.print_image_regression_models()


# In[5]:


model = vis.image_regression_model('pretrained_resnet50', train_data, val_data)


# In[6]:


# wrap model and data in Learner object
learner = ktrain.get_learner(model=model, train_data=train_data, val_data=val_data, 
                             workers=8, use_multiprocessing=False, batch_size=64)


# ## STEP 3: Estimate Learning Rate
# 
# We will select a learning rate associated with falling loss from the plot displayed.

# In[7]:


learner.lr_find(max_epochs=2)


# In[8]:


learner.lr_plot()


# From the plot above, we choose a learning rate of `1e-4`.
# 
# ## STEP 4: Train Model
# 
# We will begin by training the model for 3 epochs using a [1cycle](https://arxiv.org/abs/1803.09820) learning rate policy.

# In[9]:


learner.fit_onecycle(1e-4, 3)


# In[10]:


learner.freeze(15) # unfreeze all but the first 15 layers


# In[11]:


learner.fit_onecycle(1e-4, 2)


# After only 5 epochs, our validation MAE is 6.57.  That is, on average, our age predictions are off about 6 1/2 years.  Since it does not appear that we are overfitting yet, we could try training further for further improvement, but we will stop here for now.

# ## Make Predictions
# 
# Let's make predictions on individual photos. We could either randomly select from the entire image directory or select just from the validation images.

# In[12]:


# get a Predictor instance that wraps model and Preprocessor object
predictor = ktrain.get_predictor(learner.model, preproc)


# In[13]:


# get some random file names of images
get_ipython().getoutput('ls {DATADIR} | sort -R |head -10')


# In[41]:


# how to get validation filepaths
val_data.filenames[10:20]


# In[42]:


def show_prediction(fname):
    fname = DATADIR+'/'+fname
    predicted = round(predictor.predict_filename(fname)[0])
    actual = int(p.search(fname).group(1))
    vis.show_image(fname)
    print('predicted:%s | actual: %s' % (predicted, actual))


# In[37]:


show_prediction('25_1_3_20170104231305129.jpg')


# In[44]:


show_prediction('10_0_0_20170110225013755.jpg')


# In[45]:


show_prediction('71_0_0_20170111204446510.jpg')


# In[48]:


predictor.save('/tmp/age_predictor')


# In[49]:


reloaded_predictor = ktrain.load_predictor('/tmp/age_predictor')


# In[52]:


reloaded_predictor.predict_filename(DATADIR+'/60_0_3_20170119205726487.jpg')


# In[53]:


vis.show_image(DATADIR+'/60_0_3_20170119205726487.jpg')