def get_models(args):
sigmoid_flag = 1
if args.gan == 'lsgan':
sigmoid_flag = 0
if args.model == 'scribbler':
netG = scribbler.Scribbler(5, 3, 32)
elif args.model == 'texturegan':
netG = texturegan.TextureGAN(5, 3, 32)
elif args.model == 'pix2pix':
netG = define_G(5, 3, 32)
elif args.model == 'scribbler_dilate_128':
netG = scribbler_dilate_128.ScribblerDilate128(5, 3, 32)
else:
print(args.model + ' not support. Using Scribbler model')
netG = scribbler.Scribbler(5, 3, 32)
if args.color_space == 'lab':
netD = discriminator.Discriminator(1, 32, sigmoid_flag)
netD_local = discriminator.LocalDiscriminator(2, 32, sigmoid_flag)
elif args.color_space == 'rgb':
netD = discriminator.Discriminator(3, 32, sigmoid_flag)
if args.load == -1:
netG.apply(weights_init)
else:
load_network(netG, 'G', args.load_epoch, args.load, args)
if args.load_D == -1:
netD.apply(weights_init)
else:
load_network(netD, 'D', args.load_epoch, args.load_D, args)
load_network(netD_local, 'D_local', args.load_epoch, args.load_D, args)
return netG, netD, netD_local
def initialize(self):
if os.path.isabs(self.net_path):
net_path_full = self.net_path
else:
net_path_full = os.path.join(env_settings().network_path,
self.net_path)
self.net, _ = load_network(net_path_full, backbone_pretrained=False)
if self.use_gpu:
self.net.cuda()
self.net.eval()
# self.iou_predictor = self.net.bb_regressor # original
self.rgb_bb_regressor = self.net.rgb_bb_regressor
self.t_bb_regressor = self.net.t_bb_regressor
self.iou_guess = self.net.iou_guess
self.layer_stride = {
'conv1': 2,
'layer1': 4,
'layer2': 8,
'layer3': 16,
'layer4': 32,
'classification': 16,
'fc': None
}
self.layer_dim = {
'conv1': 64,
'layer1': 64,
'layer2': 128,
'layer3': 256,
'layer4': 512,
'classification': 256,
'fc': None
}
self.iounet_feature_layers = self.net.bb_regressor_layer
if isinstance(self.pool_stride, int) and self.pool_stride == 1:
self.pool_stride = [1] * len(self.output_layers)
self.feature_layers = sorted(
list(set(self.output_layers + self.iounet_feature_layers)))
self.mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, -1, 1, 1)
self.std = torch.Tensor([0.229, 0.224, 0.225]).view(1, -1, 1, 1)