|
|
#coding:utf-8
|
|
|
# 环境（均为CPU）
|
|
|
# paddlehub 1.8.0
|
|
|
# paddlepaddle 2.3.0
|
|
|
|
|
|
|
|
|
import argparse
|
|
|
import os
|
|
|
import ast
|
|
|
import paddle.fluid as fluid
|
|
|
import paddlehub as hub
|
|
|
import paddle
|
|
|
import numpy as np
|
|
|
import matplotlib
|
|
|
matplotlib.rcParams['font.family']='simHei' #显示汉字
|
|
|
import matplotlib.pyplot as plt
|
|
|
import matplotlib.image as mpimg
|
|
|
from paddlehub.dataset.base_cv_dataset import BaseCVDataset
|
|
|
|
|
|
# yapf: disable
|
|
|
parser = argparse.ArgumentParser(__doc__)
|
|
|
parser.add_argument("--num_epoch", type=int, default=1, help="Number of epoches for fine-tuning.")
|
|
|
parser.add_argument("--use_gpu", type=ast.literal_eval, default=False, help="Whether use GPU for fine-tuning.")
|
|
|
parser.add_argument("--checkpoint_dir", type=str, default="paddlehub_finetune_ckpt", help="Path to save log data.")
|
|
|
parser.add_argument("--batch_size", type=int, default=16, help="Total examples' number in batch for training.")
|
|
|
parser.add_argument("--module", type=str, default="resnet50", help="Module used as feature extractor.")
|
|
|
parser.add_argument("--dataset", type=str, default="flowers", help="Dataset to fine-tune.")
|
|
|
parser.add_argument("--use_data_parallel", type=ast.literal_eval, default=True, help="Whether use data parallel.")
|
|
|
# yapf: enable.
|
|
|
|
|
|
module_map = {
|
|
|
"resnet50": "resnet_v2_50_imagenet",
|
|
|
"resnet101": "resnet_v2_101_imagenet",
|
|
|
"resnet152": "resnet_v2_152_imagenet",
|
|
|
"mobilenet": "mobilenet_v2_imagenet",
|
|
|
"nasnet": "nasnet_imagenet",
|
|
|
"pnasnet": "pnasnet_imagenet"
|
|
|
}
|
|
|
|
|
|
class DemoDataset(BaseCVDataset):
|
|
|
def __init__(self):
|
|
|
# 数据集存放位置（换成自己电脑数据集的位置）
|
|
|
self.dataset_dir = "D:\clsGarbageCode\garbage"
|
|
|
super(DemoDataset, self).__init__(
|
|
|
base_path=self.dataset_dir,
|
|
|
train_list_file="train_list.txt",
|
|
|
validate_list_file="val_list.txt",
|
|
|
test_list_file="test_list.txt",
|
|
|
# predict_file="predict_list.txt",
|
|
|
label_list_file="label_list.txt"
|
|
|
)
|
|
|
|
|
|
def finetune(args):
|
|
|
# Load Paddlehub pretrained model
|
|
|
paddle.enable_static()
|
|
|
module = hub.Module(name=args.module)
|
|
|
input_dict, output_dict, program = module.context(trainable=True)
|
|
|
|
|
|
dataset = DemoDataset()
|
|
|
|
|
|
# Use ImageClassificationReader to read dataset
|
|
|
data_reader = hub.reader.ImageClassificationReader(
|
|
|
image_width=module.get_expected_image_width(),
|
|
|
image_height=module.get_expected_image_height(),
|
|
|
images_mean=module.get_pretrained_images_mean(),
|
|
|
images_std=module.get_pretrained_images_std(),
|
|
|
dataset=dataset)
|
|
|
|
|
|
feature_map = output_dict["feature_map"]
|
|
|
|
|
|
# Setup feed list for data feeder
|
|
|
feed_list = [input_dict["image"].name]
|
|
|
|
|
|
# Setup RunConfig for PaddleHub Fine-tune API
|
|
|
config = hub.RunConfig(
|
|
|
use_data_parallel=args.use_data_parallel,
|
|
|
use_cuda=False,
|
|
|
# use_cuda=args.use_gpu,
|
|
|
num_epoch=args.num_epoch,
|
|
|
batch_size=args.batch_size,
|
|
|
checkpoint_dir=args.checkpoint_dir,
|
|
|
strategy=hub.finetune.strategy.DefaultFinetuneStrategy())
|
|
|
|
|
|
# Define a image classification task by PaddleHub Fine-tune API
|
|
|
task = hub.ImageClassifierTask(
|
|
|
data_reader=data_reader,
|
|
|
feed_list=feed_list,
|
|
|
feature=feature_map,
|
|
|
num_classes=dataset.num_labels,
|
|
|
config=config)
|
|
|
|
|
|
#Fine-tune by PaddleHub's API
|
|
|
#task.finetune_and_eval()
|
|
|
return task
|
|
|
|
|
|
#暂时注释的代码
|
|
|
'''def testShow(task):
|
|
|
# "任选一张图片路径/xx.jpg"
|
|
|
data = ["/media/lsz/lll/DLcode/paddle/flower_photos/dandelion/7355522_b66e5d3078_m.jpg"]
|
|
|
dataset = DemoDataset()
|
|
|
label_map = dataset.label_dict()
|
|
|
index = 0
|
|
|
# get classification result
|
|
|
run_states = task.predict(data=data)
|
|
|
results = [run_state.run_results for run_state in run_states]
|
|
|
for batch_result in results:
|
|
|
# get predict index
|
|
|
batch_result = np.argmax(batch_result, axis=2)[0]
|
|
|
for result in batch_result:
|
|
|
index += 1
|
|
|
result = label_map[result]
|
|
|
print("input %i is %s, and the predict result is ( %s )" %
|
|
|
(index, data[index - 1], result))
|
|
|
d=plt.imread("/media/lsz/lll/DLcode/paddle/flower_photos/dandelion/7355522_b66e5d3078_m.jpg")
|
|
|
plt.imshow(d)'''
|
|
|
|
|
|
def testShow(task):
|
|
|
# "任选一张图片路径/xx.jpg"
|
|
|
data = ["D:\clsGarbageCode\garbage\show\images.jpg"]#你要判断的垃圾类型
|
|
|
#print('图片类型是',end='')
|
|
|
#print(task.predict(data=data,return_result=True))
|
|
|
matplotlib.use('TkAgg')
|
|
|
d=mpimg.imread(data[0])
|
|
|
#plt.ylabel('垃圾类型为')
|
|
|
plt.title(task.predict(data=data,return_result=True))
|
|
|
plt.imshow(d)
|
|
|
plt.show()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
args = parser.parse_args()
|
|
|
if not args.module in module_map:
|
|
|
hub.logger.error("module should in %s" % module_map.keys())
|
|
|
exit(1)
|
|
|
args.module = module_map[args.module]
|
|
|
|
|
|
task = finetune(args) # 微调模型
|
|
|
testShow(task) # 传入微调模型，使用任意图片进行测试
|
|
|
|
