2023-09-08

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)

©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

推荐阅读更多精彩内容