Update cnn_ela_training.py

cnn_ela_training.py

@@ -18,7 +18,7 @@ import tensorflow as tf
 from sklearn.model_selection import train_test_split
 from sklearn.metrics import confusion_matrix
 import itertools
-from tensorflow.keras.utils import to_categorical # convert to one-hot-encoding
+from tensorflow.keras.utils import to_categorical 
 from keras.models import Sequential
 from keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPool2D
 from tensorflow.keras.optimizers.legacy import RMSprop
@@ -57,6 +57,7 @@ def convert_to_ela_image(path, quality, output_dir, resize=(256, 256)):
     
     return ela_im
 
+
 def shuffle_and_split_data(dataframe, test_size=0.2, random_state=59):
     # Shuffle the DataFrame
     shuffled_df = dataframe.sample(frac=1, random_state=random_state).reset_index(drop=True)
@@ -84,10 +85,8 @@ def labeling(path_real, path_fake):
 
 
 if __name__ == "__main__":
-    
-   
-
-     
+    ##############################################################
+    # handling the dataset , set it and label it
     np.random.seed(22)
     tf.random.set_seed(9)
 
@@ -95,34 +94,29 @@ if __name__ == "__main__":
     traning_real_folder = 'datasets/training_set/real/'
 
 
-
     traning_ela_output = 'datasets/training_set/ela_output/'
-
-
-
-
-
-
     traning_set = labeling(traning_real_folder, traning_fake_folder)
         
 
     X = []
     Y = []
-        
+
+    #################################################################
+    # preprocess the images using ELA method and storing the output.
+    
     for index, row in traning_set.iterrows():
        X.append(array(convert_to_ela_image(row[0], 90,traning_ela_output).resize((128, 128))).flatten() / 255.0)
        Y.append(row[1])
 
-
     X = np.array(X)
     Y = to_categorical(Y, 2)
-
     X = X.reshape(-1, 128, 128, 3)
-
     X_train, X_val, Y_train, Y_val = train_test_split(X, Y, test_size = 0.2, random_state=1,shuffle=True)
 
-    model = Sequential()
 
+    ################################################################################
+    # Cnn network creation
+    model = Sequential()
     model.add(Conv2D(filters = 32, kernel_size = (5,5),padding = 'valid', 
                      activation ='relu', input_shape = (128,128,3)))
     print("Input: ", model.input_shape)
@@ -146,9 +140,11 @@ if __name__ == "__main__":
 
     model.summary()
 
+    #Define optimizer .
     optimizer = RMSprop(lr=0.0005, rho=0.9, epsilon=1e-08, decay=0.0)
+    #setting the model , loss func , mertics , optimizer.
     model.compile(optimizer = optimizer , loss = "categorical_crossentropy", metrics=["accuracy"])
-
+    #setting early stopping to train faster.
     early_stopping = EarlyStopping(monitor='val_acc',
                                   min_delta=0,
                                   patience=2,
@@ -157,11 +153,14 @@ if __name__ == "__main__":
     
     epochs = 22
     batch_size = 100
-    
-    
+    #####################################################
+    #running the model , adding the validation set
     history = model.fit(X_train, Y_train, batch_size = batch_size, epochs = epochs, 
               validation_data = (X_val, Y_val), verbose = 2, callbacks=[early_stopping])
-    
+
+
+    #####################################################
+    #plots and metrics
     plt.plot(history.history['accuracy'])
     plt.plot(history.history['val_accuracy'])
     plt.title('Model accuracy')
@@ -179,8 +178,8 @@ if __name__ == "__main__":
     plt.legend(['Train', 'Validation'], loc='upper left')
     plt.show()
 
-# every training can give different results , we got the best training score so no need to run again
-#    model.save('ELA_CNN_ART_V2.h5')
+# every training can give different results , you can mark the next line as comment when you got the best result running the test set.
+    model.save('ELA_CNN_ART_V2.h5')