Spaces:
Sleeping
Sleeping
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 |