tomat / src /model_loader.py
Reyal's picture
Update src/model_loader.py
6c87ba0 verified
raw
history blame contribute delete
999 Bytes
import tensorflow as tf
import numpy as np
from PIL import Image
MODEL_PATH = "/app/src/best_model.h5"
IMG_SIZE = 224
model = tf.keras.models.load_model(MODEL_PATH)
def preprocess(image):
image = image.resize((IMG_SIZE, IMG_SIZE))
image = np.array(image)
image = np.expand_dims(image, axis=0)
return image
def load_model_and_predict(image):
img = preprocess(image)
pred = model.predict(img)
class_index = np.argmax(pred)
confidence = np.max(pred)
# class mapping inside function (safe fallback)
class_names = [
'Tomato___Bacterial_spot',
'Tomato___Early_blight',
'Tomato___Late_blight',
'Tomato___Leaf_Mold',
'Tomato___Septoria_leaf_spot',
'Tomato___Spider_mites Two-spotted_spider_mite',
'Tomato___Target_Spot',
'Tomato___Tomato_Yellow_Leaf_Curl_Virus',
'Tomato___Tomato_mosaic_virus',
'Tomato___healthy'
]
return class_names[class_index], float(confidence)