def run(opt, model, data_loader, mode):
if torch.cuda.is_available():
device = torch.device('cuda:'+str(opt.gpu_id))
else:
device = torch.device('cpu')
model = model.to(device)
tryon_dir = os.path.join(opt.data_root, mode, 'tryon-person')
mkdir(tryon_dir)
visual_dir = os.path.join(opt.out_dir, opt.name, mode)
mkdir(visual_dir)
data_iter = tqdm(data_loader, total=len(data_loader), bar_format='{l_bar}{r_bar}')
for _data in data_iter:
data = {}
for key, value in _data.items():
if not 'name' in key:
data[key] = value.to(device) # Load data on GPU
else:
data[key] = value
cloth = data['cloth']
cloth_mask = data['cloth_mask']
person = data['person']
cloth_name = data['cloth_name']
outputs = model(torch.cat([data['feature'], cloth],1)) # (batch, channel, height, width)
rendered_person, composition_mask = torch.split(outputs, 3,1)
rendered_person = torch.tanh(rendered_person)
composition_mask = torch.sigmoid(composition_mask)
tryon_person = cloth*composition_mask + rendered_person*(1-composition_mask)
visuals = [[data['head'], data['shape'], data['pose']],
[cloth, cloth_mask*2-1, composition_mask*2-1],
[rendered_person, tryon_person, person]]
save_images(tryon_person, cloth_name, tryon_dir)
save_visual(visuals, cloth_name, visual_dir)
def run(opt, model, data_loader, mode):
if torch.cuda.is_available():
device = torch.device('cuda:' + str(opt.gpu_id))
else:
device = torch.device('cpu')
model = model.to(device)
warp_cloth_dir = os.path.join(opt.data_root, mode, 'warp-cloth')
warp_cloth_mask_dir = os.path.join(opt.data_root, mode, 'warp-cloth-mask')
visual_dir = os.path.join(opt.out_dir, opt.name, mode)
dirs = [warp_cloth_dir, warp_cloth_mask_dir, visual_dir]
for d in dirs:
mkdir(d)
data_iter = tqdm(data_loader,
total=len(data_loader),
bar_format='{l_bar}{r_bar}')
for _data in data_iter:
data = {}
for key, value in _data.items():
if not 'name' in key:
data[key] = value.to(device) # Load data on GPU
else:
data[key] = value
cloth = data['cloth']
person = data['person']
body_mask = data['body_mask']
cloth_name = data['cloth_name']
grid, _ = model(data['feature'], cloth)
warped_cloth = F.grid_sample(cloth, grid, padding_mode='border')
warped_cloth_mask = F.grid_sample(data['cloth_mask'],
grid,
padding_mode='zeros')
save_images(warped_cloth, cloth_name, warp_cloth_dir)
save_images(warped_cloth_mask * 2 - 1, cloth_name, warp_cloth_mask_dir)
if mode == 'train': # No visuals
continue
warped_grid = F.grid_sample(data['grid'], grid, padding_mode='zeros')
warped_person = body_mask * person + (1 - body_mask) * warped_cloth
gt = body_mask * person + (1 - body_mask) * data['cloth_parse']
visuals = [[data['head'], data['shape'], data['pose']],
[cloth, warped_cloth, warped_grid],
[warped_person, gt, person]]
save_visual(visuals, cloth_name, visual_dir)
def main():
""" Training
run python main.py --dataset toumingzhi --model myganomaly --load_final_weights --batchsize 30
"""
opt = Options().parse()
dataloader = load_data(opt)
model = load_model(opt, dataloader)
print(model.device)
model.load_weights()
model.netd.eval()
model.netg.eval()
epoch = 0
print(
f">> Evaluating {model.name} on {model.opt.dataset} to detect {model.opt.abnormal_class}"
)
threshold = 0.016
# the code runs much slower in the first time, so run twice to measure its real speed
for i in range(2):
time_start = time.time()
with torch.no_grad():
# Create big error tensor for the test set.
an_scores = torch.zeros(size=(len(
model.dataloader.valid.dataset), ),
dtype=torch.float32,
device=model.device)
gt_labels = torch.zeros(size=(len(
model.dataloader.valid.dataset), ),
dtype=torch.long,
device=model.device)
times = []
for i, data in enumerate(model.dataloader.valid, 0):
time_i = time.time()
inputdata = data[0].to(model.device)
label = data[1].to(model.device)
fake, latent_i, latent_o = model.netg(inputdata)
_, feat_real = model.netd(inputdata)
_, feat_fake = model.netd(fake)
error = torch.mean(torch.pow((latent_i - latent_o), 2), dim=1)
time_o = time.time()
an_scores[i * model.opt.batchsize:i * model.opt.batchsize +
error.size(0)] = error.reshape(error.size(0))
gt_labels[i * model.opt.batchsize:i * model.opt.batchsize +
error.size(0)] = label.reshape(error.size(0))
if model.opt.save_test_images and i == 0:
test_img_dst = os.path.join(model.opt.outfolder,
model.opt.name,
model.opt.abnormal_class,
'test', 'images')
visualize.save_images(test_img_dst, epoch, inputdata, fake)
# if model.opt.visulize_feature and i == 0:
# feature_img_dst = os.path.join(model.opt.outtrain_dir, 'features')
# visualize.tsne_3D(feature_img_dst, epoch, 'feature',
# feat_real.reshape(feat_real.size(0), -1).cpu().numpy(),
# label.reshape(label.size(0), -1).cpu().numpy())
# visualize.tsne_2D(feature_img_dst, epoch, 'feature',
# feat_real.reshape(feat_real.size(0), -1).cpu().numpy(),
# label.reshape(label.size(0), -1).cpu().numpy())
times.append(time_o - time_i)
times = np.array(times)
times = np.mean(times[:100] * 1000)
# Scale error vector between [0, 1]
an_scores = (an_scores - torch.min(an_scores)) / (
torch.max(an_scores) - torch.min(an_scores))
nrm_trn_idx = torch.from_numpy(
np.where(an_scores.cpu().numpy() < threshold)[0])
""" predicte labels for test images, 0 is normal, 1 is anomaly
"""
pre_labels = torch.ones_like(gt_labels)
pre_labels[nrm_trn_idx] = 0
auc = evaluate(gt_labels, an_scores, metric=model.opt.metric)
# normal_scores = an_scores[np.where(gt_labels.cpu() == 0)[0]]
# abnormal_scores = an_scores[np.where(gt_labels.cpu() != 0)[0]]
# print("normal_scores:")
# print(normal_scores)
# print('abnormal_scores:')
# print(abnormal_scores)
# print(gt_labels)
print('pre_labels:')
print(pre_labels)
performance = OrderedDict([('Avg Run Time (ms/batch)', times),
('AUC', auc)])
print(performance)
print('total time:')
print((time.time() - time_start) * 1000)
def train(self):
opt = self.opt
# device to be finished==================================================
real_label = torch.ones(size=(opt.batchsize, ),
dtype=torch.float32,
device=self.device)
fake_label = torch.zeros(size=(opt.batchsize, ),
dtype=torch.float32,
device=self.device)
optimizer_d = optim.Adam(self.netd.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizer_g = optim.Adam(self.netg.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
l_bce = nn.BCELoss()
l_adv = l2_loss
l_con = nn.L1Loss()
l_enc = l2_loss
if opt.load_best_weights or opt.load_final_weights:
self.load_weights()
print(
f">> Training {self.name} on {opt.dataset} to detect {opt.abnormal_class}"
)
best_auc = 0
loss_d = []
loss_g = []
for epoch in range(opt.nepoch):
iternum = 0
running_loss_g = 0.0
running_loss_d = 0.0
self.netd.train()
self.netg.train()
for i, data in enumerate(
tqdm(self.dataloader.train,
leave=False,
total=len(self.dataloader.train),
ncols=80)):
iternum = iternum + 1
inputdata = data[0].to(self.device)
fake, latent_i, latent_o = self.netg(inputdata)
print(latent_i.shape)
# update netD
optimizer_d.zero_grad()
pred_real, feat_real = self.netd(inputdata)
pred_fake, feat_fake = self.netd(fake.detach())
err_d_real = l_bce(pred_real, real_label)
err_d_fake = l_bce(pred_fake, fake_label)
err_d = (err_d_real + err_d_fake) * 0.5
err_d.backward()
optimizer_d.step()
# update netG
optimizer_g.zero_grad()
pred_real, feat_real = self.netd(inputdata)
pred_fake, feat_fake = self.netd(fake)
err_g_adv = opt.w_adv * l_adv(feat_fake, feat_real)
err_g_con = opt.w_con * l_con(fake, inputdata)
err_g_lat = opt.w_lat * l_enc(latent_o, latent_i)
err_g = err_g_adv + err_g_con + err_g_lat
err_g.backward()
optimizer_g.step()
# record loss
running_loss_d += err_d.item()
running_loss_g += err_g.item()
if iternum % opt.loss_iter == 0:
# print('GLoss: {:.8f} DLoss: {:.8f}'
# .format(running_loss_g / opt.loss_iter, running_loss_d / opt.loss_iter))
loss_d.append(running_loss_d / opt.loss_iter)
loss_g.append(running_loss_g / opt.loss_iter)
running_loss_d = 0
running_loss_g = 0
if opt.save_train_images and i == 0:
train_img_dst = os.path.join(opt.outtrain_dir, 'images')
visualize.save_images(train_img_dst, epoch, inputdata,
fake)
print(">> Training model %s. Epoch %d/%d" %
(self.name, epoch + 1, opt.nepoch))
if epoch % opt.eva_epoch == 0:
performance = self.evaluate(epoch)
if performance['AUC'] > best_auc:
best_auc = performance['AUC']
if opt.save_best_weight:
self.save_weights(epoch, is_best=True)
# log performance
now = time.strftime("%c")
traintime = f'================ {now} ================\n'
visualize.write_to_log_file(
os.path.join(opt.outclass_dir, 'auc.log'), traintime)
visualize.log_current_performance(opt.outclass_dir,
performance, best_auc)
print(performance)
if opt.save_loss_curve:
visualize.plot_loss_curve(opt.outclass_dir, 'loss_d', loss_d)
visualize.plot_loss_curve(opt.outclass_dir, 'loss_g', loss_g)
if opt.save_final_weight:
self.save_weights(opt.nepoch, is_best=False)
print(">> Training model %s.[Done]" % self.name)
def evaluate(self, epoch):
self.netd.eval()
self.netg.eval()
print(
f">> Evaluating {self.name} on {self.opt.dataset} to detect {self.opt.abnormal_class}"
)
with torch.no_grad():
# Create big error tensor for the test set.
an_scores = torch.zeros(size=(len(
self.dataloader.valid.dataset), ),
dtype=torch.float32,
device=self.device)
gt_labels = torch.zeros(size=(len(
self.dataloader.valid.dataset), ),
dtype=torch.long,
device=self.device)
times = []
for i, data in enumerate(self.dataloader.valid, 0):
time_i = time.time()
inputdata = data[0].to(self.device)
label = data[1].to(self.device)
fake, latent_i, latent_o = self.netg(inputdata)
_, feat_real = self.netd(inputdata)
_, feat_fake = self.netd(fake)
error = torch.mean(torch.pow((latent_i - latent_o), 2), dim=1)
time_o = time.time()
an_scores[i * self.opt.batchsize:i * self.opt.batchsize +
error.size(0)] = error.reshape(error.size(0))
gt_labels[i * self.opt.batchsize:i * self.opt.batchsize +
error.size(0)] = label.reshape(error.size(0))
if self.opt.save_test_images and i == 0:
test_img_dst = os.path.join(self.opt.outfolder,
self.opt.name,
self.opt.abnormal_class,
'test', 'images')
visualize.save_images(test_img_dst, epoch, inputdata, fake)
if self.opt.visulize_feature and i == 0:
feature_img_dst = os.path.join(self.opt.outtrain_dir,
'features')
visualize.tsne_3D(
feature_img_dst, epoch, 'feature',
feat_real.reshape(feat_real.size(0), -1).cpu().numpy(),
label.reshape(label.size(0), -1).cpu().numpy())
visualize.tsne_2D(
feature_img_dst, epoch, 'feature',
feat_real.reshape(feat_real.size(0), -1).cpu().numpy(),
label.reshape(label.size(0), -1).cpu().numpy())
times.append(time_o - time_i)
times = np.array(times)
times = np.mean(times[:100] * 1000)
# Scale error vector between [0, 1]
an_scores = (an_scores - torch.min(an_scores)) / (
torch.max(an_scores) - torch.min(an_scores))
auc = evaluate(gt_labels, an_scores, metric=self.opt.metric)
performance = OrderedDict([('Avg Run Time (ms/batch)', times),
('AUC', auc)])
return performance
def train(self):
opt = self.opt
real_label = torch.ones(size=(opt.batchsize, ),
dtype=torch.float32,
device=self.device)
fake_label = torch.zeros(size=(opt.batchsize, ),
dtype=torch.float32,
device=self.device)
optimizer_de = optim.Adam(self.decoder.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizer_d = optim.Adam(self.netd.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizer_en = optim.Adam(self.encoder.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
l_bce = nn.BCELoss()
if opt.load_best_weights or opt.load_final_weights:
self.load_GAN_weights()
print(
f">> Training {self.name} on {opt.dataset} to detect {opt.abnormal_class}"
)
loss_de = []
loss_d = []
# train GAN
for epoch in range(opt.nepoch):
iternum = 0
lossd = 0
lossde = 0
self.decoder.train()
self.netd.train()
for i, data in enumerate(
tqdm(self.dataloader.train,
leave=False,
total=len(self.dataloader.train),
ncols=80)):
iternum = iternum + 1
inputdata = data[0].to(self.device)
z = torch.randn(inputdata.shape[0], opt.nz).to(self.device)
z = z.unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(z)
# update netd
optimizer_d.zero_grad()
pred_real, feat_real = self.netd(inputdata)
pred_fake, feat_fake = self.netd(fake_image)
err_d_real = l_bce(pred_real, real_label)
err_d_fake = l_bce(pred_fake, fake_label)
err_d = (err_d_real + err_d_fake) * 0.5
err_d.backward()
optimizer_d.step()
# update netde
optimizer_de.zero_grad()
z = torch.randn(inputdata.shape[0], opt.nz).to(self.device)
z = z.unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(z)
pred_fake, feat_fake = self.netd(fake_image)
err_g = l_bce(pred_fake, real_label)
err_g.backward()
optimizer_de.step()
# record loss
lossd += err_d.item()
lossde += err_g.item()
# record loss
if iternum % opt.loss_iter == 0:
# print('GLoss: {:.8f} DLoss: {:.8f}'
# .format(running_loss_g / opt.loss_iter, running_loss_d / opt.loss_iter))
loss_d.append(lossd / opt.loss_iter)
loss_de.append(lossde / opt.loss_iter)
lossd = 0
lossde = 0
if opt.save_train_images and i == 0:
train_img_dst = os.path.join(opt.outtrain_dir, 'images')
visualize.save_images(train_img_dst, epoch, inputdata,
fake_image)
print(">> Training model %s. Epoch %d/%d" %
(self.name, epoch + 1, opt.nepoch))
if opt.save_final_weight:
self.save_GAN_weights(opt.nepoch)
if opt.save_loss_curve:
visualize.plot_loss_curve(opt.outclass_dir, 'loss_d', loss_d)
visualize.plot_loss_curve(opt.outclass_dir, 'loss_de', loss_de)
# train encoder
print(
f">> Training {self.name} on {opt.dataset} to detect {opt.abnormal_class}"
)
if opt.load_best_en_weights or opt.load_final_en_weights:
self.load_encoder_weights()
best_auc = 0
loss_en = []
self.netd.eval()
self.decoder.eval()
for epoch in range(opt.nenepoch):
iternum = 0
lossen = 0
self.encoder.train()
for i, data in enumerate(
tqdm(self.dataloader.train,
leave=False,
total=len(self.dataloader.train),
ncols=80)):
iternum = iternum + 1
inputdata = data[0].to(self.device)
# update encoder
optimizer_en.zero_grad()
latent_z = self.encoder(inputdata).squeeze()
latent_z = latent_z.unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(latent_z)
pred_real, feat_real = self.netd(inputdata)
pred_fake, feat_fake = self.netd(fake_image)
k = 1.0
error_en = 1 / (inputdata.view(inputdata.shape[0], -1).shape[1]) * l2_loss(inputdata, fake_image) \
+ k * 1 / (feat_real.view(feat_real.shape[0], -1).shape[1]) * l2_loss(feat_real, feat_fake)
error_en.backward()
optimizer_en.step()
# record loss
lossen += error_en.item()
# record loss
if iternum % opt.loss_iter == 0:
# print('GLoss: {:.8f} DLoss: {:.8f}'
# .format(running_loss_g / opt.loss_iter, running_loss_d / opt.loss_iter))
loss_en.append(lossen / opt.loss_iter)
lossen = 0
if opt.save_train_images and i == 0:
train_img_dst = os.path.join(opt.outtrain_dir, 'images2')
visualize.save_images(train_img_dst, epoch, inputdata,
fake_image)
print(">> Training model %s. Epoch %d/%d" %
(self.name, epoch + 1, opt.nenepoch))
if epoch % opt.eva_epoch == 0:
performance = self.evaluate(epoch)
if performance['AUC'] > best_auc:
best_auc = performance['AUC']
if opt.save_best_en_weight:
self.save_encoder_weights(epoch, is_best=True)
# log performance
now = time.strftime("%c")
traintime = f'================ {now} ================\n'
visualize.write_to_log_file(
os.path.join(opt.outclass_dir, 'auc.log'), traintime)
visualize.log_current_performance(opt.outclass_dir,
performance, best_auc)
print(performance)
if opt.save_final_en_weight:
self.save_encoder_weights(opt.nepoch)
if opt.save_loss_curve:
visualize.plot_loss_curve(opt.outclass_dir, 'loss_en', loss_en)
print(">> Training model %s.[Done]" % self.name)
def train(self):
opt = self.opt
real_label = torch.ones(size=(opt.batchsize, ),
dtype=torch.float32,
device=self.device)
fake_label = torch.zeros(size=(opt.batchsize, ),
dtype=torch.float32,
device=self.device)
optimizer_f = optim.Adam(self.netf.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizer_de = optim.Adam(self.decoder.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizer_en = optim.Adam(self.encoder.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizer_d = optim.Adam(self.netd.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
l_bce = nn.BCELoss()
l_adv = l2_loss
l_con = nn.L1Loss()
l_enc = l2_loss
if opt.load_best_weights or opt.load_final_weights:
self.load_weights()
print(
f">> Training {self.name} on {opt.dataset} to detect {opt.abnormal_class}"
)
best_auc = 0
loss_eds = []
loss_fgs = []
loss_egs = []
for epoch in range(opt.nepoch):
iternum = 0
lossed = 0.0
lossfg = 0.0
losseg = 0.0
self.netf.train()
self.decoder.train()
self.encoder.train()
self.netd.train()
for i, data in enumerate(
tqdm(self.dataloader.train,
leave=False,
total=len(self.dataloader.train),
ncols=80)):
iternum = iternum + 1
data[0].requires_grad = True
inputdata = data[0].to(self.device)
z = torch.randn(inputdata.shape[0], opt.nz).to(self.device)
latent_i = self.netf(z).unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(latent_i)
latent_o = self.encoder(fake_image).squeeze()
real_latent = self.encoder(inputdata).squeeze()
pred_real = self.netd(real_latent)
pred_fake = self.netd(latent_o)
optimizer_en.zero_grad()
optimizer_d.zero_grad()
loss = loss_ed(pred_real, pred_fake)
# loss = losses.discriminator_logistic_simple_gp(pred_fake, pred_real, inputdata)
lossed += loss.item()
loss.backward()
optimizer_en.step()
optimizer_d.step()
z = torch.randn(inputdata.shape[0], opt.nz).to(self.device)
latent_i = self.netf(z).unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(latent_i)
latent_o = self.encoder(fake_image).squeeze()
pred_fake = self.netd(latent_o)
optimizer_f.zero_grad()
optimizer_de.zero_grad()
loss = loss_fg(pred_fake)
# loss = losses.generator_logistic_non_saturating(pred_fake)
lossfg += loss.item()
loss.backward()
optimizer_f.step()
optimizer_de.step()
z = torch.randn(inputdata.shape[0], opt.nz).to(self.device)
latent_i = self.netf(z).unsqueeze(2).unsqueeze(3)
fake_image = self.decoder(latent_i)
latent_o = self.encoder(fake_image).squeeze()
optimizer_en.zero_grad()
optimizer_de.zero_grad()
loss = l2_loss(latent_i, latent_o)
# loss = losses.reconstruction(latent_i, latent_o)
losseg += loss.item()
loss.backward()
optimizer_en.step()
optimizer_de.step()
# record loss
if iternum % opt.loss_iter == 0:
# print('GLoss: {:.8f} DLoss: {:.8f}'
# .format(running_loss_g / opt.loss_iter, running_loss_d / opt.loss_iter))
loss_eds.append(lossed / opt.loss_iter)
loss_fgs.append(lossfg / opt.loss_iter)
loss_egs.append(losseg / opt.loss_iter)
lossed = 0
lossfg = 0
losseg = 0
if opt.save_train_images and i == 0:
train_img_dst = os.path.join(opt.outtrain_dir, 'images')
visualize.save_images(train_img_dst, epoch, inputdata,
fake_image)
print(">> Training model %s. Epoch %d/%d" %
(self.name, epoch + 1, opt.nepoch))
# if epoch % opt.eva_epoch == 0:
# performance = self.evaluate(epoch)
# if performance['AUC'] > best_auc:
# best_auc = performance['AUC']
# if opt.save_best_weight:
# self.save_weights(epoch, is_best=True)
# # log performance
# now = time.strftime("%c")
# traintime = f'================ {now} ================\n'
# visualize.write_to_log_file(os.path.join(opt.outclass_dir, 'auc.log'), traintime)
# visualize.log_current_performance(opt.outclass_dir, performance, best_auc)
# print(performance)
if opt.save_loss_curve:
visualize.plot_loss_curve(opt.outclass_dir, 'loss_ed', loss_eds)
visualize.plot_loss_curve(opt.outclass_dir, 'loss_fg', loss_fgs)
visualize.plot_loss_curve(opt.outclass_dir, 'loss_eg', loss_egs)
# if opt.save_final_weight:
# self.save_weights(opt.nepoch, is_best=False)
print(">> Training model %s.[Done]" % self.name)