8月AI實(shí)戰(zhàn)：工業(yè)視覺缺陷檢測

--基于tflite的yolov8模型優(yōu)化和推理

操作視頻見B站連接：aidlux模型優(yōu)化+工業(yè)缺陷檢測～～完美用我的華為手機(jī)實(shí)現(xiàn)缺陷檢測的推理bilibiliaidlux模型優(yōu)化+工業(yè)缺陷檢測～～完美用我的華為手機(jī)實(shí)現(xiàn)缺陷檢測的推理

1 模型優(yōu)化

將onnx模型轉(zhuǎn)化為tflite模型

打開網(wǎng)站：http://aimo.aidlux.com/

輸入試用賬號和密碼：賬號：AIMOTC001 ，密碼：AIMOTC001

通過頁面中的提示AI Model Optimizer，依次執(zhí)行步驟①上傳模型②選擇目標(biāo)平臺③參數(shù)設(shè)置④轉(zhuǎn)換結(jié)果。

通過上述①-④可將onnx模型轉(zhuǎn)為tflite模型

模型轉(zhuǎn)換過程包含如下日志信息

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，轉(zhuǎn)化完模型路徑及名稱，如下

# 模型

model_path = "/home/lesson3/yolov8_slimneck_SIOU_tflite/yolov8_slimneck_SIOU_fp32.tflite"

# 測試圖片路徑

image_path = "/home/lesson3/test"

模型推理過程包含如下步驟：

初始化aidlite類并創(chuàng)建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

將圖片轉(zhuǎn)換為模型輸入的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)

傳入模型輸入數(shù)據(jù)

aidlite.setInput_Float32(img)

執(zhí)行推理

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)

審核編輯 黃宇