Spaces:
Running
Running
File size: 2,669 Bytes
67a9b5d b27c653 67a9b5d b27c653 67a9b5d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import os
import copy
import khandy
import numpy as np
from .base import OnnxModel
from collections import OrderedDict
from .base import check_image_dtype_and_shape
class InsectIdentifier(OnnxModel):
def __init__(self):
current_dir = os.path.dirname(os.path.abspath(__file__))
model_path = os.path.join(
current_dir,
"__pycache__/Genius-Society/insecta/quarrying_insect_identifier.onnx",
)
label_map_path = os.path.join(
current_dir,
"__pycache__/Genius-Society/insecta/quarrying_insectid_label_map.txt",
)
super(InsectIdentifier, self).__init__(model_path)
self.label_name_dict = self._get_label_name_dict(label_map_path)
self.names = [
self.label_name_dict[i]["chinese_name"]
for i in range(len(self.label_name_dict))
]
self.num_classes = len(self.label_name_dict)
@staticmethod
def _get_label_name_dict(filename):
records = khandy.load_list(filename)
label_name_dict = {}
for record in records:
label, chinese_name, latin_name = record.split(",")
label_name_dict[int(label)] = OrderedDict(
[("chinese_name", chinese_name), ("latin_name", latin_name)]
)
return label_name_dict
@staticmethod
def _preprocess(image):
check_image_dtype_and_shape(image)
# image size normalization
image = khandy.letterbox_image(image, 224, 224)
# image channel normalization
image = khandy.normalize_image_channel(image, swap_rb=True)
# image dtype normalization
# image dtype and value range normalization
mean, stddev = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
image = khandy.normalize_image_value(image, mean, stddev, "auto")
# to tensor
image = np.transpose(image, (2, 0, 1))
image = np.expand_dims(image, axis=0)
return image
def predict(self, image):
inputs = self._preprocess(image)
logits = self.forward(inputs)
probs = khandy.softmax(logits)
return probs
def identify(self, image, topk=5):
assert isinstance(topk, int)
if topk <= 0 or topk > self.num_classes:
topk = self.num_classes
probs = self.predict(image)
topk_probs, topk_indices = khandy.top_k(probs, topk)
results = []
for ind, prob in zip(topk_indices[0], topk_probs[0]):
one_result = copy.deepcopy(self.label_name_dict[ind])
one_result["probability"] = prob
results.append(one_result)
return results
|