Spaces:
Paused
Paused
| """ | |
| Created on Dec 13, 2018 | |
| @author: Yuedong Chen | |
| """ | |
| from data import create_dataloader | |
| from model import create_model | |
| from visualizer import Visualizer | |
| import copy | |
| import time | |
| import os | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| def create_solver(opt): | |
| instance = Solver() | |
| instance.initialize(opt) | |
| return instance | |
| class Solver(object): | |
| """docstring for Solver""" | |
| def __init__(self): | |
| super(Solver, self).__init__() | |
| def initialize(self, opt): | |
| self.opt = opt | |
| self.visual = Visualizer() | |
| self.visual.initialize(self.opt) | |
| def run_solver(self): | |
| if self.opt.mode == "train": | |
| self.train_networks() | |
| else: | |
| self.test_networks(self.opt) | |
| def train_networks(self): | |
| # init train setting | |
| self.init_train_setting() | |
| # for every epoch | |
| for epoch in range(self.opt.epoch_count, self.epoch_len + 1): | |
| # train network | |
| self.train_epoch(epoch) | |
| # update learning rate | |
| self.cur_lr = self.train_model.update_learning_rate() | |
| # save checkpoint if needed | |
| if epoch % self.opt.save_epoch_freq == 0: | |
| self.train_model.save_ckpt(epoch) | |
| # save the last epoch | |
| self.train_model.save_ckpt(self.epoch_len) | |
| def init_train_setting(self): | |
| self.train_dataset = create_dataloader(self.opt) | |
| self.train_model = create_model(self.opt) | |
| self.train_total_steps = 0 | |
| self.epoch_len = self.opt.niter + self.opt.niter_decay | |
| self.cur_lr = self.opt.lr | |
| def train_epoch(self, epoch): | |
| epoch_start_time = time.time() | |
| epoch_steps = 0 | |
| last_print_step_t = time.time() | |
| for idx, batch in enumerate(self.train_dataset): | |
| self.train_total_steps += self.opt.batch_size | |
| epoch_steps += self.opt.batch_size | |
| # train network | |
| self.train_model.feed_batch(batch) | |
| self.train_model.optimize_paras(train_gen=(idx % self.opt.train_gen_iter == 0)) | |
| # print losses | |
| if self.train_total_steps % self.opt.print_losses_freq == 0: | |
| cur_losses = self.train_model.get_latest_losses() | |
| avg_step_t = (time.time() - last_print_step_t) / self.opt.print_losses_freq | |
| last_print_step_t = time.time() | |
| # print loss info to command line | |
| info_dict = {'epoch': epoch, 'epoch_len': self.epoch_len, | |
| 'epoch_steps': idx * self.opt.batch_size, 'epoch_steps_len': len(self.train_dataset), | |
| 'step_time': avg_step_t, 'cur_lr': self.cur_lr, | |
| 'log_path': os.path.join(self.opt.ckpt_dir, self.opt.log_file), | |
| 'losses': cur_losses | |
| } | |
| self.visual.print_losses_info(info_dict) | |
| # plot loss map to visdom | |
| if self.train_total_steps % self.opt.plot_losses_freq == 0 and self.visual.display_id > 0: | |
| cur_losses = self.train_model.get_latest_losses() | |
| epoch_steps = idx * self.opt.batch_size | |
| self.visual.display_current_losses(epoch - 1, epoch_steps / len(self.train_dataset), cur_losses) | |
| # display image on visdom | |
| if self.train_total_steps % self.opt.sample_img_freq == 0 and self.visual.display_id > 0: | |
| cur_vis = self.train_model.get_latest_visuals() | |
| self.visual.display_online_results(cur_vis, epoch) | |
| # latest_aus = model.get_latest_aus() | |
| # visual.log_aus(epoch, epoch_steps, latest_aus, opt.ckpt_dir) | |
| def test_networks(self, opt): | |
| self.init_test_setting(opt) | |
| self.test_ops() | |
| def init_test_setting(self, opt): | |
| self.test_dataset = create_dataloader(opt) | |
| self.test_model = create_model(opt) | |
| def test_ops(self): | |
| for batch_idx, batch in enumerate(self.test_dataset): | |
| with torch.no_grad(): | |
| # interpolate several times | |
| faces_list = [batch['src_img'].float().numpy()] | |
| paths_list = [batch['src_path'], batch['tar_path']] | |
| for idx in range(self.opt.interpolate_len): | |
| cur_alpha = (idx + 1.) / float(self.opt.interpolate_len) | |
| cur_tar_aus = cur_alpha * batch['tar_aus'] + (1 - cur_alpha) * batch['src_aus'] | |
| # print(batch['src_aus']) | |
| # print(cur_tar_aus) | |
| test_batch = {'src_img': batch['src_img'], 'tar_aus': cur_tar_aus, 'src_aus':batch['src_aus'], 'tar_img':batch['tar_img']} | |
| self.test_model.feed_batch(test_batch) | |
| self.test_model.forward() | |
| cur_gen_faces = self.test_model.fake_img.cpu().float().numpy() | |
| faces_list.append(cur_gen_faces) | |
| faces_list.append(batch['tar_img'].float().numpy()) | |
| self.test_save_imgs(faces_list, paths_list) | |
| def test_save_imgs(self, faces_list, paths_list): | |
| for idx in range(len(paths_list[0])): | |
| src_name = os.path.splitext(os.path.basename(paths_list[0][idx]))[0] | |
| tar_name = os.path.splitext(os.path.basename(paths_list[1][idx]))[0] | |
| if self.opt.save_test_gif: | |
| import imageio | |
| imgs_numpy_list = [] | |
| for face_idx in range(len(faces_list) - 1): # remove target image | |
| cur_numpy = np.array(self.visual.numpy2im(faces_list[face_idx][idx])) | |
| imgs_numpy_list.extend([cur_numpy for _ in range(3)]) | |
| saved_path = os.path.join(self.opt.results, "%s_%s.gif" % (src_name, tar_name)) | |
| imageio.mimsave(saved_path, imgs_numpy_list) | |
| else: | |
| # concate src, inters, tar faces | |
| concate_img = np.array(self.visual.numpy2im(faces_list[0][idx])) | |
| for face_idx in range(1, len(faces_list)): | |
| concate_img = np.concatenate((concate_img, np.array(self.visual.numpy2im(faces_list[face_idx][idx]))), axis=1) | |
| concate_img = Image.fromarray(concate_img) | |
| # save image | |
| saved_path = os.path.join(self.opt.results, "%s_%s.jpg" % (src_name, tar_name)) | |
| concate_img.save(saved_path) | |
| print("[Success] Saved images to %s" % saved_path) | |