from test2 import *

import tensorflow as tf
from sklearn.metrics import f1_score, precision_score, recall_score, accuracy_score
from matplotlib import pyplot as plt
import numpy as np


td = tf.keras.utils.image_dataset_from_directory(
        "../dataset-ultimate-2",
        image_size = (300, 300),
        labels = "inferred",
        label_mode = "binary",
        batch_size = 96)

m.load_weights("model3.keras")

p = [i >= 0 for i in m.predict(td)]
r = np.concatenate([y for x, y in td])

#print(r)
#exit()

def __prep_conf_matr(data):
    output_matrix = np.zeros([2, 2])

    for p, r in zip(*data):
        #print(p, r)
        output_matrix[int(p)][int(r[0])] += 1

    return output_matrix


def __plot_conf_matr(data):
    matr = __prep_conf_matr(data)

    _, ax = plt.subplots()

    ax.matshow(matr, cmap = plt.cm.Blues)

    for i, x in enumerate(matr):
        for j, y in enumerate(x):
            ax.text(i,
                    j,
                    str(round(y)),
                    va = "center",
                    ha = "center")

    plt.show()

#p = mod.predict(ds)
#r = np.concatenate([y for x, y in ds])

print(f"Accuracy  : {accuracy_score(p, r)}")
print(f"Precision : {precision_score(p, r, average = 'binary')}")
print(f"Recall    : {recall_score(p, r, average = 'binary')}")
print(f"F1 Score  : {f1_score(p, r, average = 'binary')}")

__plot_conf_matr([p, r])