def predict(model):
# 读入模型
model = load_checkpoint(model)
print('..... Finished loading model! ......')
##将模型放置在gpu上运行
if torch.cuda.is_available():
model.cuda()
pred_list, _id = [], []
for i in tqdm(range(len(imgs))):
img_path = imgs[i].strip()
# print(img_path)
_id.append(int(os.path.basename(img_path).split('.')[0]))
img = Image.open(img_path).convert('RGB')
# print(type(img))
img = get_test_transform(size=cfg.INPUT_SIZE)(img).unsqueeze(0)
if torch.cuda.is_available():
img = img.cuda()
with torch.no_grad():
out = model(img)
prediction = torch.argmax(out, dim=1).cpu().item()
pred_list.append(prediction)
return _id, pred_list
def save_feature(model, feature_path, label_path):
'''
提取特征，保存为pkl文件
'''
model = load_checkpoint(model)
# print(model)
print('..... Finished loading model! ......')
##将模型放置在gpu上运行
if torch.cuda.is_available():
model.cuda()
## 特征的维度需要自己根据特定的模型调整，我这里采用的是哪一个我也忘了
nb_features = NB_features
features = np.empty((len(imgs), nb_features))
labels = []
for i in tqdm(range(len(imgs))):
img_path = imgs[i].strip().split(' ')[0]
label = imgs[i].strip().split(' ')[1]
# print(img_path)
img = Image.open(img_path).convert('RGB')
# print(type(img))
img = get_test_transform(size=cfg.INPUT_SIZE)(img).unsqueeze(0)
if torch.cuda.is_available():
img = img.cuda()
with torch.no_grad():
out = model.extract_features(img)
# print(out.size())
out2 = nn.AdaptiveAvgPool2d(1)(out)
feature = out2.view(out.size(1), -1).squeeze(1)
# print(out3.size())
# print(out2.size())
features[i, :] = feature.cpu().numpy()
labels.append(label)
pickle.dump(features, open(feature_path, 'wb'))
pickle.dump(labels, open(label_path, 'wb'))
print('CNN features obtained and saved.')
def predict():
# initialize the data dictionary that will be returned from the
# view
data = {"success": False}
# print(data)
# ensure an image was properly uploaded to our endpoint
if request.method == "POST":
# print("Hello")
if request.files.get("image"):
# print("world")
now = time.strftime("%Y-%m-%d-%H_%M_%S",
time.localtime(time.time()))
# read the image in PIL format
image = request.files["image"].read()
image = Image.open(io.BytesIO(image)).convert('RGB')
image.save(now + '.jpg')
# preprocess the image and prepare it for classification
img = get_test_transform(mean, std, input_size)(image).unsqueeze(0)
# classify the input image and then initialize the list
# of predictions to return to the client
out = model(img)
# print(out)
pred_label = torch.max(out, 1)[1].item()
# print(pred_label)
data["predictions"] = []
data["predictions"].append(label_id_name_dict[str(pred_label)])
# indicate that the request was a success
data["success"] = True
# print(data["success"])
# return the data dictionary as a JSON response
return jsonify(data)
def __init__(self, model_name, model_path):
self.model_name = model_name
self.model_path = model_path
self.model = make_model(args)
#self.model = models.__dict__['resnet50'](num_classes=54)
self.use_cuda = False
if torch.cuda.is_available():
print('Using GPU for inference')
self.use_cuda = True
self.model = torch.nn.DataParallel(self.model).cuda()
checkpoint = torch.load(self.model_path)
#self.model.load_state_dict(checkpoint['state_dict'])
self.model.load_state_dict(checkpoint['state_dict'])
else:
print('Using CPU for inference')
checkpoint = torch.load(self.model_path, map_location='cpu')
state_dict = OrderedDict()
# 训练脚本 main.py 中保存了'epoch', 'arch', 'state_dict', 'best_acc1', 'optimizer'五个key值，
# 其中'state_dict'对应的value才是模型的参数。
# 训练脚本 main.py 中创建模型时用了torch.nn.DataParallel，因此模型保存时的dict都会有‘module.’的前缀，
# 下面 tmp = key[7:] 这行代码的作用就是去掉‘module.’前缀
for key, value in checkpoint['state_dict'].items():
tmp = key[7:]
state_dict[tmp] = value
self.model.load_state_dict(state_dict)
self.model.eval()
#self.idx_to_class = checkpoint['idx_to_class']
#self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
# std=[0.229, 0.224, 0.225])
#self.transforms = transforms.Compose([
# transforms.Resize(256),
# transforms.CenterCrop(224),
# transforms.ToTensor(),
# self.normalize
#])
self.transforms = get_test_transform([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225],
224)
self.label_id_name_dict = \
{
"0": "工艺品/仿唐三彩",
"1": "工艺品/仿宋木叶盏",
"2": "工艺品/布贴绣",
"3": "工艺品/景泰蓝",
"4": "工艺品/木马勺脸谱",
"5": "工艺品/柳编",
"6": "工艺品/葡萄花鸟纹银香囊",
"7": "工艺品/西安剪纸",
"8": "工艺品/陕历博唐妞系列",
"9": "景点/关中书院",
"10": "景点/兵马俑",
"11": "景点/南五台",
"12": "景点/大兴善寺",
"13": "景点/大观楼",
"14": "景点/大雁塔",
"15": "景点/小雁塔",
"16": "景点/未央宫城墙遗址",
"17": "景点/水陆庵壁塑",
"18": "景点/汉长安城遗址",
"19": "景点/西安城墙",
"20": "景点/钟楼",
"21": "景点/长安华严寺",
"22": "景点/阿房宫遗址",
"23": "民俗/唢呐",
"24": "民俗/皮影",
"25": "特产/临潼火晶柿子",
"26": "特产/山茱萸",
"27": "特产/玉器",
"28": "特产/阎良甜瓜",
"29": "特产/陕北红小豆",
"30": "特产/高陵冬枣",
"31": "美食/八宝玫瑰镜糕",
"32": "美食/凉皮",
"33": "美食/凉鱼",
"34": "美食/德懋恭水晶饼",
"35": "美食/搅团",
"36": "美食/枸杞炖银耳",
"37": "美食/柿子饼",
"38": "美食/浆水面",
"39": "美食/灌汤包",
"40": "美食/烧肘子",
"41": "美食/石子饼",
"42": "美食/神仙粉",
"43": "美食/粉汤羊血",
"44": "美食/羊肉泡馍",
"45": "美食/肉夹馍",
"46": "美食/荞面饸饹",
"47": "美食/菠菜面",
"48": "美食/蜂蜜凉粽子",
"49": "美食/蜜饯张口酥饺",
"50": "美食/西安油茶",
"51": "美食/贵妃鸡翅",
"52": "美食/醪糟",
"53": "美食/金线油塔"
}