134 lines
4.4 KiB
Python
Executable File
134 lines
4.4 KiB
Python
Executable File
import collections
|
|
import os
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torchvision.models as models
|
|
import torchvision.transforms as transforms
|
|
from PIL import Image
|
|
|
|
|
|
def preprocess(img, scale_size=None):
|
|
"""PILimg: RGB: HxWxC"""
|
|
if scale_size is not None:
|
|
scale_transforms = MaxScale(scale_size)
|
|
img = scale_transforms(img)
|
|
|
|
prep_transforms = transforms.Compose(
|
|
[
|
|
transforms.ToTensor(),
|
|
transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])]), # turn to BGR
|
|
transforms.Normalize(mean=[0.40760392, 0.45795686, 0.48501961], std=[1, 1, 1]), # subtract imagenet mean
|
|
transforms.Lambda(lambda x: x.mul_(255)),
|
|
]
|
|
)
|
|
img = prep_transforms(img)
|
|
return img
|
|
|
|
|
|
def deprocess(img):
|
|
post_transforms_a = transforms.Compose(
|
|
[
|
|
transforms.Lambda(lambda x: x.mul_(1.0 / 255)),
|
|
transforms.Normalize(mean=[-0.40760392, -0.45795686, -0.48501961], std=[1, 1, 1]), # add imagenet mean
|
|
transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])]), # turn to RGB
|
|
]
|
|
)
|
|
post_transforms_b = transforms.Compose([transforms.ToPILImage()])
|
|
img = post_transforms_a(img)
|
|
img[img > 1] = 1
|
|
img[img < 0] = 0
|
|
img = post_transforms_b(img)
|
|
return img
|
|
|
|
|
|
def get_renamed_vgg():
|
|
cache_file = "data/vgg19_conv.pth"
|
|
vgg = models.vgg19().features
|
|
renamed_vgg = nn.Sequential()
|
|
part_idx, layer_idx = 1, 1
|
|
for layer in list(vgg):
|
|
if isinstance(layer, nn.Conv2d):
|
|
name = "conv{}_{}".format(part_idx, layer_idx)
|
|
renamed_vgg.add_module(name, layer)
|
|
elif isinstance(layer, nn.ReLU):
|
|
name = "relu{}_{}".format(part_idx, layer_idx)
|
|
renamed_vgg.add_module(name, layer)
|
|
layer_idx += 1
|
|
elif isinstance(layer, nn.MaxPool2d):
|
|
name = "pool{}".format(part_idx)
|
|
renamed_vgg.add_module(name, layer)
|
|
part_idx += 1
|
|
layer_idx = 1
|
|
renamed_vgg.load_state_dict(torch.load(cache_file))
|
|
|
|
return renamed_vgg
|
|
|
|
|
|
def get_renamed_vgg_johnson():
|
|
cache_file = "data/vgg19-d01eb7cb.pth"
|
|
vgg = models.vgg19()
|
|
model_dict = vgg.state_dict()
|
|
pretrained_dict = torch.load(cache_file)
|
|
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
|
|
model_dict.update(pretrained_dict)
|
|
vgg.load_state_dict(model_dict)
|
|
vgg = vgg.features
|
|
renamed_vgg = nn.Sequential()
|
|
part_idx, layer_idx = 1, 1
|
|
for layer in list(vgg):
|
|
if isinstance(layer, nn.Conv2d):
|
|
name = "conv{}_{}".format(part_idx, layer_idx)
|
|
renamed_vgg.add_module(name, layer)
|
|
elif isinstance(layer, nn.ReLU):
|
|
name = "relu{}_{}".format(part_idx, layer_idx)
|
|
renamed_vgg.add_module(name, layer)
|
|
layer_idx += 1
|
|
elif isinstance(layer, nn.MaxPool2d):
|
|
name = "pool{}".format(part_idx)
|
|
renamed_vgg.add_module(name, layer)
|
|
part_idx += 1
|
|
layer_idx = 1
|
|
|
|
return renamed_vgg
|
|
|
|
|
|
class MaxScale(object):
|
|
"""Rescale the input PIL.Image to the given size.
|
|
|
|
Args:
|
|
size (sequence or int): Desired output size. If size is a sequence like
|
|
(w, h), output size will be matched to this. If size is an int,
|
|
smaller edge of the image will be matched to this number.
|
|
i.e, if height > width, then image will be rescaled to
|
|
(size * height / width, size)
|
|
interpolation (int, optional): Desired interpolation. Default is
|
|
``PIL.Image.BILINEAR``
|
|
"""
|
|
|
|
def __init__(self, size, interpolation=Image.BILINEAR):
|
|
assert isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2)
|
|
self.size = size
|
|
self.interpolation = interpolation
|
|
|
|
def __call__(self, img):
|
|
"""
|
|
Args:
|
|
img (PIL.Image): Image to be scaled.
|
|
|
|
Returns:
|
|
PIL.Image: Rescaled image.
|
|
"""
|
|
if not isinstance(self.size, int):
|
|
return img.resize(self.size, self.interpolation)
|
|
w, h = img.size
|
|
if (w <= h == self.size) or (h <= w == self.size):
|
|
return img
|
|
if w < h:
|
|
oh = self.size
|
|
ow = int(self.size * w / h)
|
|
else:
|
|
ow = self.size
|
|
oh = int(self.size * h / w)
|
|
return img.resize((ow, oh), self.interpolation)
|