Upload functions.py

functions.py

@@ -0,0 +1,47 @@
+import numpy as np
+import json
+import onnxruntime as rt
+
+model_path = 'model/model.onnx'
+class_path = 'model/birds_name_mapping.json'
+
+normalise_means = [0.4914, 0.4822, 0.4465]
+normalise_stds = [0.2023, 0.1994, 0.2010]
+
+def normalise_image(image):
+    image = image.copy()
+    for i in range(3):
+        image[:, i, :, :] = (image[:, i, :, :] - normalise_means[i]) / normalise_stds[i]
+    return image
+
+def load_class_names():
+    with open(class_path, 'r') as f:
+        class_names = json.load(f)
+    return class_names
+
+def predict(inp_image):
+    
+    class_names = load_class_names()
+
+    image = inp_image
+    image = image.transpose((2, 0, 1))
+
+    image = image / 255.0
+    image = np.expand_dims(image, axis=0)
+    image = normalise_image(image)
+    image = image.astype(np.float32)
+
+    sess = rt.InferenceSession(model_path)
+
+    input_name = sess.get_inputs()[0].name
+    output_name = sess.get_outputs()[0].name
+
+    output = sess.run([output_name], {input_name: image})[0]
+    prob = np.exp(output) / np.sum(np.exp(output), axis=1, keepdims=True)
+
+    top5 = np.argsort(prob[0])[-5:][::-1]
+
+    class_probs = {class_names[str(i)]: float(prob[0][i]) for i in top5}
+    print(class_probs)
+
+    return class_probs