8 月 AI 实战:工业视觉缺点检测
– 基于 tflite 的 yolov8 模型优化和推理
操作视频见 B 站连贯:aidlux 模型优化 + 工业缺点检测~~完满用我的华为手机实现缺点检测的推理 bilibiliaidlux 模型优化 + 工业缺点检测~~完满用我的华为手机实现缺点检测的推理
1 模型优化
将 onnx 模型转化为 tflite 模型
关上网站:http://aimo.aidlux.com/
输出试用账号和明码:账号:AIMOTC001,明码:AIMOTC001
通过页面中的提醒 AI Model Optimizer,顺次执行步骤①上传模型②抉择指标平台③参数设置④转换后果。
通过上述①-④可将 onnx 模型转为 tflite 模型
模型转换过程蕴含如下日志信息
2023-09-07 19:47:05,969 - INFO : Optimization started.
2023-09-07 19:47:05,970 - INFO : [ONNX-SIM] Clean ONNX Model input node.
2023-09-07 19:47:06,733 - INFO : [ONNX2TFLITE] Start converting to TFLITE.
2023-09-07 19:47:28,511 - INFO : Model optimization done.
2 推理的 py 文件
模型采纳课程中提供的 yolov8_slimneck_SIOU.ONNX,转化完模型门路及名称,如下
# 模型
model_path = "/home/lesson3/yolov8_slimneck_SIOU_tflite/yolov8_slimneck_SIOU_fp32.tflite"
# 测试图片门路
image_path = "/home/lesson3/test"
模型推理过程蕴含如下步骤:
- 初始化 aidlite 类并创立 aidlite 对象
aidlite = aidlite_gpu.aidlite()
print("ok")
- 加载模型
value = aidlite.ANNModel(model_path, [640 * 640 * 3 * 4], [8400 * 11 * 4], 4, 0)
print("gpu:", value)
蕴含遍历每一张图片
for root, dirs, files in os.walk(image_path):
num = 0
for file in files:
file = os.path.join(root, file)
frame = cv2.imread(file)
x_scale = frame.shape[1] / 640
y_scale = frame.shape[0] / 640
将图片转换为模型输出的 640*640 尺寸
img = cv2.resize(frame, (640, 640))
# img_copy=img.co
img = img / 255.0
img = np.expand_dims(img, axis=0)
img = img.astype(dtype=np.float32)
print(img.shape)
- 传入模型输出数据
aidlite.setInput_Float32(img)
- 执行推理
start = time.time()
aidlite.invoke()
end = time.time()
timerValue = 1000 * (end - start)
print("infer time(ms):{0}", timerValue)
- 获取输入
pred = aidlite.getOutput_Float32(0)
# print(pred.shape)
pred = np.array(pred)
print(pred.shape)
pred = np.reshape(pred, (8400, 11))
print(pred.shape) # shape=(8400,11)
- 后处理, 解析输入
boxes, scores, classes = postProcess(pred, confThresh, NmsThresh)
- 绘制保留图像
ret_img = draw(frame, x_scale, y_scale, boxes, scores, classes)
ret_img = ret_img[:, :, ::-1]
num += 1
image_file_name = "/home/result/res" + str(num) + ".jpg"
8. 保留图片
cv2.imwrite(image_file_name, ret_img)