from PIL import Image
from io import BytesIO
import base64
import uuid
import os
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import json

class Classify:
    def __init__(self, model_path, class_mapping):
        self.model_path = model_path
        self.class_mapping = class_mapping
    
    
    def _preprocess_image(self, img_path):
        img = image.load_img(img_path, target_size=(150, 150))
        img_array = image.img_to_array(img)
        img_array = np.expand_dims(img_array, axis=0)
        img_array /= 255.0  # Normalizar os valores dos pixels para o intervalo [0, 1]
        return img_array
    
    def _predict_class(self, img_path):
        img_array = self._preprocess_image(img_path)
        try:
            model = load_model(self.model_path)
            prediction = model.predict(img_array)
            predicted_class_index = np.argmax(prediction)
            return predicted_class_index
        except Exception as e:
            # Se ocorrer uma exceção, imprima uma mensagem de erro
            return f"Erro ao carregar o modelo: {str(e)}"
       
        
    
    def classify(self, file):
        predicted_class = self._predict_class(file)
        return predicted_class
        
 
    def init(self, data):
        image_bytes = data['data'].split(',')[1].encode('utf-8')
        image = Image.open(BytesIO(base64.b64decode(image_bytes)))
        random_filename = str(uuid.uuid4()) + '.png'
        image.save('./tmp/' + random_filename)
        image_path = './tmp/' + random_filename
        result = self.classify(image_path)
        
        
        with open(self.class_mapping, 'r') as file:
            data = json.load(file)
            
        os.remove(image_path)
        result_class = data[str(result)]
        return result_class