许多嵌入式爱好者喜爱使用树莓派进行AI算法的实验，但事实上，我们的智能手机同样能够运行众多算法。没有必要专门购买树莓派，本文将展示如何通过Python在智能手机上实现图像分割。
模型转换
我们从unet开源项目上下载 onnx 模型，然后上传到aimo.aidlux.com网站上进行模型转换，
目标平台为 TFLite，参数设置为FP32，下图为转换结果。
aidlux 推理
下面的代码示例展示了如何使用AIDLite库在GPU上进行模型推理，并且将推理结果输出为图像。
依赖库
- aidlite_gpu：用于加载和运行TFLite模型。
- cv2：用于图像处理。
- numpy：用于数组计算。
- PIL/Image：用于图像操作。
- matplotlib.pyplot：用于图像展示（尽管在当前代码中未使用）。
函数说明
- mask_to_image(mask: np.ndarray)将一个Numpy数组（mask）转换为一个PIL图像。如果数组是二维的，将其直接转换为图像；如果是三维的，取最高概率通道作为图像。
- aidlux_tflite_infer(model_path, img_path, save_path) 执行TFLite模型的推理，并将输出保存为图像。
- model_path: TFLite模型的路径。
- img_path: 待处理的输入图像路径。
- save_path: 保存推理结果的路径。
主要步骤
- 初始化AIDLite类和对象。
- 加载模型。
- 设置输入数据：读取图像，调整大小，归一化，增加维度。
- 执行模型推理。
- 获取并处理推理结果。
- 保存或展示处理后的图像。
在运行此代码前，请确保：
- 安装了所有必要的库。
- 替换了model_path和img_path为实际的路径。
- 有一个有效的TFLite模型文件。
- 运行此代码将执行推理并将结果保存在指定的路径。
# -*- coding: UTF-8 -*-
import aidlite_gpu
import cv2
import os
import time
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
def mask_to_image(mask: np.ndarray):
if mask.ndim == 2:
return Image.fromarray((mask * 255).astype(np.uint8))
elif mask.ndim == 3:
return Image.fromarray((np.argmax(mask, axis=0) * 255 / mask.shape[0]).astype(np.uint8))
def aidlux_tflite_infer(model_path, img_path, save_path):
# step1: 初始化aidlite类并创建aidlite对象
aidlite = aidlite_gpu.aidlite()
print('model initial success!!')
# step2: 加载模型
inp_shape = [256*256*1*4]
out_shape = [256*256*2*4]
value = aidlite.ANNModel(model_path, inp_shape, out_shape, 4, 0)
# step3: 传入模型输入数据
img = cv2.imread(img_path, 0)
img = cv2.resize(img, (256, 256))
img = img[np.newaxis, ...]
img = img / 255.0
img = np.expand_dims(img, axis=0)
img = img.astype(dtype=np.float32)
print("image shape is ", img.shape)
aidlite.setInput_Float32(img)
# step4: 执行推理
start = time.time()
aidlite.invoke()
end = time.time()
print("infer time(ms):{0}", 1000 * (end - start))
# step5: 获取输出
pred = aidlite.getOutput_Float32(0)
# step6: 后处理
pred = np.array(pred)
pred = np.reshape(pred,(2,256,256))
mask_img = mask_to_image(pred)
mask_img.save(save_path)
# mask_img = np.array(mask_img)
# cv2.imshow('mask_img', mask_img)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
if __name__ == '__main__':
model_path = "/home/dataset2aidlux/unetmodel_fp32.tflite"
img_path = "/home/dataset2aidlux/test_imgs/0597.PNG"
save_path = '/home/dataset2aidlux/test_imgs/result_0597.png'
aidlux_tflite_infer(model_path, img_path, save_path)
推理结果
原图
推理结果
