import tensorflow as tf
import numpy as np
import string
from tensorflow.keras.preprocessing.image import load_img

def image_preprocessing(image):
    # Hyperparameters
    img_height = 220
    img_width = 220
    
    # Resize gambar
    image_resized = image.resize((img_height, img_width))
    
    # Konversi gambar menjadi array numpy
    img_array = np.array(image_resized)
    
    # Tambahkan dimensi batch (1 gambar) ke array
    img_array = np.expand_dims(img_array, axis=0)
    
    return img_array


class AnimalClassifier:
    def __init__(self):
        self.class_mapping = {0:'Cheetah', 1:'Leopard', 2:'Lion', 3:'Puma', 4:'Tiger'}
        self.unknown_class = 'Unknown'

    def model_pred(self, model, feature):
        pred = model.predict(feature)
        index = []
        for i in pred:
            index.append(np.argmax(i) if i[np.argmax(i)] > 0.5 else len(self.class_mapping))
            
        # Convert the index to a class name using the specified mapping
        index = [self.class_mapping.get(i, self.unknown_class) for i in index]
        
        return index