ml_predict.py

from PIL import Image
from tensorflow import cast
import tensorflow
from tensorflow.image import resize
import numpy as np
from tensorflow.keras.models import load_model
import os


def preprocess_img(image):
    image = cast(image, tensorflow.float32)
    image = image / 255.
    image = resize(image, [200, 200])
    return image


def predict_defect(image):
    reference = {0:"Crazing", 1:"Inclusion", 2:"Patches", 3:"Pitted", 4:"Rolled", 5:"Scratches"}
    image = preprocess_img(image)
    subdir='model_files'
    path = os.path.join(subdir,"Metal_surface_defect_detector.h5")
    model=load_model(path)
    result = np.argmax(model.predict(np.array([image])))
    # return {reference[i]: float(result[i]) for i in range(6)}
    return reference[result]