def setup(model_dir):
global model
opt = {}
opt = SimpleNamespace(**opt)
opt.nThreads = 1
opt.batchSize = 1
opt.serial_batches = True
opt.no_flip = True
opt.name = 'pretrained'
opt.checkpoints_dir = '.'
opt.model = 'pix2pix'
opt.which_direction = 'AtoB'
opt.norm = 'batch'
opt.input_nc = 3
opt.output_nc = 1
opt.which_model_netG = 'resnet_9blocks'
opt.no_dropout = True
opt.isTrain = False
opt.use_cuda = True
opt.ngf = 64
opt.ndf = 64
opt.init_type = 'normal'
opt.which_epoch = 'latest'
opt.pretrain_path = model_dir
model = create_model(opt)
return model
opt = TrainOptions().parse()
n_frames_G, n_frames_D = opt.n_frames_G, opt.n_frames_D
t_scales = opt.n_scales_temporal
input_nc = opt.input_nc
output_nc = opt.output_nc
visualizer = Visualizer(opt)
### initialize dataset
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
### initialize models
modelG, modelD, flowNet = create_model(opt)
iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt')
if opt.continue_train:
try:
start_epoch, epoch_iter = np.loadtxt(iter_path,
delimiter=',',
dtype=int)
except:
start_epoch, epoch_iter = 1, 0
if epoch_iter > 0:
### initialize dataset again
if opt.serial_batches:
data_loader = CreateDataLoader(opt, epoch_iter)
dataset = data_loader.load_data()
visualizer.vis_print('Resuming from epoch %d at iteration %d' %
def train():
'''load data'''
transform_mask = transforms.Compose([
# transforms.Resize((opt.fineSize,opt.fineSize)),
transforms.ToTensor(),
])
transform = transforms.Compose([
# transforms.RandomHorizontalFlip(),
# transforms.Resize((opt.fineSize,opt.fineSize)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3)
])
dataset_train = Data_load(opt.dataroot, opt.maskroot, transform,
transform_mask)
iterator_train = (data.DataLoader(dataset_train,
batch_size=opt.batchSize,
shuffle=True))
dataset_test = Data_load(opt.valinputroot, opt.valmask, transform,
transform_mask)
iterator_test = (data.DataLoader(dataset_test,
batch_size=opt.batchSize,
shuffle=False))
print(len(dataset_train))
print(len(dataset_test))
model = create_model(opt)
total_steps = 0
if not os.path.exists(opt.checkpoints_dir):
os.makedirs(opt.checkpoints_dir)
'''if resuming'''
# model.load(epoch_number/best)
best_score = 0
gts = []
preds = []
for i in range(100):
img_name = "{}.jpg".format(401 + i)
gts.append(os.path.join(opt.valroot, img_name))
pred_name = "{}.jpg".format(401 + i)
preds.append(os.path.join("pred", pred_name))
for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):
# for epoch in range(loaded_epoch_#+1, opt.niter + opt.niter_decay + 1):
print("this is a new epoch!")
epoch_start_time = time.time()
epoch_iter = 0
for ori_image, ori_mask in iterator_train:
align_corners = True
img_size = 512 # Or any other size
# Shrinking the image to a square. # ori_imgs is the original image
image = F.interpolate(ori_image,
img_size,
mode='bicubic',
align_corners=align_corners)
mask = F.interpolate(ori_mask,
img_size,
mode='bicubic',
align_corners=align_corners
) # mask looks pretty complete at this stage
image = image.clamp(min=-1, max=1)
mask = (mask > 0).type(torch.FloatTensor)
image = image.cuda()
mask = mask.cuda()
mask = mask[0][0]
mask = torch.unsqueeze(mask, 0)
mask = torch.unsqueeze(mask, 1)
mask = mask.byte()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
model.set_input(
image, mask) # sets both input data with mask and latent mask.
model.set_gt_latent()
model.optimize_parameters()
if total_steps % opt.display_freq == 0:
real_A, real_B, fake_B = model.get_current_visuals()
# real_A=input, real_B=ground truth fake_b=output
pic = (torch.cat([real_A, real_B, fake_B], dim=0) + 1) / 2.0
print("saving image!")
torchvision.utils.save_image(
pic,
'%s/Epoch_(%d)_(%dof%d).jpg' %
(opt.checkpoints_dir, epoch, total_steps + 1,
len(dataset_train)),
nrow=2)
if total_steps % 1 == 0:
errors = model.get_current_errors()
print(errors)
if epoch % opt.save_epoch_freq == 0:
print('saving the model at the end of epoch %d, iters %d' %
(epoch, total_steps))
model.save(epoch)
# evaluation #
generate_images(model, iterator_test)
mse_avg, ssim_avg = get_average_mse_ssim(gts, preds)
score = 1 - mse_avg / 100 + ssim_avg
if score > best_score:
best_score = score
model.save("best")
print(
f"best_score {best_score} found in epoch {epoch}. MSE:{mse_avg} SSIM:{ssim_avg}"
)
### end evaluation ###
print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch, opt.niter + opt.niter_decay,
time.time() - epoch_start_time))
model.update_learning_rate()
def test(opt):
opt.nThreads = 1 # test code only supports nThreads = 1
opt.batchSize = 1 # test code only supports batchSize = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#test batches = %d' % (int(dataset_size / len(opt.sort_order))))
visualizer = Visualizer(opt)
model = create_model(opt)
model.eval()
# create webpage
if opt.random_seed != -1:
exp_dir = '%s_%s_seed%s' % (opt.phase, opt.which_epoch, str(opt.random_seed))
else:
exp_dir = '%s_%s' % (opt.phase, opt.which_epoch)
web_dir = os.path.join(opt.results_dir, opt.name, exp_dir)
if opt.traverse or opt.deploy:
if opt.traverse:
out_dirname = 'traversal'
else:
out_dirname = 'deploy'
output_dir = os.path.join(web_dir,out_dirname)
if not os.path.isdir(output_dir):
os.makedirs(output_dir)
for image_path in opt.image_path_list:
print(image_path)
data = dataset.dataset.get_item_from_path(image_path)
visuals = model.inference(data)
if opt.traverse and opt.make_video:
out_path = os.path.join(output_dir, os.path.splitext(os.path.basename(image_path))[0] + '.mp4')
visualizer.make_video(visuals, out_path)
elif opt.traverse or (opt.deploy and opt.full_progression):
if opt.traverse and opt.compare_to_trained_outputs:
out_path = os.path.join(output_dir, os.path.splitext(os.path.basename(image_path))[0] + '_compare_to_{}_jump_{}.png'.format(opt.compare_to_trained_class, opt.trained_class_jump))
else:
out_path = os.path.join(output_dir, os.path.splitext(os.path.basename(image_path))[0] + '.png')
visualizer.save_row_image(visuals, out_path, traverse=opt.traverse)
else:
out_path = os.path.join(output_dir, os.path.basename(image_path[:-4]))
visualizer.save_images_deploy(visuals, out_path)
else:
webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))
# test
for i, data in enumerate(dataset):
if i >= opt.how_many:
break
visuals = model.inference(data)
img_path = data['Paths']
rem_ind = []
for i, path in enumerate(img_path):
if path != '':
print('process image... %s' % path)
else:
rem_ind += [i]
for ind in reversed(rem_ind):
del img_path[ind]
visualizer.save_images(webpage, visuals, img_path)
webpage.save()
train_list_dir_portrait = root + '/phoenix/S6/zl548/MegaDpeth_code/train_list/portrait/'
input_height = 320
input_width = 240
data_loader_p = CreateDataLoader(root, train_list_dir_portrait, input_height,
input_width)
dataset_p = data_loader_p.load_data()
dataset_size_p = len(data_loader_p)
print('========================= training portrait images = %d' %
dataset_size_p)
_isTrain = False
batch_size = 32
num_iterations_L = (dataset_size_l) / batch_size
num_iterations_P = (dataset_size_p) / batch_size
model = create_model(opt, _isTrain)
model.switch_to_train()
best_loss = 100
print("num_iterations ", num_iterations_L, num_iterations_P)
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
best_epoch = 0
total_iteration = 0
print(
"================================= BEGIN TRAINING ====================================="
)
start_epoch, epoch_iter = 1, 0
if opt.debug:
opt.display_freq = 1
opt.print_freq = 1
opt.niter = 1
opt.niter_decay = 0
opt.max_dataset_size = 10
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
num_batches = int(dataset_size / opt.batchSize)
model = create_model(opt, num_batches)
visualizer = Visualizer(opt)
total_steps = (start_epoch - 1) * dataset_size + epoch_iter
for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
if epoch != start_epoch:
epoch_iter = epoch_iter % num_batches
iter_start = epoch_iter
for i, data in enumerate(dataset, start=epoch_iter):
iter_start_time = time.time()
total_steps += opt.batchSize
epoch_iter += 1 #opt.batchSize
##################
## Forward Pass ##
elif opt.dataset == 'scannet':
eval_list_path = root + '/phoenix/S3/zl548/ScanNet/test_scannet_normal_list.txt'
eval_num_threads = 2
test_data_loader = CreateScanNetDataLoader(opt, eval_list_path, False,
EVAL_BATCH_SIZE,
eval_num_threads)
test_dataset = test_data_loader.load_data()
test_data_size = len(test_data_loader)
print('========================= ScanNet eval #images = %d =========' %
test_data_size)
else:
print('INPUT DATASET DOES NOT EXIST!!!')
sys.exit()
model = create_model(opt, _isTrain=False)
model.switch_to_train()
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
global_step = 0
def test_numerical(model, dataset, global_step):
rot_e_list = []
roll_e_list = []
pitch_e_list = []
count = 0.0
model.switch_to_eval()
def train():
opt = TrainOptions().parse()
if opt.distributed:
init_dist()
opt.batchSize = opt.batchSize // len(opt.gpu_ids)
### setup dataset
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
### setup trainer
trainer = Trainer(opt, data_loader)
### setup models
model, flowNet = create_model(opt, trainer.start_epoch)
flow_gt = conf_gt = [None] * 3
ref_idx_fix = torch.zeros([opt.batchSize])
for epoch in tqdm(range(trainer.start_epoch, opt.niter + opt.niter_decay + 1)):
trainer.start_of_epoch(epoch, model, data_loader)
n_frames_total, n_frames_load = data_loader.dataset.n_frames_total, opt.n_frames_per_gpu
for idx, data in enumerate(tqdm(dataset), start=trainer.epoch_iter):
trainer.start_of_iter()
if not opt.warp_ani:
data.update({'ani_image':None, 'ani_lmark':None, 'cropped_images':None, 'cropped_lmarks':None })
if not opt.no_flow_gt:
data_list = [data['tgt_mask_images'], data['cropped_images'], data['warping_ref'], data['ani_image']]
flow_gt, conf_gt = flowNet(data_list, epoch)
data_list = [data['tgt_label'], data['tgt_image'], data['tgt_template'], data['cropped_images'], flow_gt, conf_gt]
data_ref_list = [data['ref_label'], data['ref_image']]
data_prev = [None, None, None]
data_ani = [data['warping_ref_lmark'], data['warping_ref'], data['ori_warping_refs'], data['ani_lmark'], data['ani_image']]
############## Forward Pass ######################
prevs = {"raw_images":[], "synthesized_images":[], \
"prev_warp_images":[], "prev_weights":[], \
"ani_warp_images":[], "ani_weights":[], \
"ref_warp_images":[], "ref_weights":[], \
"ref_flows":[], "prev_flows":[], "ani_flows":[], \
"ani_syn":[]}
for t in range(0, n_frames_total, n_frames_load):
data_list_t = get_data_t(data_list, n_frames_load, t) + data_ref_list + \
get_data_t(data_ani, n_frames_load, t) + data_prev
g_losses, generated, data_prev, ref_idx = model(data_list_t, save_images=trainer.save, mode='generator', ref_idx_fix=ref_idx_fix)
g_losses = loss_backward(opt, g_losses, model.module.optimizer_G)
d_losses, _ = model(data_list_t, mode='discriminator', ref_idx_fix=ref_idx_fix)
d_losses = loss_backward(opt, d_losses, model.module.optimizer_D)
# store previous
store_prev(generated, prevs)
loss_dict = dict(zip(model.module.lossCollector.loss_names, g_losses + d_losses))
output_data_list = [prevs] + [data['ref_image']] + data_ani + data_list + [data['tgt_mask_images']]
if trainer.end_of_iter(loss_dict, output_data_list, model):
break
# pdb.set_trace()
trainer.end_of_epoch(model)
def train(opt):
iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt')
if opt.continue_train:
if opt.which_epoch == 'latest':
try:
start_epoch, epoch_iter = np.loadtxt(iter_path,
delimiter=',',
dtype=int)
except:
start_epoch, epoch_iter = 1, 0
else:
start_epoch, epoch_iter = int(opt.which_epoch), 0
print('Resuming from epoch %d at iteration %d' %
(start_epoch, epoch_iter))
for update_point in opt.decay_epochs:
if start_epoch < update_point:
break
opt.lr *= opt.decay_gamma
else:
start_epoch, epoch_iter = 0, 0
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
model = create_model(opt)
visualizer = Visualizer(opt)
total_steps = (start_epoch) * dataset_size + epoch_iter
display_delta = total_steps % opt.display_freq
print_delta = total_steps % opt.print_freq
save_delta = total_steps % opt.save_latest_freq
bSize = opt.batchSize
#in case there's no display sample one image from each class to test after every epoch
if opt.display_id == 0:
dataset.dataset.set_sample_mode(True)
dataset.num_workers = 1
for i, data in enumerate(dataset):
if i * opt.batchSize >= opt.numClasses:
break
if i == 0:
sample_data = data
else:
for key, value in data.items():
if torch.is_tensor(data[key]):
sample_data[key] = torch.cat(
(sample_data[key], data[key]), 0)
else:
sample_data[key] = sample_data[key] + data[key]
dataset.num_workers = opt.nThreads
dataset.dataset.set_sample_mode(False)
for epoch in range(start_epoch, opt.epochs):
epoch_start_time = time.time()
if epoch != start_epoch:
epoch_iter = 0
for i, data in enumerate(dataset, start=epoch_iter):
iter_start_time = time.time()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
# whether to collect output images
save_fake = (total_steps % opt.display_freq
== display_delta) and (opt.display_id > 0)
############## Network Pass ########################
model.set_inputs(data)
disc_losses = model.update_D()
gen_losses, gen_in, gen_out, rec_out, cyc_out = model.update_G(
infer=save_fake)
loss_dict = dict(gen_losses, **disc_losses)
##################################################
############## Display results and errors ##########
### print out errors
if total_steps % opt.print_freq == print_delta:
errors = {
k: v.item()
if not (isinstance(v, float) or isinstance(v, int)) else v
for k, v in loss_dict.items()
}
t = (time.time() - iter_start_time) / opt.batchSize
visualizer.print_current_errors(epoch + 1, epoch_iter, errors,
t)
if opt.display_id > 0:
visualizer.plot_current_errors(
epoch,
float(epoch_iter) / dataset_size, opt, errors)
### display output images
if save_fake and opt.display_id > 0:
class_a_suffix = ' class {}'.format(data['A_class'][0])
class_b_suffix = ' class {}'.format(data['B_class'][0])
classes = None
visuals = OrderedDict()
visuals_A = OrderedDict([('real image' + class_a_suffix,
util.tensor2im(gen_in.data[0]))])
visuals_B = OrderedDict([('real image' + class_b_suffix,
util.tensor2im(gen_in.data[bSize]))])
A_out_vis = OrderedDict([('synthesized image' + class_b_suffix,
util.tensor2im(gen_out.data[0]))])
B_out_vis = OrderedDict([('synthesized image' + class_a_suffix,
util.tensor2im(gen_out.data[bSize]))
])
if opt.lambda_rec > 0:
A_out_vis.update([('reconstructed image' + class_a_suffix,
util.tensor2im(rec_out.data[0]))])
B_out_vis.update([('reconstructed image' + class_b_suffix,
util.tensor2im(rec_out.data[bSize]))])
if opt.lambda_cyc > 0:
A_out_vis.update([('cycled image' + class_a_suffix,
util.tensor2im(cyc_out.data[0]))])
B_out_vis.update([('cycled image' + class_b_suffix,
util.tensor2im(cyc_out.data[bSize]))])
visuals_A.update(A_out_vis)
visuals_B.update(B_out_vis)
visuals.update(visuals_A)
visuals.update(visuals_B)
ncols = len(visuals_A)
visualizer.display_current_results(visuals, epoch, classes,
ncols)
### save latest model
if total_steps % opt.save_latest_freq == save_delta:
print('saving the latest model (epoch %d, total_steps %d)' %
(epoch + 1, total_steps))
model.save('latest')
np.savetxt(iter_path, (epoch, epoch_iter),
delimiter=',',
fmt='%d')
if opt.display_id == 0:
model.eval()
visuals = model.inference(sample_data)
visualizer.save_matrix_image(visuals, 'latest')
model.train()
# end of epoch
iter_end_time = time.time()
print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch + 1, opt.epochs, time.time() - epoch_start_time))
### save model for this epoch
if (epoch + 1) % opt.save_epoch_freq == 0:
print('saving the model at the end of epoch %d, iters %d' %
(epoch + 1, total_steps))
model.save('latest')
model.save(epoch + 1)
np.savetxt(iter_path, (epoch + 1, 0), delimiter=',', fmt='%d')
if opt.display_id == 0:
model.eval()
visuals = model.inference(sample_data)
visualizer.save_matrix_image(visuals, epoch + 1)
model.train()
### multiply learning rate by opt.decay_gamma after certain iterations
if (epoch + 1) in opt.decay_epochs:
model.update_learning_rate()
from data.data_loader import CreateDataLoader
from models.models import create_model
import utils.util as util
from utils.visualizer import Visualizer
from utils import html
from torch.autograd import Variable
import time
opt = TestOptions().parse(save=False)
opt.nThreads = 1
opt.batchSize = 1
opt.serial_batches = True # no shuffle
data_loader = CreateDataLoader(opt)
dataset, _ = data_loader.load_data()
model = create_model(opt, data_loader.dataset)
visualizer = Visualizer(opt)
# create website
web_dir = os.path.join(opt.results_dir, opt.name,
'%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(web_dir, '%s: %s' % (opt.name, opt.which_epoch))
img_dir = os.path.join(web_dir, 'images')
# test
label_trues, label_preds = [], []
model.model.eval()
tic = time.time()
accs = []
D_mesh = D_mesh.unsqueeze_(0).expand(batchSize, dpt, hgt, wdt)
H_mesh = H_mesh.unsqueeze_(0).expand(batchSize, dpt, hgt, wdt)
W_mesh = W_mesh.unsqueeze_(0).expand(batchSize, dpt, hgt, wdt)
D_upmesh = dDepth.float() + D_mesh
H_upmesh = dHeight.float() + H_mesh
W_upmesh = dWidth.float() + W_mesh
return torch.stack([D_upmesh, H_upmesh, W_upmesh], dim=1)
if __name__ == '__main__':
opt = TestOptions().parse()
opt.nThreads = 1
opt.batchSize = 1
model_regist = create_model(opt)
visualizer = Visualizer(opt)
stn = Dense3DSpatialTransformer()
# create website
web_dir = os.path.join(opt.results_dir, opt.name,
'%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(
web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' %
(opt.name, opt.phase, opt.which_epoch))
datafiles = []
dataFiles = sorted(os.listdir(opt.dataroot))
for isub, dataName in enumerate(dataFiles):
datafiles.append(os.path.join(opt.dataroot, dataName))
def train(config, writer, logger):
config = config.opt
config.distributed = False
if 'WORLD_SIZE' in os.environ:
config.distributed = int(os.environ['WORLD_SIZE']) > 1
config.gpu = 0
config.world_size = 1
if config.distributed:
config.gpu = config.local_rank
torch.cuda.set_device(config.gpu)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
config.world_size = torch.distributed.get_world_size()
data_set = CreateDataLoader(config).load_data()
total_steps = config.epochs * len(data_set)
if config.gpu == 0:
print(len(data_set), "# of Training Images")
model = create_model(config)
visualizer = Visualizer(config)
model.named_buffers = lambda: []
average_tensor = utils.load_average_img(config)
average_tensor = average_tensor.view(1, *average_tensor.shape).cuda()
if config.fp16:
from apex import amp
from apex.parallel import DistributedDataParallel as DDP
model = model.cuda()
model, [optimizer_G, optimizer_D] = \
amp.initialize(model, [model.optimizer_G, model.optimizer_D],
opt_level='O1')
if config.distributed:
model = DDP(model, delay_allreduce=True)
else:
model = torch.nn.DataParallel(model)
step = 0
for epoch in range(config.epochs):
if config.gpu == 0:
print("epoch: ", epoch)
for i, data in enumerate(data_set):
save_gen = (step + 1) % config.display_freq == 0
if i == 0 or (not config.no_temporal_smoothing and \
np.random.random() < config.prob_restart_sequence):
if config.no_temporal_smoothing:
prev_generated = prev_label = prev_real = None
else:
prev_generated = average_tensor
prev_label = torch.zeros_like(data['label']).cuda()
prev_real = average_tensor
data['label'] = data['label'][:, :1]
assert data['label'].shape[1] == 1
losses, generated = model(Variable(data['label']).cuda(),
Variable(data['inst'].cuda()),
Variable(data['image'].cuda()),
Variable(data['feat'].cuda()),
prev_label,
prev_generated,
prev_real,
infer=save_gen or \
(not config.no_temporal_smoothing))
# average=average_tensor)
# sum per device losses
losses = [
torch.mean(x) if not isinstance(x, int) else x for x in losses
]
loss_dict = dict(zip(model.module.loss_names, losses))
# calculate final loss scalar
loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5
loss_G = loss_dict['G_GAN'] + loss_dict.get('G_GAN_Feat', 0) + \
loss_dict.get('G_VGG', 0)
# update generator weights\n",
model.module.optimizer_G.zero_grad()
if config.fp16:
with amp.scale_loss(loss_G, optimizer_G) as scaled_loss:
scaled_loss.backward()
else:
loss_G.backward()
model.module.optimizer_G.step()
# update discriminator weights\n",
model.module.optimizer_D.zero_grad()
if config.fp16:
with amp.scale_loss(loss_D, optimizer_D) as scaled_loss:
scaled_loss.backward()
else:
loss_D.backward()
model.module.optimizer_D.step()
if (not config.distributed) or torch.distributed.get_rank() == 0:
if (step +
1) % config.print_freq == 0 or step == total_steps - 1:
logger.info("Train: [{:2d}/{}] Step {:03d}/{:03d}".format(
epoch + 1, config.epochs, i,
len(data_set) - 1))
logger.info("Loss D: {}, Loss G {}, Loss VGG {}".format(
loss_D.item(), loss_dict['G_GAN'].item(),
loss_dict['G_VGG'].item()))
if save_gen:
visuals = OrderedDict([
('input_label',
util.tensor2label(data['label'][0], config.label_nc)),
('synthesized_image',
util.tensor2im(generated.data[0])),
('real_image', util.tensor2im(data['image'][0]))
])
if not config.no_temporal_smoothing:
visuals['prev_label'] = util.tensor2label(
prev_label[0], config.label_nc)
visuals['prev_real'] = util.tensor2im(prev_real[0])
visuals['prev_generated'] = util.tensor2im(
prev_generated.data[0])
visualizer.display_current_results(visuals, epoch,
total_steps)
if (step + 1) % config.save_latest_freq == 0 or \
step == total_steps - 1:
model.module.save('latest')
model.module.save(epoch)
if not config.no_temporal_smoothing:
prev_generated = generated.cuda()
prev_real = data['image'].detach().cuda()
prev_label = data['label'].detach().cuda()
step += 1
### train the entire network after certain iterations
if (config.niter_fix_global != 0) and (epoch
== config.niter_fix_global):
model.module.update_fixed_params()
### linearly decay learning rate after certain iterations
if epoch > config.niter:
model.module.update_learning_rate()
def infer(n, image_label_path, image_inst_path):
opt = TestOptions().parse(save=False)
opt.nThreads = 1 # test code only supports nThreads = 1
opt.batchSize = 1 # test code only supports batchSize = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
opt.name = "label2city_1024p"
opt.netG = "local"
opt.ngf = 32
opt.resize_or_crop = "none"
data_loader = CreateOneDataLoader(opt)
dataset = data_loader.load_data(image_label_path, image_inst_path)
visualizer = Visualizer(opt)
# create website
#web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))
#webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))
# test
if not opt.engine and not opt.onnx:
model = create_model(opt)
if opt.data_type == 16:
model.half()
elif opt.data_type == 8:
model.type(torch.uint8)
if opt.verbose:
print(model)
else:
from run_engine import run_trt_engine, run_onnx
for i, data in enumerate(dataset):
if i >= opt.how_many:
break
if opt.data_type == 16:
data['label'] = data['label'].half()
data['inst'] = data['inst'].half()
elif opt.data_type == 8:
data['label'] = data['label'].uint8()
data['inst'] = data['inst'].uint8()
if opt.export_onnx:
print("Exporting to ONNX: ", opt.export_onnx)
assert opt.export_onnx.endswith(
"onnx"), "Export model file should end with .onnx"
torch.onnx.export(model, [data['label'], data['inst']],
opt.export_onnx,
verbose=True)
exit(0)
minibatch = 1
if opt.engine:
generated = run_trt_engine(opt.engine, minibatch,
[data['label'], data['inst']])
elif opt.onnx:
generated = run_onnx(opt.onnx, opt.data_type, minibatch,
[data['label'], data['inst']])
else:
generated = model.inference(data['label'], data['inst'])
visuals = OrderedDict([
('input_label', util.tensor2label(data['label'][0], opt.label_nc)),
('synthesized_image', util.tensor2im(generated.data[0]))
])
img_path = data['path']
print('process image... %s' % img_path)
visualizer.save_image(visuals, n)
def getEnlighten(input_image, cFlag):
opt = Namespace(D_P_times2=False,
IN_vgg=False,
aspect_ratio=1.0,
batchSize=1,
checkpoints_dir=baseLoc + 'weights/',
dataroot='test_dataset',
dataset_mode='unaligned',
display_id=1,
display_port=8097,
display_single_pane_ncols=0,
display_winsize=256,
fcn=0,
fineSize=256,
gpu_ids=[0],
high_times=400,
how_many=50,
hybrid_loss=False,
identity=0.0,
input_linear=False,
input_nc=3,
instance_norm=0.0,
isTrain=False,
l1=10.0,
lambda_A=10.0,
lambda_B=10.0,
latent_norm=False,
latent_threshold=False,
lighten=False,
linear=False,
linear_add=False,
loadSize=286,
low_times=200,
max_dataset_size='inf',
model='single',
multiply=False,
nThreads=1,
n_layers_D=3,
n_layers_patchD=3,
name='enlightening',
ndf=64,
new_lr=False,
ngf=64,
no_dropout=True,
no_flip=True,
no_vgg_instance=False,
noise=0,
norm='instance',
norm_attention=False,
ntest='inf',
output_nc=3,
patchD=False,
patchD_3=0,
patchSize=64,
patch_vgg=False,
phase='test',
resize_or_crop='no',
results_dir='./results/',
self_attention=True,
serial_batches=True,
skip=1.0,
syn_norm=False,
tanh=False,
times_residual=True,
use_avgpool=0,
use_mse=False,
use_norm=1.0,
use_ragan=False,
use_wgan=0.0,
vary=1,
vgg=0,
vgg_choose='relu5_3',
vgg_maxpooling=False,
vgg_mean=False,
which_direction='AtoB',
which_epoch='200',
which_model_netD='basic',
which_model_netG='sid_unet_resize',
cFlag=cFlag)
im = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
transform = get_transform(opt)
A_img = transform(im)
r, g, b = A_img[0] + 1, A_img[1] + 1, A_img[2] + 1
A_gray = 1. - (0.299 * r + 0.587 * g + 0.114 * b) / 2.
A_gray = torch.unsqueeze(A_gray, 0)
data = {
'A': A_img.unsqueeze(0),
'B': A_img.unsqueeze(0),
'A_gray': A_gray.unsqueeze(0),
'input_img': A_img.unsqueeze(0),
'A_paths': 'A_path',
'B_paths': 'B_path'
}
model = create_model(opt)
model.set_input(data)
visuals = model.predict()
out = visuals['fake_B'].astype(np.uint8)
out = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
# cv2.imwrite("/Users/kritiksoman/PycharmProjects/new/out.png", out)
return out
def test(opt):
opt.nThreads = 1 # test code only supports nThreads = 1
opt.batchSize = 1 # test code only supports batchSize = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
opt.no_instance = True
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
visualizer = Visualizer(opt)
# create website
web_dir = os.path.join(opt.results_dir,
'%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(
web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' %
(opt.name, opt.phase, opt.which_epoch))
model = create_model(opt)
ssim_sum = 0
psnr_sum = 0
mse_sum = 0
scores = np.zeros((len(dataset), 3))
indices = np.array(choose_k_images(len(dataset), 100))
saved_scores = np.zeros((len(indices), 3))
for i, data in enumerate(dataset):
generated = model.inference(data['label'], data['inst'], data['image'])
visuals = OrderedDict([
('input_label', util.tensor2label(data['label'][0], opt.label_nc)),
('synthesized_image', util.tensor2im(generated.data[0])),
('real_image', util.tensor2im(data['image'][0]))
])
img_path = data['path']
scores = save_scores(scores, data, generated, i)
if i in indices:
# save to their visualizerx
print('process image... %s' % img_path)
visualizer.save_images(webpage, visuals, img_path)
#save to my files
save_img(visuals, opt, img_path[0])
saved_scores[np.where(indices == i)[0]] = scores[i]
print(i)
webpage.save()
dic_saved = {"viz matrix": saved_scores, "mean": saved_scores.mean(axis=0)}
with open(opt.viz_dir + "/saved_results.pkl", 'wb') as f:
pickle.dump(dic_saved, f)
mean_scores = scores.mean(axis=0)
dic = {"scores matrix": scores, "mean": mean_scores}
with open(opt.results_dir + "results.pkl", 'wb') as f:
pickle.dump(dic, f)
num_gpus = torch.cuda.device_count()
for gpu_id in range(num_gpus):
torch.cuda.set_device(gpu_id)
torch.cuda.empty_cache()
torch.cuda.set_device(0)
return mean_scores
def setup(opts):
opt.name = opts['checkpoints_root'].split('/')[-1]
opt.checkpoints_dir = os.path.join(opts['checkpoints_root'], '..')
model = create_model(opt)
return model
def main():
# Create train dataset
train_set_opt = opt.datasets[0]
train_set = create_dataset(train_set_opt)
train_size = int(math.ceil(len(train_set) / train_set_opt.batch_size))
print('Number of train images: %d batches of size %d' % (train_size, train_set_opt.batch_size))
total_iters = int(opt.train.niter)
total_epoches = int(math.ceil(total_iters / train_size))
print('Total epoches needed: %d' % total_epoches)
# Create val dataset
val_set_opt = opt.datasets[1]
val_set = create_dataset(val_set_opt)
val_size = len(val_set)
print('Number of val images: %d' % val_size)
# Create dataloader
train_loader = create_dataloader(train_set, train_set_opt)
val_loader = create_dataloader(val_set, val_set_opt)
# Create model
model = create_model(opt)
# Create binarization module
import bin
global bin_op
bin_op = bin.BinOp(model.netG)
model.train()
# Create logger
logger = Logger(opt)
current_step = 0
need_make_val_dir = True
start_time = time.time()
for epoch in range(total_epoches):
for i, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
train_start_time = time.time()
# Training
model.feed_data(train_data)
# optimize_parameters 함수를 분할해서 대체함 for binarization
# model.optimize_parameters(current_step)
bin_op.binarization()
model.forward_G()
model.optimizer_G.zero_grad()
model.backward_G()
bin_op.restore()
bin_op.updateBinaryGradWeight()
model.optimizer_G.step()
train_duration = time.time() - train_start_time
if current_step % opt.logger.print_freq == 0:
losses = model.get_current_losses()
logger.print_results(losses, epoch, current_step, train_duration, 'loss')
if current_step % opt.logger.save_checkpoint_freq == 0:
print('Saving the model at the end of current_step %d' % (current_step))
model.save(current_step)
# Validation
if current_step % opt.train.val_freq == 0:
validate(val_loader, val_size, model, logger, epoch, current_step)
model.update_learning_rate(step=current_step, scheme=opt.train.lr_scheme)
print('End of Epoch %d' % epoch)
print('Saving the final model')
model.save('latest')
print('End of Training \t Time Taken: %d sec' % (time.time() - start_time))
# from util.util import tensor2im, get_display_image, one_hot
#
# from val import get_test_result
# #from eval.test_is_2 import get_IS
# import fnmatch
opt = TestOptions().parse(save=False)
opt.nThreads = 1 # test code only supports nThreads = 1
opt.batchSize = 1 # test code only supports batchSize = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
model = create_model(opt, 'resNet')
visualizer = Visualizer(opt)
dataset_size = len(data_loader)
print('#testing images = %d' % dataset_size)
def main():
web_dir = os.path.join(opt.results_dir, opt.name,
'%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(
web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' %
(opt.name, opt.phase, opt.which_epoch))
for i, data in enumerate(dataset):
if i >= opt.how_many:
break
opt = TestOptions().parse(save=False)
opt.nThreads = 1 # test code only supports nThreads = 1
opt.batchSize = 1 # test code only supports batchSize = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
opt.stage = 12 ## choose stage_I_II_dataset.py
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
opt.name = "stage_I_gan_ganFeat_noL1_oneD_Parsing_bz50_parsing20_04222"
opt.which_G = "resNet"
opt.stage = 1
opt.which_epoch = 100
model_1 = create_model(opt)
opt.name = "gan_L1_feat_vgg_notv_noparsing_afftps_05102228"
opt.which_G = "wapResNet_v3_afftps"
opt.stage = 2
opt.which_epoch = 45
model_2 = create_model(opt)
visualizer = Visualizer(opt)
dataset_size = len(data_loader)
print('#testing images = %d' % dataset_size)
geo = GeoAPI()
affTnf = GeometricTnf(geometric_model='affine', use_cuda=False)
tpsTnf = GeometricTnf(geometric_model='tps', use_cuda=False)
def train():
opt = TrainOptions().parse()
if opt.debug:
opt.display_freq = 1
opt.print_freq = 1
opt.nThreads = 1
### initialize dataset
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training frames = %d' % dataset_size)
### initialize models
models = create_model(opt)
modelG, modelD, flowNet, optimizer_G, optimizer_D, optimizer_D_T = create_optimizer(
opt, models)
### set parameters
n_gpus, tG, tD, tDB, s_scales, t_scales, input_nc_1, input_nc_2, \
start_epoch, epoch_iter, print_freq, total_steps, iter_path = init_params_composer(opt, modelG, modelD, data_loader)
visualizer = Visualizer(opt)
### real training starts here
for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
for idx, data in enumerate(dataset, start=epoch_iter):
if total_steps % print_freq == 0:
iter_start_time = time.time()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
# whether to collect output images
save_fake = total_steps % opt.display_freq == 0
n_frames_total, n_frames_load, t_len = data_loader.dataset.init_data_params(
data, n_gpus, tG)
fake_SI_prev_last, frames_all = data_loader.dataset.init_data(
t_scales)
for i in range(0, n_frames_total, n_frames_load):
input_TParsing, input_TFG, input_SPose, input_SParsing, input_SFG, input_BG, input_SFG_full, input_SI = data_loader.dataset.prepare_data_composer(
data, i)
###################################### Forward Pass ##########################
####### generator
fake_SI, fake_SI_raw, fake_sd, fake_SFG_full, fake_SFG_res, flow, weight, real_input_T, real_input_S, real_input_SFG, real_input_BG, real_SIp, real_SFG_fullp, fake_SI_last \
= modelG(input_TParsing, input_TFG, input_SPose, input_SParsing, input_SFG, input_BG, input_SFG_full, input_SI, fake_SI_prev_last)
####### discriminator
### individual frame discriminator
real_SI_prev, real_SI = real_SIp[:, :
-1], real_SIp[:,
1:] # the collection of previous and current real frames
real_SFG_full_prev, real_SFG_full = real_SFG_fullp[:, :
-1], real_SFG_fullp[:,
1:]
#flow_ref, conf_ref = flowNet(real_SI, real_SI_prev) # reference flows and confidences
flow_ref, conf_ref = flowNet(
real_SFG_full,
real_SFG_full_prev) # reference flows and confidences
fake_SI_prev = modelG.module.compute_fake_B_prev(
real_SI_prev, fake_SI_prev_last, fake_SI)
fake_SI_prev_last = fake_SI_last
real_input_BG_flag = data['BG_flag']
losses = modelD(0, [
real_SI, real_SFG_full, fake_SI, fake_SI_raw,
fake_SFG_full, fake_SFG_res, real_input_T, real_input_S,
real_input_SFG, real_input_BG, real_input_BG_flag,
real_SI_prev, fake_SI_prev, real_SFG_full_prev, flow,
weight, flow_ref, conf_ref
])
losses = [
torch.mean(x) if x is not None else 0 for x in losses
]
loss_dict = dict(zip(modelD.module.loss_names, losses))
### temporal discriminator
# get skipped frames for each temporal scale
frames_all, frames_skipped = modelD.module.get_all_skipped_frames(frames_all, \
real_SI, fake_SI, flow_ref, conf_ref, t_scales, tD, n_frames_load, i, flowNet)
# run discriminator for each temporal scale
loss_dict_T = []
for s in range(t_scales):
if frames_skipped[0][s] is not None:
losses = modelD(s + 1, [
frame_skipped[s]
for frame_skipped in frames_skipped
])
losses = [
torch.mean(x) if not isinstance(x, int) else x
for x in losses
]
loss_dict_T.append(
dict(zip(modelD.module.loss_names_T, losses)))
# collect losses
loss_G, loss_D, loss_D_T, t_scales_act = modelD.module.get_losses(
loss_dict, loss_dict_T, t_scales)
###################################### Backward Pass #################################
# update generator weights
loss_backward(opt, loss_G, optimizer_G)
# update individual discriminator weights
loss_backward(opt, loss_D, optimizer_D)
# update temporal discriminator weights
for s in range(t_scales_act):
loss_backward(opt, loss_D_T[s], optimizer_D_T[s])
if i == 0:
fake_SI_first = fake_SI[
0, 0] # the first generated image in this sequence
if opt.debug:
call([
"nvidia-smi", "--format=csv",
"--query-gpu=memory.used,memory.free"
])
############## Display results and errors ##########
### print out errors
if total_steps % print_freq == 0:
t = (time.time() - iter_start_time) / print_freq
errors = {
k: v.data.item() if not isinstance(v, int) else v
for k, v in loss_dict.items()
}
for s in range(len(loss_dict_T)):
errors.update({
k + str(s):
v.data.item() if not isinstance(v, int) else v
for k, v in loss_dict_T[s].items()
})
visualizer.print_current_errors_new(epoch, epoch_iter, errors,
modelD.module.loss_names,
t)
visualizer.plot_current_errors(errors, total_steps)
### display output images
if save_fake:
visuals = util.save_all_tensors_composer(
opt, real_input_T, real_input_S, real_input_SFG,
real_input_BG, fake_SI, fake_SI_raw, fake_SI_first,
fake_SFG_full, fake_SFG_res, fake_sd, real_SI,
real_SFG_full, flow_ref, conf_ref, flow, weight, modelD)
visualizer.display_current_results(visuals, epoch, total_steps)
### save latest model
save_models(opt, epoch, epoch_iter, total_steps, visualizer,
iter_path, modelG, modelD)
if epoch_iter > dataset_size - opt.batchSize:
epoch_iter = 0
break
# end of epoch
iter_end_time = time.time()
visualizer.vis_print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch, opt.niter + opt.niter_decay,
time.time() - epoch_start_time))
### save model for this epoch and update model params
save_models(opt,
epoch,
epoch_iter,
total_steps,
visualizer,
iter_path,
modelG,
modelD,
end_of_epoch=True)
update_models(opt, epoch, modelG, modelD, data_loader)
def main():
os.makedirs('sample', exist_ok=True)
opt = TestOptions().parse()
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('# Inference images = %d' % dataset_size)
model = create_model(opt)
for i, data in enumerate(dataset):
# add gaussian noise channel
# wash the label
t_mask = torch.FloatTensor(
(data['label'].cpu().numpy() == 7).astype(np.float))
#
# data['label'] = data['label'] * (1 - t_mask) + t_mask * 4
mask_clothes = torch.FloatTensor(
(data['label'].cpu().numpy() == 4).astype(np.int))
mask_fore = torch.FloatTensor(
(data['label'].cpu().numpy() > 0).astype(np.int))
mask_back = torch.FloatTensor(
(data['label'].cpu().numpy() == 0).astype(np.int))
img_fore = data['image'] * mask_fore
img_back = data['image'] * mask_back
all_clothes_label = changearm(data['label'])
############## Forward Pass ######################
fake_image, warped_cloth, refined_cloth = model(
Variable(data['label'].cuda()), Variable(data['edge'].cuda()),
Variable(img_fore.cuda()), Variable(mask_clothes.cuda()),
Variable(data['color'].cuda()), Variable(all_clothes_label.cuda()),
Variable(data['image'].cuda()), Variable(data['pose'].cuda()),
Variable(data['image'].cuda()), Variable(mask_fore.cuda()))
# Restore Background
fake_image = fake_image * mask_fore.cuda() + img_back.cuda()
# make output folders
output_dir = os.path.join(opt.results_dir, opt.phase)
fake_image_dir = os.path.join(output_dir, 'try-on')
os.makedirs(fake_image_dir, exist_ok=True)
warped_cloth_dir = os.path.join(output_dir, 'warped_cloth')
os.makedirs(warped_cloth_dir, exist_ok=True)
refined_cloth_dir = os.path.join(output_dir, 'refined_cloth')
os.makedirs(refined_cloth_dir, exist_ok=True)
# save output
for j in range(opt.batchSize):
print("Saving", data['name'][j])
util.save_tensor_as_image(
fake_image[j], os.path.join(fake_image_dir, data['name'][j]))
util.save_tensor_as_image(
warped_cloth[j], os.path.join(warped_cloth_dir,
data['name'][j]))
util.save_tensor_as_image(
refined_cloth[j],
os.path.join(refined_cloth_dir, data['name'][j]))
import os
from options.test_options import TestOptions
from data.data_loader import CreateDataLoader
from models.models import create_model
from util.visualizer import Visualizer
from util import html
opt = TestOptions().parse()
opt.nThreads = 1 # test code only supports nThreads = 1
opt.batchSize = 1 # test code only supports batchSize = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
model = create_model(opt)
visualizer = Visualizer(opt)
# create website
web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))
# test
for i, data in enumerate(dataset):
if i >= opt.how_many:
break
model.set_input(data)
model.test()
visuals = model.get_current_visuals()
img_path = model.get_image_paths()
print('process image... %s' % img_path)
visualizer.save_images(webpage, visuals, img_path)
from torch.autograd import Variable
from collections import OrderedDict
from options.test_options import TestOptions
from data.data_loader import CreateDataLoader
from models.models import create_model
import util.util as util
from util.visualizer import Visualizer
opt = TestOptions().parse(save=False)
opt.nThreads = 1
opt.batchSize = 1
opt.serial_batches = True
visualizer = Visualizer(opt)
modelG = create_model(opt)
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('Generating %d frames' % dataset_size)
save_dir = os.path.join(opt.results_dir, opt.name, opt.which_epoch + '_epoch',
opt.phase)
total_distance, total_pixels = 0, 0
mtotal_distance, mtotal_pixels = 0, 0
mouth_total_distance, mouth_total_pixels = 0, 0
for i, data in enumerate(dataset):
