guicon/techchallenge-pet-computer-vision-model

1 2import tensorflow as tf3import numpy as np4 5IMG_SIZE = 1606 7model = tf.keras.models.load_model("model")8 9with open("labels.txt") as f:10    class_names = [line.strip() for line in f.readlines()]11 12def predict(image):13    img = image.resize((IMG_SIZE, IMG_SIZE))14    img_array = np.array(img)15 16    if img_array.shape[-1] == 4:17        img_array = img_array[:, :, :3]18 19    img_array = np.expand_dims(img_array, axis=0)20 21    preds = model.predict(img_array, verbose=0)[0]22    idx = np.argmax(preds)23 24    return {25        "classe": class_names[idx],26        "confidence": float(preds[idx])27    }28