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