# import the necessary packages
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.preprocessing.image import load_img
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import AveragePooling2D
from tensorflow.keras.layers import Dropout
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.utils import to_categorical
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import matplotlib.pyplot as plt
import numpy as np
import os
import warnings
warnings.filterwarnings('ignore')

# initialize the initial learning rate, number of epochs to train for, and batch size
INIT_LR = 1e-5

# Number of passes of the entire training dataset the ML algo has completed
# If the bactch size is the whole dataset then number of EPOCHS is the number of iterations
EPOCHS = 22

# Datasets are usually grouped into batches especially when the datasets are very larg
BATCH_SIZ = 30
# Explanation on EPOCHS and BATCH_SIZ:
# The genera l relation where dataset size is d,
# number of epochs is e,
# number of iteration is i
# and batch size is b
# would be d*e=i*b


DIRECTORY = r"F:\Face-Mask-Detection\dataset"
CATEGORIES = ["with_mask", "without_mask"]

# grab the list of images in our dataset directory, then initialize
# the list of data (i.e., images) and class images
print("[INFO] loading images...")

img_data = []
img_labels = []

for category in CATEGORIES:
    path = os.path.join(DIRECTORY, category)
    for img in os.listdir(path):
        img_path = os.path.join(path, img)
        image = load_img(img_path, target_size=(100, 100))
        image = img_to_array(image)
        image = preprocess_input(image)

        img_data.append(image)
        img_labels.append(category)

# perform one-hot encoding on the labels
lb = LabelBinarizer()
labels = lb.fit_transform(img_labels)
labels = to_categorical(labels)

data = np.array(img_data, dtype="float64")
labels = np.array(labels)

(trainX, testX, trainY, testY) = train_test_split (data, labels,
                                 test_size=0.20, stratify=labels,
                                 random_state=28)

# load the MobileNetV2 network, ensuring the head FC layer sets are
# left off
baseModel = MobileNetV2(weights="imagenet", include_top=False,
                        input_tensor=Input(shape=(100, 100, 3)))

# construct the head of the model that will be placed on top of the
# the base model
headModel = baseModel.output
headModel = AveragePooling2D(pool_size=(4, 4))(headModel)
headModel = Flatten(name="flatten")(headModel)
headModel = Dense(128, activation="relu")(headModel)
headModel = Dropout(0.5)(headModel)
headModel = Dense(2, activation="softmax")(headModel)

# place the head FC model on top of the base model (this will become
# the actual model we will train)
model = Model(inputs=baseModel.input, outputs=headModel)
# loop over all layers in the base model and freeze them so they will
# *not* be updated during the first training process
for layer in baseModel.layers:
    layer.trainable = False

# compile our model
print("[INFO] compiling model...")
opt = Adam(lr=INIT_LR, decay=INIT_LR / EPOCHS)
model.compile(loss="binary_crossentropy", optimizer=opt,
              metrics=["accuracy"])

# construct the training image generator for data augmentation
print("[INFO] generating images by changing its properties...")
Img_Gen = ImageDataGenerator(
    rotation_range=25,
    zoom_range=0.25,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.15,
    horizontal_flip=True,
    fill_mode="nearest")

# train the head of the network
print("[INFO] training head...")
H = model.fit(
    Img_Gen.flow(trainX, trainY, batch_size=BATCH_SIZ),
    steps_per_epoch=len(trainX) // BATCH_SIZ,
    validation_data=(testX, testY),
    validation_steps=len(testX) // BATCH_SIZ,
    epochs=EPOCHS)

# make predictions on the testing set
print("[INFO] predicting model...")
predIdxs = model.predict(testX, batch_size=BATCH_SIZ)

# for each image in the testing set we need to find the index of the
# label with corresponding largest predicted probability
predIdxs = np.argmax(predIdxs, axis=1)

# show a nicely formatted classification report
print(classification_report(testY.argmax(axis=1), predIdxs,
                            target_names=lb.classes_))

# serialize the model to disk
print("[INFO] saving mask detector model...")
model.save("mask_detector.model", save_format="h5")

# plot the training loss and accuracy
N = EPOCHS
plt.style.use("ggplot")
plt.figure(figsize=(12,8))
plt.plot(np.arange(0, N), H.history["loss"], label="train_loss")
plt.plot(np.arange(0, N), H.history["val_loss"], label="val_loss")
plt.plot(np.arange(0, N), H.history["acc"], label="train_acc")
plt.plot(np.arange(0, N), H.history["val_acc"], label="val_acc")
plt.title("Training Loss and Accuracy")
plt.xlabel("Epochs...")
plt.ylabel("Loss/Accuracy")
plt.legend(loc="center right")
plt.savefig("accuracy plot.png")


