RK3568笔记四：基于TensorFlow花卉图像分类部署

若该文为原创文章，转载请注明原文出处。

基于正点原子的ATK-DLRK3568部署测试。

花卉图像分类任务，使用使用 tf.keras.Sequential 模型，简单构建模型，然后转换成 RKNN 模型部署到ATK-DLRK3568板子上。

在 PC 使用 Windows 系统安装 tensorflow，并创建虚拟环境进行训练，然后切换到VM下的RK3568环境，使用rknn-toolkit2把模型转成rknn模型部署到RK3568板子上测试。

一、介绍

TensorFlow 是一个基于数据流编程（dataflow programming）的符号数学系统，被广泛应用于机器学习（machine learning）算法的编程实现，其前身是谷歌的神经网络算法库 DistBelief。

使用 tf.keras.Sequential 模型对花卉图像进行分类。

二、环境搭建

1、创建虚拟环境

 conda create -n tensorflow_env python=3.8 -y

2、激活环境

conda activate tensorflow_env

3、安装环境

pip install numpypip install tensorflowpip install pillow

三、训练

1、下载数据集

https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz

数据集不好下载，自行处理。

2、训练

tensorflow_classification.py

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential# 获取
import pathlib
#dataset_url = "https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz"
#data_dir = tf.keras.utils.get_file('flower_photos', origin=dataset_url, untar=True)
data_dir = './flower_photos'
data_dir = pathlib.Path(data_dir)batch_size = 32
img_height = 180
img_width = 180# 划分数据
train_ds = tf.keras.utils.image_dataset_from_directory(data_dir,validation_split=0.2,subset="training",seed=123,image_size=(img_height, img_width),batch_size=batch_size)val_ds = tf.keras.utils.image_dataset_from_directory(data_dir,validation_split=0.2,subset="validation",seed=123,image_size=(img_height, img_width),batch_size=batch_size)class_names = train_ds.class_names
#print(class_names)# 处理数据
normalization_layer = layers.Rescaling(1./255)
train_ds = train_ds.map(lambda x, y: (normalization_layer(x), y))
val_ds = val_ds.map(lambda x, y: (normalization_layer(x), y))
num_classes = len(class_names)data_augmentation = keras.Sequential([layers.RandomFlip("horizontal",input_shape=(img_height,img_width,3)),layers.RandomRotation(0.1),layers.RandomZoom(0.1),]
)model = Sequential([data_augmentation,layers.Conv2D(16, 3, padding='same', activation='relu'),layers.MaxPooling2D(),layers.Conv2D(32, 3, padding='same', activation='relu'),layers.MaxPooling2D(),layers.Conv2D(64, 3, padding='same', activation='relu'),layers.MaxPooling2D(),layers.Dropout(0.2),layers.Flatten(),layers.Dense(128, activation='relu'),layers.Dense(num_classes, name="outputs")
])model.compile(optimizer='adam',loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),metrics=['accuracy'])model.summary()# 训练模型
epochs=15
history = model.fit(train_ds,validation_data=val_ds,epochs=epochs,
)# 测试模型
#sunflower_url = "https://storage.googleapis.com/download.tensorflow.org/example_images/592px-Red_sunflower.jpg"
#sunflower_path = tf.keras.utils.get_file('Red_sunflower', origin=sunflower_url)
sunflower_path = './test_180.jpg'img = tf.keras.utils.load_img(sunflower_path, target_size=(img_height, img_width)
)
img_array = tf.keras.utils.img_to_array(img)
img_array = tf.expand_dims(img_array, 0) # Create a batchpredictions = model.predict(img_array)
score = tf.nn.softmax(predictions[0])print("This image most likely belongs to {} with a {:.2f} percent confidence.".format(class_names[np.argmax(score)], 100 * np.max(score))
)# Convert the model.
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()# Save the model.
with open('model.tflite', 'wb') as f:f.write(tflite_model)

代码有点需要注意，代码屏蔽了下载的功能，所以需要预先下载数据集，如果没有下载数据集，就需要把下载的代码开启。

#dataset_url = "https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz"
#data_dir = tf.keras.utils.get_file('flower_photos', origin=dataset_url, untar=True)

执行下面命令开始训练：

python tensorflow_classification.py

等待一会，会生成model.tflite模型文件。

四、RKNN模型转换

转换代码通过下面代码：

rknn_transfer.py

import numpy as np
import cv2
from rknn.api import RKNN
import tensorflow as tfimg_height = 180
img_width = 180
IMG_PATH = 'test.jpg'
class_names = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']if __name__ == '__main__':# Create RKNN object#rknn = RKNN(verbose='Debug')rknn = RKNN()# Pre-process configprint('--> Config model')rknn.config(mean_values=[0, 0, 0], std_values=[255, 255, 255], target_platform='rk3568')print('done')# Load modelprint('--> Loading model')ret = rknn.load_tflite(model='model.tflite')if ret != 0:print('Load model failed!')exit(ret)print('done')# Build modelprint('--> Building model')ret = rknn.build(do_quantization=False)#ret = rknn.build(do_quantization=True,dataset='./dataset.txt')if ret != 0:print('Build model failed!')exit(ret)print('done')# Export rknn modelprint('--> Export rknn model')ret = rknn.export_rknn('./model.rknn')if ret != 0:print('Export rknn model failed!')exit(ret)print('done')#Init runtime environment
print('--> Init runtime environment')
ret = rknn.init_runtime()
#    if ret != 0:
#        print('Init runtime environment failed!')
#        exit(ret)
print('done')img = cv2.imread(IMG_PATH)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img,(180,180))
img = np.expand_dims(img, 0)#print('--> Accuracy analysis')
#rknn.accuracy_analysis(inputs=['./test.jpg'])
#print('done')print('--> Running model')
outputs = rknn.inference(inputs=[img])
print(outputs)
outputs = tf.nn.softmax(outputs)
print(outputs)print("This image most likely belongs to {} with a {:.2f} percent confidence.".format(class_names[np.argmax(outputs)], 100 * np.max(outputs))
)
#print("图像预测是:", class_names[np.argmax(outputs)])
print('--> done')rknn.release()

运行后会生成RKNN模型

五、部署

把rknnlite_inference.py和图片，及模型model.rknn拷贝到开发板上，终端运行即可。

rknnlite_inference.py源码：

import numpy as np
import cv2
from rknnlite.api import RKNNLiteIMG_PATH = 'test.jpg'
RKNN_MODEL = 'model.rknn'
img_height = 180
img_width = 180
class_names = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']# Create RKNN object
rknn_lite = RKNNLite()# load RKNN model
print('--> Load RKNN model')
ret = rknn_lite.load_rknn(RKNN_MODEL)
if ret != 0:print('Load RKNN model failed')exit(ret)
print('done')# Init runtime environment
print('--> Init runtime environment')
ret = rknn_lite.init_runtime()
if ret != 0:print('Init runtime environment failed!')exit(ret)
print('done')# load image
img = cv2.imread(IMG_PATH)
img = cv2.resize(img,(180,180))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.expand_dims(img, 0)# runing model
print('--> Running model')
outputs = rknn_lite.inference(inputs=[img])
print("result: ", outputs)
print("This image most likely belongs to {}.".format(class_names[np.argmax(outputs)])
)rknn_lite.release()

终端中执行：python rknnlite_inference.py