|
from numpy.lib.arraysetops import isin |
|
import torchvision.models as models |
|
import torch.nn as nn |
|
import torch |
|
|
|
|
|
|
|
class Vgg19(nn.Module): |
|
def __init__(self): |
|
super(Vgg19, self).__init__() |
|
self.vgg19 = self.get_vgg19().eval() |
|
vgg_mean = torch.tensor([0.485, 0.456, 0.406]).float() |
|
vgg_std = torch.tensor([0.229, 0.224, 0.225]).float() |
|
self.mean = vgg_mean.view(-1, 1 ,1) |
|
self.std = vgg_std.view(-1, 1, 1) |
|
|
|
def to(self, device): |
|
new_self = super(Vgg19, self).to(device) |
|
new_self.mean = new_self.mean.to(device) |
|
new_self.std = new_self.std.to(device) |
|
return new_self |
|
|
|
def forward(self, x): |
|
return self.vgg19(self.normalize_vgg(x)) |
|
|
|
@staticmethod |
|
def get_vgg19(last_layer='conv4_4'): |
|
vgg = models.vgg19(weights=models.VGG19_Weights.IMAGENET1K_V1).features |
|
model_list = [] |
|
|
|
i = 0 |
|
j = 1 |
|
for layer in vgg.children(): |
|
if isinstance(layer, nn.MaxPool2d): |
|
i = 0 |
|
j += 1 |
|
|
|
elif isinstance(layer, nn.Conv2d): |
|
i += 1 |
|
|
|
name = f'conv{j}_{i}' |
|
|
|
if name == last_layer: |
|
model_list.append(layer) |
|
break |
|
|
|
model_list.append(layer) |
|
|
|
|
|
model = nn.Sequential(*model_list) |
|
return model |
|
|
|
|
|
def normalize_vgg(self, image): |
|
''' |
|
Expect input in range -1 1 |
|
''' |
|
image = (image + 1.0) / 2.0 |
|
return (image - self.mean) / self.std |
|
|
|
|
|
if __name__ == '__main__': |
|
from PIL import Image |
|
import numpy as np |
|
from utils.image_processing import normalize_input |
|
|
|
image = Image.open("example/10.jpg") |
|
image = image.resize((224, 224)) |
|
np_img = np.array(image).astype('float32') |
|
np_img = normalize_input(np_img) |
|
|
|
img = torch.from_numpy(np_img) |
|
img = img.permute(2, 0, 1) |
|
img = img.unsqueeze(0) |
|
|
|
vgg = Vgg19() |
|
|
|
feat = vgg(img) |
|
|
|
print(feat.shape) |
