from PIL import Image
import numpy as np
import cv2
import os
import random

def input(data_set, batch_size):

img_list = (os.listdir('Face_Image/' + data_set + '/ro'))

people = ['messy', 'son', 'ro'] # 나중에 랜덤으로 뽑음.
img_batch = [] # 3차원의 이미지를 쑤셔 넣을 공간.
    Y_Labels = np.zeros((batch_size, 1)) # label 정보가 저장될 공간, 크기는 입력된 batch_size임
    # Y_Labels = np.zeros((batch_size, 3))
    for i in range(int(batch_size)): # 입력된 batch_size만큼 반복됨.
        people_random = random.choice(people) 메시, 손흥민, 호날두를 랜덤으로 뽑음.
        img = Image.open('Face_Image/' + data_set + '/' + people_random + '/' + str(random.choice(img_list))) # 입력된 data_set안에 랜덤으로 뽑은 사람의 사진중 랜덤으로 가져옴.
        img.thumbnail((64, 64)) # 100x100 사이즈를 64x64 사이즈로 압축
        # img_np = np.array(img) # 불러온 이미지를 np.array를 사용하여 [64, 64, 3] 배열로 변환
        tmp = random.randrange(1, 5) # 1~4 랜덤
        blur = cv2.blur(np.array(img), (int(tmp), int(tmp))) # 이미지에 1~4중 랜덤으로 노이즈를 줌.
        img_batch.append(np.array(blur)) # 3차원 이미지를 쑤셔 넣는다.

        if people_random == 'messy': # 처리한 이미지의 폴더가 messy일 경우 Label를 0으로.
            Y_Labels[i][0] = 0
        elif people_random == 'son': # 처리한 이미지의 폴더가 son일 경우 Label를 1으로.
            Y_Labels[i][0] = 1
        elif people_random == 'ro':# 처리한 이미지의 폴더가 ro일 경우 Label를 2으로.
            Y_Labels[i][0] = 2
        print(i) # 처리되고 있는 순서.

    result = np.concatenate([img_temp.reshape(1, 64, 64, 3) for img_temp in img_batch]) # 3차원 이미지를 [1, 64, 64, 3]으로 reshape한 후 result에 쑤셔 넣는다.

    return result, Y_Labels # model로 4차원 data_set과 2차원 Label을 반환한다.

from PIL import Image
import numpy as np
import cv2
import os
import random


def input(data_set, batch_size):


    img_list = (os.listdir('Face_Image/' + data_set + '/ro'))
    people = ['messy', 'son', 'ro']
    img_batch = []
    Y_Labels = np.zeros((batch_size, 1))
    # Y_Labels = np.zeros((batch_size, 3))
    for i in range(int(batch_size)):
        people_random = random.choice(people)
        img = Image.open('Face_Image/' + data_set + '/' + people_random + '/' + str(random.choice(img_list)))
        img.thumbnail((64, 64))
        # img_np = np.array(img)
        tmp = random.randrange(1, 5)
        blur = cv2.blur(np.array(img), (int(tmp), int(tmp)))
        img_batch.append(np.array(blur))

        if people_random == 'messy':
            Y_Labels[i][0] = 0
        elif people_random == 'son':
            Y_Labels[i][0] = 1
        elif people_random == 'ro':
            Y_Labels[i][0] = 2
        print(i)


    result = np.concatenate([img_temp.reshape(1, 64, 64, 3) for img_temp in img_batch])

    return result, Y_Labels

import input
import tensorflow as tf
import input
import numpy as np

def next_batch(num, data, labels):
    idx = np.arange(0, len(data))
    np.random.shuffle(idx)
    idx = idx[:num]
    data_shuffle = [data[i] for i in idx]
    labels_shuffle = [labels[i] for i in idx]

    return np.asarray(data_shuffle), np.asarray(labels_shuffle)


def weight_variable(shape):
    initial = tf.Variable(tf.truncated_normal(shape=shape, stddev=0.05))
    return tf.Variable(initial)


def bias_variable(shape):
    initial = tf.Variable(tf.truncated_normal(shape=shape, stddev=0.05))
    return tf.Variable(initial)

def conv2d(x, w, bias):
    return tf.nn.conv2d(x, w, strides=[1, 1, 1, 1], padding='SAME') + bias


def relu(x):
    return tf.nn.relu(x)


def max_pool_2x2(x):
    return tf.nn.max_pool(x, ksize=[1, 3, 3, 1], strides=[1, 2, 2, 1], padding='SAME')

X = tf.placeholder(tf.float32, shape=[None, 64, 64, 3])
Y_Label = tf.placeholder(tf.float32, shape=[None, 3])
keep_prob = tf.placeholder(tf.float32)


Conv1 = conv2d(X, weight_variable([5, 5, 3, 64]), bias_variable([64]))
Relu1 = relu(Conv1)
Pool1 = max_pool_2x2(Relu1)
# 32x32


Conv2 = conv2d(Pool1, weight_variable([5, 5, 64, 64]), bias_variable([64]))
Relu2 = relu(Conv2)
Pool2 = max_pool_2x2(Relu2)
# 16x16


Conv3 = conv2d(Pool2, weight_variable([4, 4, 64, 128]), bias_variable([128]))
Relu3 = relu(Conv3)
Pool3 = max_pool_2x2(Relu3)
# 8x8


Conv4 = conv2d(Pool3, weight_variable([4, 4, 128, 128]), bias_variable([128]))
Relu4 = relu(Conv4)
Pool4 = max_pool_2x2(Relu4)
# 4x4

w_fc1 = tf.Variable(tf.truncated_normal(shape=[4 * 4 * 128, 512], stddev=5e-2))
b_fc1 = tf.Variable(tf.constant(0.1, shape=[512]))

h_conv5_flat = tf.reshape(Pool4, [-1, 4 * 4 * 128])
h_fc1 = tf.nn.relu(tf.matmul(h_conv5_flat, w_fc1) + b_fc1)

h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

w_fc2 = tf.Variable(tf.truncated_normal(shape=[512, 3], stddev=5e-2))
b_fc2 = tf.Variable(tf.constant(0.1, shape=[3]))
logits = tf.matmul(h_fc1_drop, w_fc2) + b_fc2
y_pred = tf.nn.softmax(logits)

x_train, y_train = input.input('train', 5000)
x_test, y_test = input.input('eval', 10000)

y_train_one_hot = tf.squeeze(tf.one_hot(y_train, 3), axis=1)
y_test_one_hot = tf.squeeze(tf.one_hot(y_test, 3), axis=1)

loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=Y_Label, logits=logits))
train_step = tf.train.RMSPropOptimizer(1e-3).minimize(loss)

correct_prediction = tf.equal(tf.argmax(y_pred, 1), tf.argmax(Y_Label, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))


with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    # save_path = "./model/model_1.ckpt"
    # saver.restore(sess, save_path)

    for step in range(1000):
        batch = next_batch(64, x_train, y_train_one_hot.eval())

        if step % 10 == 0:
            test_batch = next_batch(1000, x_test, y_test_one_hot.eval())
            accuracy_print = accuracy.eval(feed_dict={X: test_batch[0], Y_Label: test_batch[1], keep_prob: 1.0})
            Loss_print = loss.eval(feed_dict={X: batch[0], Y_Label: batch[1], keep_prob: 1.0})
            print(step, accuracy_print, Loss_print)



        sess.run(train_step, feed_dict={X: batch[0], Y_Label: batch[1], keep_prob: 0.8})