|
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 |