# -*- coding: utf-8 -*-
"""CNN projekt házifeladat v2.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1eFVYq9wLHAt4OoacsMEPoScg4CjUXj5D
"""

import numpy as np
import matplotlib.pyplot as plt
import os
import random
import cv2
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPooling2D
from keras import optimizers
from keras.utils import plot_model

!unzip /content/sample_data/dataset.zip

categories = ["anime", "cartoon"]
data_set_folder = "/content/dataset"
IMG_SIZE = 64
training_data = []

def create_training_data():
  print("Creating training data...")
  for category in categories:
    print("Create data for " + category + "...")
    path = os.path.join(data_dir, category)
    class_num = categories.index(category) #0 anime, 1 cartoon
    for folder in os.listdir(path):
      path2 = os.path.join(path, folder)
      for img in os.listdir(path2):
        try:
          img_array = cv2.imread(os.path.join(path2, img))
          img_array = cv2.resize(img_array, (IMG_SIZE, IMG_SIZE))
          training_data.append([img_array, class_num])
        except Exception as e:
          pass
  random.shuffle(training_data)

def generate_data():
  print("Generate train_x, train_y...")
  train_x, train_y = [], []
  for features, label in training_data:
    train_x.append(features)
    train_y.append(label)
  train_x = np.array(train_x).reshape(-1, IMG_SIZE, IMG_SIZE, 3) #3 mert színes
  train_y = np.array(train_y)
  train_x = train_x/255.0 #normalizáljuk a pixeleket
  return train_x, train_y

create_training_data()
train_x, train_y = generate_data()
print(train_x.shape)

print("Build network...")
model = Sequential()
model.add(Conv2D(64, (3,3), activation = "relu", input_shape = train_x.shape[1:])) #mert a -1 shape nem kell
model.add(MaxPooling2D(pool_size = (2,2)))
model.add(Conv2D(64, (3,3), activation = "relu"))
model.add(MaxPooling2D(pool_size=(2,2)))
model.add(Conv2D(64, (3,3), activation = "relu"))
model.add(Flatten())
model.add(Dense(64, activation = "relu"))
model.add(Dropout(0.5))
model.add(Dense(1, activation = "sigmoid"))
model.summary()
print("Start training...")
model.compile(loss = "binary_crossentropy", optimizer = "adam", metrics = ["accuracy"])
history = model.fit(train_x, train_y, batch_size = 32, epochs = 5, validation_split = 0.3) #lehet csinálni olyat, hogy validation_split

plt.plot(history.history["accuracy"], label = "accuracy")
plt.plot(history.history["val_accuracy"], label = "val_accuracy")
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.ylim([0.5, 1])
plt.legend(loc = "lower right")

print("Test some interesting examples...")
test_path = "/content/test"
test_data = list()
for i, img in enumerate(os.listdir(test_path)):
  try:
    img_array = cv2.imread(os.path.join(test_path, img))
    plt.imshow(cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB))
    plt.figure(i+1)
    plt.show()
    img_array = cv2.resize(img_array, (IMG_SIZE, IMG_SIZE))
    img_array = np.array(img_array).reshape(-1, IMG_SIZE, IMG_SIZE, 3)
    img_array = img_array/255.0
    prediction = model.predict(img_array)
    percentage = prediction[0][0]*100
    print(os.path.join(test_path, img)+ ":: Cartoon: " + str(round(percentage, 2)) + "% , Anime: " + str(round(100 - percentage, 2)) + "%")
  except Exception as e:
    pass

plot_model(model, to_file='model.png', show_shapes=True)