# Checking your version of Python
import sys
print(sys.version)

#import required libraries
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras import layers, models
from keras.models import Sequential

import os

# The original source of fake and real image is the OpenForensics++ Dataset
# available at https://zenodo.org/records/5528418#.YpdlS2hBzDd under the CC-BY 4.0 License.

# A subset of above dataset is available at Kaggle here:
#https://www.kaggle.com/code/diegoguizanlopez/it-s-an-ai-or-a-person/input


#The dataset used in this implementation is a subset of Kaggle dataset
# I have created this folder "shorter_df_images" to contain subset of images downloaded from Kaggle
# You can use entire dataset downloaded from Kaggle and save it using any name. Use the name of your folder 
# as value of "imagedataset_path"

imagedataset_path = "shorter_df_images"

# Training images are in the subfolder called "Train"
train_dir = os.path.join(imagedataset_path,"Train")

# Test images are in the subfolder called "Test"
test_dir = os.path.join(imagedataset_path,"Test")

# Validation images are in the subfolder called "Validation"
val_dir = os.path.join(imagedataset_path,"Validation")

#Loading Images
#Number of images in one batch
batch_size = 32 

#image height and width
img_height  = 128
img_width = 128


#Create a tensor of training images present in the train_dir
train_ds = tf.keras.utils.image_dataset_from_directory(
           train_dir,
           labels='inferred',
           color_mode="rgb",
           seed=42,
           batch_size=batch_size,
           image_size=(img_height, img_width))

#Create a tensor of test images present in the test_dir
test_ds = tf.keras.utils.image_dataset_from_directory(
           test_dir,
           labels='inferred',
           color_mode="rgb",
           seed=42,
           batch_size=batch_size,
           image_size=(img_height, img_width))

#Create a tensor of validation images present in the val_dir
validation_ds = tf.keras.utils.image_dataset_from_directory(
           val_dir,
           labels='inferred',
           color_mode="rgb",
           seed=42,
           batch_size=batch_size,
           image_size=(img_height, img_width))

# Print the unique class labels from images. This comes from the unique folder names 
#within Train, Test and Validation folders.
class_names = train_ds.class_names
print(class_names)

# DisplayPlot first 9 images from training dataset using matplotlib.pyplot
import matplotlib.pyplot as plt
plt.figure(figsize=(10, 10))
for images, labels in train_ds.take(1):
    for i in range(9):
        ax = plt.subplot(3, 3, i + 1)
        plt.imshow(images[i].numpy().astype("uint8"))
        plt.title(class_names[labels[i]])
        plt.axis("off")

#print the batch size and shape of each image in a batch of training data
for image_batch, labels_batch in train_ds:
    print(image_batch.shape)
    #print the number of labels in each batch
    print(labels_batch.shape)
    break

# let us not define a deep CNN model for image feature extraction and classification
img_width = 128
img_height = 128
#Input() defines the input layer of the CNN
#Its shape argument specified image width, height and number of channels.
inputs = Input(shape = (img_width, img_height, 3))
# Sequential() initializes a sequential deep model.
model = models.Sequential()
# add the input layer to the model
model.add(inputs)
#add first convolutional block with two convolutional layers, and a MaxPooling2D layer.
model.add(Conv2D(32,(3,3),activation="relu", padding='same'))
model.add(Conv2D(32,(3,3),activation="relu", padding='same'))
model.add(MaxPooling2D((2,2), strides=(2, 2)))
#add second convolutional block with two convolutional layers, and a MaxPooling2D layer.
model.add(Conv2D(64,(3,3),activation="relu", padding='same'))
model.add(Conv2D(64,(3,3),activation="relu", padding='same'))
model.add(MaxPooling2D((2,2), strides=(2, 2)))
#add third convolutional block with two convolutional layers, and a MaxPooling2D layer.
model.add(Conv2D(128,(3,3),activation="relu", padding='same'))
model.add(Conv2D(128,(3,3),activation="relu", padding='same'))
model.add(MaxPooling2D((2,2), strides=(2, 2)))
#add 4th convolutional block with two convolutional layers, and a MaxPooling2D layer.
model.add(Conv2D(128,(3,3),activation="relu", padding='same'))
model.add(Conv2D(128,(3,3),activation="relu", padding='same'))
model.add(MaxPooling2D((2,2), strides=(2, 2)))
#add 5th convolutional block with two convolutional layers, and a MaxPooling2D layer.
model.add(Conv2D(128,(3,3),activation="relu", padding='same'))
model.add(Conv2D(128,(3,3),activation="relu", padding='same'))
model.add(MaxPooling2D((2,2), strides=(2, 2)))
#add a Flatten() layer to convert extracted features into a row vector.
model.add(Flatten())
#add a dropout layer to reduce overfitting.
model.add(Dropout(0.5))
#add two fully connected layers with 128 and 256 neurons respectively
model.add(Dense(128,activation="relu"))
model.add(Dense(256,activation="relu"))
#add a final layer with a single neuron
#this layer gives the final classification label of the test image
model.add(Dense(1,activation="sigmoid"))
#print the model structure summary
model.summary()

#compile the model defined above
#we are using 'adam' as the optimizer, BinaryCrossEntropy() loss as the loss function 
#since this is a two class classification
# problem and accuracy is the evaluation metric.
model.compile('adam',loss=tf.keras.losses.BinaryCrossentropy(), metrics=['accuracy'])

#Let us train the model for 25 epochs
epochs = 25
history = model.fit(train_ds,epochs=epochs,batch_size=32,validation_data=validation_ds)

# plot training accuracy, validation accuracy, training loss and validation loss over 25 epochs
acc = history.history['accuracy']
val_acc = history.history['val_accuracy']

loss = history.history['loss']
val_loss = history.history['val_loss']

epochs = 25
epochs_range = range(epochs)

plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label='Training Accuracy')
plt.plot(epochs_range, val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')

plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label='Training Loss')
plt.plot(epochs_range, val_loss, label='Validation Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()

#Evaluate the trained model on test data
model.evaluate(test_ds)

#Save the trained model to your hard drive
model.save("deepfake_imagedetector_shorterdataset.keras")