|
|
@@ -1,7 +1,9 @@
|
|
|
import argparse
|
|
|
import sophon.sail as sail
|
|
|
import cv2
|
|
|
+import os
|
|
|
import logging
|
|
|
+import json
|
|
|
import numpy as np
|
|
|
|
|
|
class Predictor:
|
|
|
@@ -52,8 +54,8 @@ class Predictor:
|
|
|
def predict(self, input_img):
|
|
|
input_data = {self.input_name: input_img}
|
|
|
outputs = self.net.process(self.graph_name, input_data)
|
|
|
- print('predict fun:', outputs)
|
|
|
- print('predict fun return:', list(outputs.values())[0])
|
|
|
+ # print('predict fun:', outputs)
|
|
|
+ # print('predict fun return:', list(outputs.values())[0])
|
|
|
return list(outputs.values())[0]
|
|
|
|
|
|
def preprocess(self, img):
|
|
|
@@ -68,39 +70,61 @@ class Predictor:
|
|
|
return img
|
|
|
|
|
|
def postprocess(self, outputs):
|
|
|
-
|
|
|
outputs_exp = np.exp(outputs)
|
|
|
- print('exp res:', outputs_exp)
|
|
|
+ # print('exp res:', outputs_exp)
|
|
|
outputs = outputs_exp / np.sum(outputs_exp, axis=1)[:, None]
|
|
|
- print('softmax res:', outputs)
|
|
|
- predictions = np.argmax(outputs, axis=1)
|
|
|
- print('预测结果:', predictions)
|
|
|
- return outputs
|
|
|
-
|
|
|
-def main(args):
|
|
|
- predictor = Predictor()
|
|
|
- filename = args.input
|
|
|
- src_img = cv2.imread(filename, cv2.IMREAD_COLOR)
|
|
|
- src_img = predictor.preprocess(src_img)
|
|
|
- src_img = np.stack([src_img])
|
|
|
- print('图像输入shape:',src_img.shape)
|
|
|
- if src_img is None:
|
|
|
- logging.error("{} imread is None.".format(filename))
|
|
|
- return
|
|
|
- res = predictor(src_img)
|
|
|
- print('预测结果res:', res)
|
|
|
- predictor.postprocess(res)
|
|
|
-
|
|
|
+ # print('softmax res:', outputs)
|
|
|
+ predictions = np.argmax(outputs, axis=1) # 返回最大概率的类别
|
|
|
+ # print('预测结果:', predictions)
|
|
|
+ return predictions
|
|
|
|
|
|
def argsparser():
|
|
|
parser = argparse.ArgumentParser(prog=__file__)
|
|
|
- parser.add_argument('--input', type=str, default='./000000_256_512_1.jpg', help='path of input, must be image directory')
|
|
|
+ parser.add_argument('--input', type=str, default='./test', help='path of input, must be image directory')
|
|
|
parser.add_argument('--bmodel', type=str, default='./shufflenet_f32.bmodel', help='path of bmodel')
|
|
|
parser.add_argument('--dev_id', type=int, default=0, help='tpu id')
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
return args
|
|
|
|
|
|
+def main(args):
|
|
|
+ # 加载推理引擎
|
|
|
+ predictor = Predictor()
|
|
|
+ # 获取图片
|
|
|
+ cls = os.listdir(args.input)
|
|
|
+ # cls = os.listdir('./label_data/test')
|
|
|
+ filename = {}
|
|
|
+ for c in cls:
|
|
|
+ # filename[c] = [os.path.join('./label_data/test', p) for p in os.listdir(os.path.join('./label_data/test', c))]
|
|
|
+ filename[c] = [os.path.join(args.input,c, p) for p in os.listdir(os.path.join(args.input, c))]
|
|
|
+ print('find data:')
|
|
|
+ for k, v in filename.items():
|
|
|
+ print(f'{k}: {len(v)}')
|
|
|
+ print('开始推理...')
|
|
|
+ res_list=[]
|
|
|
+ for k, v in filename.items():
|
|
|
+ for i in v:
|
|
|
+ src_img = cv2.imread(i, cv2.IMREAD_COLOR)
|
|
|
+ src_img = predictor.preprocess(src_img)
|
|
|
+ res = [-1]
|
|
|
+ if src_img is not None:
|
|
|
+ src_img = np.stack([src_img])
|
|
|
+ res = predictor(src_img)
|
|
|
+ res = predictor.postprocess(res)
|
|
|
+ print(f'{k}: {res[0]}')
|
|
|
+ res_list.append([i, k, res[0]])
|
|
|
+ # 保存结果为 json
|
|
|
+ res_dict = {}
|
|
|
+ for path, real_value, pred_value in res_list:
|
|
|
+ res_dict[path] = {'y':int(real_value), 'x':int(pred_value)}
|
|
|
+ res_save_path = f'{predictor.graph_name}_prediction_result.json'
|
|
|
+ with open(res_save_path, 'w') as f:
|
|
|
+ json.dump(res_dict, f)
|
|
|
+ print(f'保存结果为 {res_save_path} 成功!')
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
if __name__ == '__main__':
|
|
|
args = argsparser()
|
|
|
main(args)
|