Классификация с предварительно обученной моделью pytorch vgg16 и ее классами - PullRequest
Новая
       188

Классификация с предварительно обученной моделью pytorch vgg16 и ее классами

0 голосов
/ 20 июня 2020

Я написал модель классификации изображений vgg с предварительно обученной моделью vgg16 pytorch. 

import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image
import urllib
from skimage.transform import resize
from skimage import io
import yaml

# Downloading imagenet 1000 classes list
file = urllib. request. urlopen("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt")
classes = ''
for f in file:
  classes = classes +  f.decode("utf-8")
classes = yaml.load(classes)

# Downloading pretrained vgg16 model
model = torch.hub.load('pytorch/vision:v0.6.0', 'vgg16', pretrained=True)

print(model)

for param in model.parameters():
    param.requires_grad = False


url, filename = ("https://raw.githubusercontent.com/pytorch/hub/master/dog.jpg", "dog.jpg")

image=io.imread(url)

plt.imshow(image)
plt.show()

# resize to 224x224x3
img = resize(image,(224,224,3))

plt.imshow(img)
plt.show()
# Normalizing input for vgg16
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
img1 = mean*img+std
img1 = np.clip(img1,0,1)

img1 = torch.from_numpy(img1).unsqueeze(0)
img1 = img1.permute(0,3,2,1) # batch_size x channels x height x width

model.eval()
pred = model(img1.float())
print(classes[torch.argmax(pred).numpy().tolist()])

Код работает нормально, но выводит неправильные классы. Я не уверен, где я ошибся, но если мне нужно предположить, это может быть список классов imag enet yaml или нормализующее входное изображение. Кто-нибудь может сказать мне, где я делаю ошибки?

1 Ответ

1 голос
/ 20 июня 2020

Есть некоторые проблемы с предварительной обработкой изображения. Во-первых, нормализация рассчитывается как (value - mean) / std), а не value * mean + std. Во-вторых, значения не должны быть обрезаны до [0, 1], нормализация намеренно смещает значения от [0, 1]. Во-вторых, изображение в виде массива NumPy имеет форму [высота, ширина, 3] , когда вы переставляете размеры, вы меняете местами высоту и ширину, создавая тензор с формой [batch_size, каналы , width, height] . 

img = resize(image,(224,224,3))


# Normalizing input for vgg16
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
img1 = (img1 - mean) / std

img1 = torch.from_numpy(img1).unsqueeze(0)
img1 = img1.permute(0, 3, 1, 2) # batch_size x channels x height x width

Вместо того, чтобы делать это вручную, вы можете использовать torchvision.transforms. 

from torchvision import transforms

preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

img = resize(image,(224,224,3))
img1 = preprocess(img)
img1 = img1.unsqueeze(0)

Если вы используете PIL для загрузки изображений вы также можете изменить размер изображений, добавив transforms.Resize((224, 224)) в конвейер предварительной обработки, или вы даже можете добавить transforms.ToPILImage(), чтобы сначала преобразовать изображение в изображение PIL (transforms.Resize требуется образ PIL).

...