import os
import numpy as np
import tensorflow as tf
from tensorflow import keras
import cv2
import random


class VacuumRecognizer:
    model = keras.models.load_model("D:/Image_dataset/model_color.h5")
    #model = keras.models.load_model("D:/Image_dataset/model.h5")


    def recognize(self, image_path) -> str:
        class_names = ['Banana', 'Cat', 'Earings', 'Plant']


        img = cv2.imread(image_path)
        #img = cv2.imread(image_path, flags=cv2.IMREAD_GRAYSCALE)
        # print(img.shape)
        cv2.imshow("img", img)
        cv2.waitKey(0)
        img = (np.expand_dims(img, 0))

        predictions = self.model.predict(img)[0].tolist()

        print(class_names)
        print(predictions)
        print(max(predictions))
        print(predictions.index(max(predictions)))

        return class_names[predictions.index(max(predictions))]


image_paths = []
image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Banana/')
image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Cat/')
image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Earings/')
image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Plant/')

# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/train/Banana/')
# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Cat/')
# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Earings/')
# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Plant/')
uio = VacuumRecognizer()


for image_path in image_paths:
    dirs = os.listdir(image_path)
    for i in range(10):
        print(uio.recognize(image_path + dirs[random.randint(0, len(dirs)-1)]))