import numpy as np
from scipy.special import expit, softmax
from tensorflow import lite
from pathlib import Path
from skimage.io import imread
from termcolor import colored
np.set_printoptions(suppress=True)
""" 加载tflite """
interpreter = lite.Interpreter(model_path=str('mobile_yolo.tflite'))
interpreter.allocate_tensors()
input_index = [details["index"] for details in interpreter.get_input_details()]
output_index = [details["index"] for details in interpreter.get_output_details()]
""" 加载图片 """
img_paths = list(Path('test_logs').glob('*.jpg'))
img_paths.sort()
imgs = np.array([imread(str(path)) for path in img_paths])
imgs = imgs / 255
""" 推理 """
def infer(img: np.ndarray) -> [np.ndarray, np.ndarray]:
inp = img[np.newaxis, ...].astype('float32')
interpreter.set_tensor(input_index[0], inp)
interpreter.invoke()
predictions = [interpreter.get_tensor(idx)[0] for idx in output_index]
return predictions
tf_res = [infer(img) for img in imgs]
""" 加载bin """
bin_paths = list(Path('output').glob('*.bin'))
bin_paths.sort()
def parser_bin(path: Path) -> [np.ndarray, np.ndarray]:
content = path.open('rb').read()
assert len(content) / 4 == 7 * 10 * 75 + 14 * 20 * 75
out = np.array(np.frombuffer(content, '<f4'))
out = [np.transpose(out[:7 * 10 * 75].reshape(75, 7, 10), (1, 2, 0)),
np.transpose(out[7 * 10 * 75:].reshape(75, 14, 20), (1, 2, 0))]
return out
kmd_res = [parser_bin(path) for path in bin_paths]
""" 解析输出 """
inshape = (224, 320)
anchors = np.array([[[81, 82], [135, 169], [344, 319]], [[10, 14], [23, 27], [37, 58]]])
for i in range(len(kmd_res)):
for j in range(len(output_index)):
grid_shape = tf_res[i][j].shape[0:2]
a = tf_res[i][j].reshape(grid_shape + (3, 25))
b = kmd_res[i][j].reshape(grid_shape + (3, 25))
""" 解析xy """
grid_y = np.tile(np.reshape(np.arange(0, grid_shape[0]), [-1, 1, 1, 1]), [1, grid_shape[1], 1, 1])
grid_x = np.tile(np.reshape(np.arange(0, grid_shape[1]), [1, -1, 1, 1]), [grid_shape[0], 1, 1, 1])
grid = np.concatenate([grid_x, grid_y], axis=-1)
a[..., 0:2] = (expit(a[..., 0:2]) + grid) / grid_shape[::-1] * inshape[::-1]
b[..., 0:2] = (expit(b[..., 0:2]) + grid) / grid_shape[::-1] * inshape[::-1]
""" 解析wh """
a[..., 2:4] = np.exp(a[..., 2:4]) * anchors[j]
b[..., 2:4] = np.exp(b[..., 2:4]) * anchors[j]
""" 解析confendice """
a[..., 4:5] = expit(a[..., 4:5])
b[..., 4:5] = expit(b[..., 4:5])
""" 解析类别 """
a[..., 5] = np.argmax(softmax(a[..., 5:], axis=-1), axis=-1)
b[..., 5] = np.argmax(softmax(b[..., 5:], axis=-1), axis=-1)
print(colored(f'img {i} layer {j} tflite :\n', 'blue'), a[np.where(a[..., 4] > .6)][:, :5])
print(colored(f'img {i} layer {j} kmodel :\n', 'green'), b[np.where(b[..., 4] > .6)][:, :5])
|
