157 lines
7.1 KiB
Python
157 lines
7.1 KiB
Python
import os
|
|
import cv2
|
|
import torch
|
|
import numpy as np
|
|
import torch.nn as nn
|
|
from torch.nn import init
|
|
import torch.nn.functional as F
|
|
from torch.optim import lr_scheduler
|
|
|
|
|
|
def save_imgs(imgs, names, path):
|
|
if not os.path.exists(path):
|
|
os.mkdir(path)
|
|
for img, name in zip(imgs, names):
|
|
img = tensor2img(img)
|
|
img = np.array(img)
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
|
cv2.imwrite(os.path.join(path, name + '.jpg'), img)
|
|
|
|
|
|
def tensor2img(img):
|
|
img = img[0].cpu().float().numpy()
|
|
if img.shape[0] == 1:
|
|
img = np.tile(img, (3, 1, 1))
|
|
img = (np.transpose(img, (1, 2, 0)) + 1) / 2.0 * 255.0
|
|
return img.astype(np.uint8)
|
|
|
|
####################################################################
|
|
# ------------------------- Basic Functions -------------------------
|
|
####################################################################
|
|
|
|
def get_scheduler(optimizer, opts, cur_ep=-1):
|
|
if opts.lr_policy == 'lambda':
|
|
def lambda_rule(ep):
|
|
lr_l = 1.0 - max(0, ep - opts.n_ep_decay) / float(opts.n_ep - opts.n_ep_decay + 1)
|
|
return lr_l
|
|
|
|
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule, last_epoch=cur_ep)
|
|
elif opts.lr_policy == 'step':
|
|
scheduler = lr_scheduler.StepLR(optimizer, step_size=opts.n_ep_decay, gamma=0.1, last_epoch=cur_ep)
|
|
elif opts.lr_policy == 'plateau':
|
|
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
|
|
else:
|
|
return NotImplementedError('learning rate policy [%s] is not implemented', opts.lr_policy)
|
|
return scheduler
|
|
|
|
def init_weights(net, init_type, gain):
|
|
def init_func(m):
|
|
classname = m.__class__.__name__
|
|
if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
|
|
if init_type == 'normal':
|
|
init.normal_(m.weight.data, 0.0, gain)
|
|
elif init_type == 'xavier':
|
|
init.xavier_normal_(m.weight.data, gain=gain)
|
|
elif init_type == 'kaiming':
|
|
init.kaiming_normal_(m.weight.data, a=0, mode='fainplanes')
|
|
elif init_type == 'orthogonal':
|
|
init.orthogonal_(m.weight.data, gain=gain)
|
|
else:
|
|
raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
|
|
if hasattr(m, 'bias') and m.bias is not None:
|
|
init.constant_(m.bias.data, 0.0)
|
|
elif classname.find('BatchNorm2d') != -1:
|
|
init.normal_(m.weight.data, 1.0, gain)
|
|
init.constant_(m.bias.data, 0.0)
|
|
|
|
print('initialize network with %s' % init_type)
|
|
net.apply(init_func)
|
|
|
|
def init_net(net, gpu, init_type='normal', gain=0.02):
|
|
assert (torch.cuda.is_available())
|
|
net.to(gpu)
|
|
init_weights(net, init_type, gain)
|
|
return net
|
|
|
|
####################################################################
|
|
# ------------------------- VGG -------------------------
|
|
####################################################################
|
|
class VGG19_feature(nn.Module):
|
|
'''
|
|
NOTE: there is no need to pre-process the input
|
|
input tensor should range in [0,1]
|
|
'''
|
|
|
|
def __init__(self, pool='max', vgg_normal_correct=True, ic=3):
|
|
super(VGG19_feature, self).__init__()
|
|
self.vgg_normal_correct = vgg_normal_correct
|
|
|
|
self.conv1_1 = nn.Conv2d(ic, 64, kernel_size=3, padding=1)
|
|
self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
|
|
self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
|
|
self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
|
|
self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
|
|
self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
|
|
self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
|
|
self.conv3_4 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
|
|
self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
|
|
self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
|
|
self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
|
|
self.conv4_4 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
|
|
self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
|
|
self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
|
|
self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
|
|
self.conv5_4 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
|
|
if pool == 'max':
|
|
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
|
|
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
|
|
self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
|
|
self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)
|
|
self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2)
|
|
elif pool == 'avg':
|
|
self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)
|
|
self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)
|
|
self.pool3 = nn.AvgPool2d(kernel_size=2, stride=2)
|
|
self.pool4 = nn.AvgPool2d(kernel_size=2, stride=2)
|
|
self.pool5 = nn.AvgPool2d(kernel_size=2, stride=2)
|
|
|
|
def forward(self, x, out_keys, preprocess=True):
|
|
'''
|
|
NOTE: input tensor should range in [0,1]
|
|
'''
|
|
out = {}
|
|
if preprocess:
|
|
x = vgg_preprocess(x, vgg_normal_correct=self.vgg_normal_correct)
|
|
out['r11'] = F.relu(self.conv1_1(x))
|
|
out['r12'] = F.relu(self.conv1_2(out['r11']))
|
|
out['p1'] = self.pool1(out['r12'])
|
|
out['r21'] = F.relu(self.conv2_1(out['p1']))
|
|
out['r22'] = F.relu(self.conv2_2(out['r21']))
|
|
out['p2'] = self.pool2(out['r22'])
|
|
out['r31'] = F.relu(self.conv3_1(out['p2']))
|
|
out['r32'] = F.relu(self.conv3_2(out['r31']))
|
|
out['r33'] = F.relu(self.conv3_3(out['r32']))
|
|
out['r34'] = F.relu(self.conv3_4(out['r33']))
|
|
out['p3'] = self.pool3(out['r34'])
|
|
out['r41'] = F.relu(self.conv4_1(out['p3']))
|
|
out['r42'] = F.relu(self.conv4_2(out['r41']))
|
|
out['r43'] = F.relu(self.conv4_3(out['r42']))
|
|
out['r44'] = F.relu(self.conv4_4(out['r43']))
|
|
out['p4'] = self.pool4(out['r44'])
|
|
out['r51'] = F.relu(self.conv5_1(out['p4']))
|
|
out['r52'] = F.relu(self.conv5_2(out['r51']))
|
|
out['r53'] = F.relu(self.conv5_3(out['r52']))
|
|
out['r54'] = F.relu(self.conv5_4(out['r53']))
|
|
out['p5'] = self.pool5(out['r54'])
|
|
return [out[key] for key in out_keys]
|
|
|
|
def vgg_preprocess(tensor, vgg_normal_correct=True):
|
|
if vgg_normal_correct:
|
|
tensor = (tensor + 1) / 2
|
|
# input is RGB tensor which ranges in [0,1]
|
|
# output is BGR tensor which ranges in [0,255]
|
|
tensor_bgr = torch.cat((tensor[:, 2:3, :, :], tensor[:, 1:2, :, :], tensor[:, 0:1, :, :]), dim=1)
|
|
# tensor_bgr = tensor[:, [2, 1, 0], ...]
|
|
tensor_bgr_ml = tensor_bgr - torch.Tensor([0.40760392, 0.45795686, 0.48501961]).type_as(tensor_bgr).view(1, 3, 1, 1)
|
|
tensor_rst = tensor_bgr_ml * 255
|
|
return tensor_rst |