predict.py

import tensorflow as tf
import numpy as np

IMG_SIZE = 160

model = tf.keras.models.load_model("model")

with open("labels.txt") as f:
    class_names = [line.strip() for line in f.readlines()]

def predict(image):
    img = image.resize((IMG_SIZE, IMG_SIZE))
    img_array = np.array(img)

    if img_array.shape[-1] == 4:
        img_array = img_array[:, :, :3]

    img_array = np.expand_dims(img_array, axis=0)

    preds = model.predict(img_array, verbose=0)[0]
    idx = np.argmax(preds)

    return {
        "classe": class_names[idx],
        "confidence": float(preds[idx])
    }