这是我的推理脚本
import torch
import torchvision.models as models
from torchvision import transforms
from torch.autograd import Variable
from PIL import Image
import torch.nn as nn
MODEL_PATH = "/content/deepfashion-dataset/models/atr-recognition-stage-2-resnet34.pkl"
DATA_PATH = "/content/deepfashion-dataset/"
CLASSES_PATH = "/content/deepfashion-dataset/clothes_categories/attribute-classes.txt"
class ClassificationModel():
def __init__(self):
return
def load(self, model_path, labels_path, eval=False):
self.model = torch.load(model_path)
self.model = nn.Sequential(self.model)
self.labels = open(labels_path, 'r').read().splitlines()
if eval:
print(model.eval())
return
def predict(self, image_path):
device = torch.device("cpu")
img = Image.open(image_path)
test_transforms = transforms.Compose([transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
image_tensor = test_transforms(img).float()
image_tensor = image_tensor.unsqueeze_(0)
inp = Variable(image_tensor)
inp = inp.to(device)
output = self.model(inp)
index = output.data.cpu().numpy().argmax()
return self.labels[index]
learner = ClassificationModel()
learner.load(MODEL_PATH, CLASSES_PATH)
print(learner.predict(DATA_PATH+"img-lk/IMG_20210530_152109_062.jpg"))
它给出此图像的输出是:
leather
在训练这样的模型后我对其进行测试时:
def predict_attribute(model, path, display_img=True):
predicted = model.predict(path)
if display_img:
size = 244,244
img=Image.open(path)
img.thumbnail(size,Image.ANTIALIAS)
display(img)
return predicted[0]
image_path = PATH + 'img-lk/IMG_20210530_152109_062.jpg'
predict_attribute(learn, image_path)
输出为:
(#2) ['faux-leather','leather']
当时给出了两个输出属性。另外,我的属性也保存在文本file
.中。
面对同一问题 请告诉您是否找到解决方案