image_optimization.py

# In[1]:

get_ipython().magic(u'matplotlib inline')
from matplotlib import pyplot as plt

from __future__ import division, print_function, absolute_import

import pandas as pd
import numpy as np
import time as t
import numpy as np

from PIL import Image

import tflearn
from tflearn.data_utils import shuffle, to_categorical
from tflearn.layers.core import input_data, dropout, fully_connected
from tflearn.layers.conv import conv_2d, max_pool_2d
from tflearn.layers.estimator import regression
from tflearn.data_preprocessing import ImagePreprocessing
from tflearn.data_augmentation import ImageAugmentation


# In[ ]:

def load_data():
    X = []
    Y = []

    # dataset 1
    for i in range(1, 196 - 1):
        im_x = Image.open('images/1_before.jpg_' + str(i) + '.jpg')
        X.append(np.asarray(im_x))
        im_x.close()

        im_y = Image.open('images/1_after.jpg_' + str(i) + '.jpg')
        Y.append(np.asarray(im_y))
        im_y.close()

    # dataset 2
    for i in range(1, 64 - 1):
        im_x = Image.open('images/2_before.jpg_' + str(i) + '.jpg')
        X.append(np.asarray(im_x))
        im_x.close()

        im_y = Image.open('images/2_after.jpg_' + str(i) + '.jpg')
        Y.append(np.asarray(im_y))
        im_y.close()

    # dataset 3
    for i in range(1, 441 - 1):
        im_x = Image.open('images/3_before.jpg_' + str(i) + '.jpg')
        X.append(np.asarray(im_x))
        im_x.close()

        im_y = Image.open('images/3_after.jpg_' + str(i) + '.jpg')
        Y.append(np.asarray(im_y))
        im_y.close()

    return zip(X, Y)


# In[ ]:

dataset = np.asarray(load_data())
np.random.shuffle(dataset)

# test 20%
test_percentage = int(len(dataset) * 0.8)
dataset, dataset_test = np.vsplit(dataset, [test_percentage])

X, Y = dataset[:, 0], dataset[:, 1]
X_test, Y_test = dataset_test[:, 0], dataset_test[:, 1]


# In[ ]:

# create DNN
net = tflearn.input_data(shape=[None, 32, 32, 3])
net = tflearn.dropout(net, 0.8)
net = tflearn.fully_connected(net, 32*32, activation='relu')
net = tflearn.dropout(net, 0.5)
net = tflearn.fully_connected(net, 32*8, activation='relu')
net = tflearn.dropout(net, 0.5)
net = tflearn.fully_connected(net, 32*32, activation='relu')
net = tflearn.dropout(net, 0.5)
net = tflearn.fully_connected(net, 32*32*3, activation='relu')

net = tflearn.reshape(net, (-1, 32, 32, 3))
net = tflearn.regression(net, optimizer='adam', learning_rate=0.001,
                         loss='mean_square', metric=None)

model = tflearn.DNN(net, tensorboard_verbose=0)


# learn
batch_size = int(len(X) / 3)
n_epoch = batch_size * 3 * 3 * 3
model.fit(X, Y, n_epoch=n_epoch, run_id="image_optimize", batch_size=batch_size, validation_set=(X_test, Y_test))