参考:
https://developers.google.com/mediapipe/solutions/customization/image_classifier
https://colab.research.google.com/github/googlesamples/mediapipe/blob/main/examples/customization/image_classifier.ipynb#scrollTo=plvO-YmcQn5g
安装:
pip install mediapipe-model-maker -i http://mirrors.aliyun.com/pypi/simple --trusted-host mirrors.aliyun.com --use-pep517
版本错误情况
1)RuntimeError: File loading is not yet supported on Windows
其中mediapipe版本要大于等于0.10.0;下图中的要升级;不然后续用python 加载文件会报:
2)ImportError: cannot import name ‘array_record_module’ from ‘array_record.python’ ;参考:https://blog.csdn.net/LQ_001/article/details/130991571;原因:包依赖关系出现问题,原来版本 tensorflow-datasets==4.9.0
pip install tensorflow-datasets==4.8.3
1、训练代码
import os
import tensorflow as tf
assert tf.__version__.startswith('2')
from mediapipe_model_maker import image_classifier
import matplotlib.pyplot as plt
image_path = os.path.join(os.path.dirname(r"C:\Users\loong\Downloads\mediapipe\flower_photos\flower_photos"), 'flower_photos') ## down data :https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz
#Review data
labels = []
for i in os.listdir(image_path):
if os.path.isdir(os.path.join(image_path, i)):
labels.append(i)
print(labels)
##plt
NUM_EXAMPLES = 5
for label in labels:
label_dir = os.path.join(image_path, label)
example_filenames = os.listdir(label_dir)[:NUM_EXAMPLES]
fig, axs = plt.subplots(1, NUM_EXAMPLES, figsize=(10,2))
for i in range(NUM_EXAMPLES):
axs[i].imshow(plt.imread(os.path.join(label_dir, example_filenames[i])))
axs[i].get_xaxis().set_visible(False)
axs[i].get_yaxis().set_visible(False)
fig.suptitle(f'Showing {NUM_EXAMPLES} examples for {label}')
plt.show()
#Create dataset;训练集、测试集
data = image_classifier.Dataset.from_folder(image_path)
train_data, remaining_data = data.split(0.8)
test_data, validation_data = remaining_data.split(0.5)
## retrain model 训练模型
spec = image_classifier.SupportedModels.MOBILENET_V2 ##有几个预训练模型,需要联网下载
hparams = image_classifier.HParams(export_dir="exported_model") ##指定模型保存位置
options = image_classifier.ImageClassifierOptions(supported_model=spec, hparams=hparams)
model = image_classifier.ImageClassifier.create(
train_data = train_data,
validation_data = validation_data,
options=options,
)
## 验证模型
loss, acc = model.evaluate(test_data)
print(f'Test loss:{loss}, Test accuracy:{acc}')
##保存模型
model.export_model()
默认训练是10epcos
##模型压缩
from mediapipe_model_maker import quantization
quantization_config = quantization.QuantizationConfig.for_int8(train_data)
model.export_model(model_name="model_int8.tflite", quantization_config=quantization_config)
从8M缩小到3M左右
2、加载推理
参考:https://blog.csdn.net/weixin_42357472/article/details/131322076
import mediapipe as mp
BaseOptions = mp.tasks.BaseOptions
ImageClassifier = mp.tasks.vision.ImageClassifier
ImageClassifierOptions = mp.tasks.vision.ImageClassifierOptions
VisionRunningMode = mp.tasks.vision.RunningMode
options = ImageClassifierOptions(
base_options=BaseOptions(model_asset_path=r"C:\User**ediapipe\model.tflite"),
max_results=5,
running_mode=VisionRunningMode.IMAGE) ##加载模型
classifier = ImageClassifier.create_from_options(options)
# Load the input image from an image file.
mp_image = mp.Image.create_from_file(r"C:\Users\loong\Downloads\sun2.jpg")
# Perform image classification on the provided single image.
classification_result = classifier.classify(mp_image)
classification_result