(五)将YOLOv5 PyTorch模型权重转换为TensorFlow Lite格式_yolov5 tensorflow

import tensorflow as tf
print(tf.__version__)
!PYTHONPATH=. python3 /content/yolov5/models/tf.py --weight
/content/yolov5/runs/train/exp2/weights/best.pt --cfg models/yolov5s.yaml --img 416 --source /content/FaceMaskDataset/images/train

传递权重的文件路径（configuration.yaml文件）时，请指示模型接受的图像尺寸和训练数据集的来源（最后一个参数是可选的）。该脚本将使用TensorFlow 2.3.1将.pt权重转换为TensorFlow格式，并且输出将保存在/content/yolov5/runs/train/exp/weights中。

# Uncomment all this if you want to follow the long path
#!pip install onnx>=1.7.0 # for ONNX export
#!pip install coremltools==4.0 # for CoreML export
#!python models/export.py --weights /content/yolov5/runs/train/exp2/weights/best.pt --img 416 --batch 1 # export at 640x640 with batch size 1
 
#!pip install tensorflow-addons
#!pip install onnx-tf
#!pip install tensorflow==2.3.0
 
#import tensorflow_addons as tfa
#import onnx
#from onnx_tf.backend import prepare
#import tensorflow as tf
#print(tf.__version__)
 
#base_model = onnx.load('/content/yolov5/runs/train/exp2/weights/best.onnx')
#to_tf = prepare(base_model)
#to_tf.export_graph("/content/yolov5/runs/train/exp2/weights/customyolov5")
 
#converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model('/content/yolov5/runs/train/exp2/weights/customyolov5')
#tflite_model = converter.convert() #just FYI: this step could go wrong and your notebook instance could crash.

在Google Colab上测试TFLite权重

!python detect.py --weight /content/yolov5/runs/train/exp2/weights/best-fp16.tflite --img 416 --source ../test_images

如果一切顺利，结果将类似于以下内容：

#Install the appropriate tensorflow version
!pip install tensorflow==2.3.1
import torch #It's required by YOLO
import tensorflow as tf
!git clone https://github.com/zldrobit/yolov5 #comment this if this is not your first run
%cd yolov5
!git checkout tf-android
print(tf.__version__)
 
#Install all dependencies indicated at requirements.txt file
#Prompt if any GPU is available
!pip install -r requirements.txt
print('All set. Using PyTorch version %s with %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

在继续之前，请记住修改detect.py文件中第157行的名称列表，并将所有下载的权重复制到YOLOv5文件夹内的/weights文件夹中。

!python detect.py --weights weights/best.pt --img 416 --conf 0.4 --source 0 --classes 0 1

它将在单独的窗口中启动网络摄像头，识别您的脸部，并检测您是否戴着口罩。这是您应该期望的：

pip3 https://github.com/google-coral/pycoral/releases/download/release-frogfish/tflite_runtime-2.5.0-cp36-cp36m-linux_x86_64.whl

如果您使用的是其他操作系统，建议您选择最适合的版本。

!python detect4pi.py --weights weights/best-fp16.tflite --img 416 --conf 0.45 --source 0 --classes 0 1

下一步