Prototype / helper /model.py
WatermelonPapaya's picture
Update helper/model.py
6e6ecde verified
import helper.optimizer_def
from keras.models import load_model
import numpy as np
CONVERT_CLASS_PRED_TO_NAME = ["Common Rust", "Gray Leaf Spot", "Leaf Blight"]
def fetch_model(opt_name: str):
return load_model(f'models/{opt_name}-model-001.h5', safe_mode=False)
def preprocess_image(img):
img = img.resize((224, 224))
img_array = np.array(img)
img_array = np.expand_dims(img_array, axis=0)
return img_array
def classify_image(model, img):
prediction = model.predict(img)
predicted_class = np.argmax(prediction, axis=1)[0]
return CONVERT_CLASS_PRED_TO_NAME[predicted_class]