Upload ml_predict.py

ml_predict.py

@@ -0,0 +1,26 @@
+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]
+