def __init__(self, args):
self.args = args
Tensor = torch.cuda.FloatTensor if args.gpu_ids else torch.Tensor
use_sigmoid = args.no_lsgan
# Global discriminator
self.netD = networks.define_D(args.input_nc, args.ndf,
args.which_model_netD, args.n_layers_D,
args.norm, use_sigmoid, args.gpu_ids)
# Local discriminator
self.netD_local = networks.define_D(args.input_nc, args.ndf,
args.which_model_netD,
args.n_layers_D, args.norm,
use_sigmoid, args.gpu_ids)
# Generator
self.netG = TransformerNet(args.norm, args.affine_state)
self.gan_loss = networks.GANLoss(use_lsgan=not args.no_lsgan,
tensor=Tensor)
self.identity_criterion = torch.nn.L1Loss()
# Resume
if args.resume_netG != '':
self.netG.load_state_dict(torch.load(args.resume_netG))
if args.resume_netD != '':
self.netD.load_state_dict(torch.load(args.resume_netD))
if args.resume_netD_local != '':
self.netD_local.load_state_dict(torch.load(args.resume_netD_local))
# optimizer
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
self.optimizer_D_local = torch.optim.Adam(self.netD_local.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
if self.args.cuda:
self.netD = self.netD.cuda()
self.netG = self.netG.cuda()
def test(**kwargs):
"""
perform style transfer
"""
opt = Config()
for k_, v_ in kwargs.items():
setattr(opt, k_, v_)
device = t.device('cuda') if opt.use_gpu else t.device('cpu')
# preprocess input image
content_image = tv.datasets.folder.default_loader(opt.content_path)
content_transform = tv.transforms.Compose(
[tv.transforms.ToTensor(),
tv.transforms.Lambda(lambda x: x.mul(255))])
content_image = content_transform(content_image)
content_image = content_image.unsqueeze(0).to(device).detach()
# model setup
style_model = TransformerNet().eval()
style_model.load_state_dict(
t.load(opt.model_path, map_location=lambda _s, _: _s))
style_model.to(device)
# style transfer and save output
output = style_model(content_image)
output_data = output.cpu().data[0]
tv.utils.save_image(((output_data / 255)).clamp(min=0, max=1),
opt.result_path)
def stylize(args):
img_list=os.listdir(args.content_image)
epoch_name=os.path.basename(args.model).split('.')[0]
experiment_name=os.path.dirname(args.output_image)
if not os.path.exists(experiment_name):
os.system('mkdir {}'.format(experiment_name))
if not os.path.exists(args.output_image):
os.system('mkdir {}'.format(args.output_image))
for img in img_list:
if is_image_file(img):
content_image = utils.load_image(os.path.join(args.content_image,img), scale=args.content_scale)
content_image=content_image.convert('RGB')
content_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))
])
content_image = content_transform(content_image)
content_image = content_image.unsqueeze(0)
if args.cuda:
content_image = content_image.cuda()
content_image = Variable(content_image, volatile=True)
style_model = TransformerNet()
style_model.load_state_dict(torch.load(args.model))
if args.cuda:
style_model.cuda()
output = style_model(content_image)
if args.cuda:
output = output.cpu()
content_image=content_image.cpu()
output_data = output.data[0]
content_image_data=content_image.data[0]
output_data=torch.cat([content_image_data,output_data],2)
output_name=os.path.join(args.output_image,epoch_name+"_result_"+img)
utils.save_image(output_name, output_data)
# test_classifier = [ClassifierNet(i, num_classes) for i in model]
# for c, p in zip(test_classifier, test_classifier_path):
# c.load_state_dict(torch.load(p))
# if gain_model_name is not None:
# if gain_model_name in pretrainedmodels.model_names:
# backbone = pretrainedmodels.__dict__[gain_model_name](pretrained=None)
# backbone = nn.Sequential(*list(backbone.children())[: -2])
# else:
# raise Exception('\nModel {} not exist'.format(model_name))
# gain = GAINSolver(backbone, num_classes, in_channels=in_channels, devices=devices)
# gain.load_model(gain_checkpoint_path)
# else:
# gain = None
unet = TransformerNet(num_classes=num_classes * 3)
attack_net = AttackNet(unet)
solver = Attack(attack_net,
targeted=targeted,
num_classes=num_classes,
as_noise=as_noise)
solver.load_model(checkpoint_path)
images = [Image.open(i) for i in data['image_path']]
labels = data['label']
targets = data['target']
result = []
epoch = math.ceil(len(images) / inference_batch_size)
for i in range(epoch):
batch_x = images[i * inference_batch_size:(i + 1) *
inference_batch_size]
batch_y = (targets[i * inference_batch_size:(i + 1) *
def train(**kwargs):
opt._parse(kwargs)
device = t.device('cuda') if opt.use_gpu else t.device('cpu')
vis = Visualizer(opt.env)
# Data loading
transfroms = tv.transforms.Compose([
tv.transforms.Resize(opt.image_size),
tv.transforms.CenterCrop(opt.image_size),
tv.transforms.ToTensor(),
tv.transforms.Lambda(lambda x: x * 255)
])
dataset = tv.datasets.ImageFolder(opt.data_root, transfroms)
dataloader = data.DataLoader(dataset, opt.batch_size)
# style transformer network
transformer = TransformerNet()
if opt.model_path:
transformer.load_state_dict(
t.load(opt.model_path, map_location=lambda _s, _: _s))
transformer.to(device)
# Vgg16 for Perceptual Loss
vgg = Vgg16().eval()
vgg.to(device)
for param in vgg.parameters():
param.requires_grad = False
# Optimizer: use Adam
optimizer = t.optim.Adam(transformer.parameters(), opt.lr)
# Get style image
style = utils.get_style_data(opt.style_path)
vis.img('style', (style.data[0] * 0.225 + 0.45).clamp(min=0, max=1))
style = style.to(device)
# print("style.shape: ", style.shape)
# gram matrix for style image
with t.no_grad():
features_style = vgg(style)
gram_style = [utils.gram_matrix(y) for y in features_style]
# Loss meter
style_meter = tnt.meter.AverageValueMeter()
content_meter = tnt.meter.AverageValueMeter()
for epoch in range(opt.epoches):
content_meter.reset()
style_meter.reset()
for ii, (x, _) in tqdm.tqdm(enumerate(dataloader)):
# Train
optimizer.zero_grad()
x = x.to(device)
y = transformer(x)
y = utils.normalize_batch(y)
x = utils.normalize_batch(x)
features_y = vgg(y)
features_x = vgg(x)
# content loss
content_loss = opt.content_weight * F.mse_loss(
features_y.relu2_2, features_x.relu2_2)
# style loss
style_loss = 0.
for ft_y, gm_s in zip(features_y, gram_style):
gram_y = utils.gram_matrix(ft_y)
style_loss += F.mse_loss(gram_y, gm_s.expand_as(gram_y))
style_loss *= opt.style_weight
total_loss = content_loss + style_loss
total_loss.backward()
optimizer.step()
# Loss smooth for visualization
content_meter.add(content_loss.item())
style_meter.add(style_loss.item())
if (ii + 1) % opt.plot_every == 0:
if os.path.exists(opt.debug_file):
ipdb.set_trace()
# visualization
vis.plot('content_loss', content_meter.value()[0])
vis.plot('style_loss', style_meter.value()[0])
# denorm input/output, since we have applied (utils.normalize_batch)
vis.img('output',
(y.data.cpu()[0] * 0.225 + 0.45).clamp(min=0, max=1))
vis.img('input', (x.data.cpu()[0] * 0.225 + 0.45).clamp(min=0,
max=1))
# save checkpoint
vis.save([opt.env])
t.save(transformer.state_dict(), 'checkpoints/%s_style.pth' % epoch)
class GAN():
def __init__(self, args):
self.args = args
Tensor = torch.cuda.FloatTensor if args.gpu_ids else torch.Tensor
use_sigmoid = args.no_lsgan
# Global discriminator
self.netD = networks.define_D(args.input_nc, args.ndf,
args.which_model_netD, args.n_layers_D,
args.norm, use_sigmoid, args.gpu_ids)
# Local discriminator
self.netD_local = networks.define_D(args.input_nc, args.ndf,
args.which_model_netD,
args.n_layers_D, args.norm,
use_sigmoid, args.gpu_ids)
# Generator
self.netG = TransformerNet(args.norm, args.affine_state)
self.gan_loss = networks.GANLoss(use_lsgan=not args.no_lsgan,
tensor=Tensor)
self.identity_criterion = torch.nn.L1Loss()
# Resume
if args.resume_netG != '':
self.netG.load_state_dict(torch.load(args.resume_netG))
if args.resume_netD != '':
self.netD.load_state_dict(torch.load(args.resume_netD))
if args.resume_netD_local != '':
self.netD_local.load_state_dict(torch.load(args.resume_netD_local))
# optimizer
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
self.optimizer_D_local = torch.optim.Adam(self.netD_local.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
if self.args.cuda:
self.netD = self.netD.cuda()
self.netG = self.netG.cuda()
def update_D(self, content, fake):
self.netD.volatile = False
# feed real image
pred_real = self.netD(content)
loss_D_real = self.gan_loss(pred_real, True)
# feed fake image
pred_fake = self.netD(fake.detach())
loss_D_fake = self.gan_loss(pred_fake, False)
self.loss_D = (loss_D_real + loss_D_fake) * 0.5 * self.args.gan_weight
return self.loss_D
def update_G(self, fake):
self.netD.volatile = True
pred_fake = self.netD(fake)
self.loss_G = self.gan_loss(pred_fake, True) * self.args.gan_weight
return self.loss_G
def update_D_local(self, content_patch, fake_patch):
self.netD_local.volatile = False
# feed real sample
pred_real = self.netD_local(content_patch)
local_D_real = self.gan_loss(pred_real, True)
# feed fake sample
pred_fake = self.netD_local(fake_patch.detach())
local_D_fake = self.gan_loss(pred_fake, False)
self.loss_D_local = (local_D_real + local_D_fake) * \
0.5 * self.args.gan_local_weight
return self.loss_D_local
def update_G_local(self, fake_patch):
self.netD_local.volatile = True
pred_fake = self.netD_local(fake_patch)
self.loss_G_local = self.gan_loss(pred_fake,
True) * self.args.gan_local_weight
return self.loss_G_local
# the cropping and padding scheme
def crop_and_pad(self, image):
# cropping
image_crop_1 = image[:, :, 16:48, 0:128]
image_crop_2 = image[:, :, 48:80, 0:128]
image_crop_3 = image[:, :, 80:112, 0:128]
# padding
image_pad_1 = F.pad(image_crop_1, (0, 0, 16, 80))
image_pad_2 = F.pad(image_crop_2, (0, 0, 48, 48))
image_pad_3 = F.pad(image_crop_3, (0, 0, 80, 16))
return [image_pad_1, image_pad_2, image_pad_3]
# Cut video into frames
for filename in os.listdir(video_path):
if (filename.endswith(".mp4") and filename == sys.argv[1]): #or .avi, .mpeg, whatever.
os.system("ffmpeg -i " + os.path.join(video_path, video_name) + " " + os.path.join(frame_path, "%d.jpg"))
found = True
# Exit if video doesn't exist
if(not found):
print("No file named " + sys.argv[1] + " in directory " + video_path)
exit()
# Change to GPU if desired
device = torch.device("cuda")
# Load PyTorch Model
model = TransformerNet()
with torch.no_grad():
state_dict = torch.load(path_to_weights)
for k in list(state_dict.keys()):
if re.search(r'in\d+\.running_(mean|var)$', k):
del state_dict[k]
model.load_state_dict(state_dict)
model.to(device)
#Iterate over frames and stylize them
index = 0
for filename in os.listdir(frame_path):
if (filename.endswith(".jpg")):
print(os.path.join(frame_path, filename))