from trainer import Trainer from importlib import import_module import math import torch from torch import optim from torch.optim import lr_scheduler import numpy as np from metrics import calculate_metrics import os import glob import torch.nn.functional as F from models.temporal_patch_gan import Discriminator import cv2 import cvbase import imageio from skimage.feature import canny from models.lafc_single import Model from data.util.flow_utils import region_fill as rf class Network(Trainer): def init_model(self): model_package = import_module('models.{}'.format(self.opt['model'])) model = model_package.Model(self.opt) dist_in = 3 discriminator = Discriminator(in_channels=dist_in, conv_type=self.opt['conv_type'], dist_cnum=self.opt['dist_cnum']) optimizer = optim.Adam(model.parameters(), lr=float(self.opt['train']['lr']), betas=(float(self.opt['train']['BETA1']), float(float(self.opt['train']['BETA2'])))) dist_optim = optim.Adam(discriminator.parameters(), lr=float(self.opt['train']['lr']), betas=(float(self.opt['train']['BETA1']), float(float(self.opt['train']['BETA2'])))) if self.rank <= 0: self.logger.info( 'Optimizer is Adam, BETA1: {}, BETA2: {}'.format(float(self.opt['train']['BETA1']), float(self.opt['train']['BETA2']))) step_size = int(math.ceil(self.opt['train']['UPDATE_INTERVAL'] / self.trainSize)) if self.rank <= 0: self.logger.info('Step size for optimizer is {} epoch'.format(step_size)) scheduler = lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=self.opt['train']['lr_decay']) dist_scheduler = lr_scheduler.StepLR(dist_optim, step_size=step_size, gamma=self.opt['train']['lr_decay']) return model, discriminator, optimizer, dist_optim, scheduler, dist_scheduler def init_flow_model(self): flow_model = Model(self.opt['flow_config']) state = torch.load(self.opt['flow_checkPoint'], map_location=lambda storage, loc: storage.cuda(self.opt['device'])) flow_model.load_state_dict(state['model_state_dict']) flow_model = flow_model.to(self.opt['device']) return flow_model def resume_training(self): gen_state = torch.load(self.opt['path']['gen_state'], map_location=lambda storage, loc: storage.cuda(self.opt['device'])) dis_state = torch.load(self.opt['path']['dis_state'], map_location=lambda storage, loc: storage.cuda(self.opt['device'])) opt_state = torch.load(self.opt['path']['opt_state'], map_location=lambda storage, loc: storage.cuda(self.opt['device'])) if self.rank <= 0: self.logger.info('Resume state is activated') self.logger.info('Resume training from epoch: {}, iter: {}'.format( opt_state['epoch'], opt_state['iteration'] )) if self.opt['finetune'] == False: start_epoch = opt_state['epoch'] current_step = opt_state['iteration'] self.optimizer.load_state_dict(opt_state['optimizer_state_dict']) self.dist_optim.load_state_dict(opt_state['dist_optim_state_dict']) self.scheduler.load_state_dict(opt_state['scheduler_state_dict']) self.dist_scheduler.load_state_dict(opt_state['dist_scheduler_state_dict']) else: start_epoch = 0 current_step = 0 self.model.load_state_dict(gen_state['model_state_dict']) self.dist.load_state_dict(dis_state['dist_state_dict']) if self.rank <= 0: self.logger.info('Resume training mode, optimizer, scheduler and model have been uploaded') return start_epoch, current_step def norm_flows(self, flows): flattened_flows = flows.flatten(3) flow_max = torch.max(flattened_flows, dim=-1, keepdim=True)[0] flows = flows / flow_max.unsqueeze(-1) return flows def _trainEpoch(self, epoch): for idx, train_data in enumerate(self.trainLoader): self.currentStep += 1 if self.currentStep > self.totalIterations: if self.rank <= 0: self.logger.info('Train process has been finished') break if self.opt['train']['WARMUP'] is not None and self.currentStep <= self.opt['train']['WARMUP'] // self.opt[ 'world_size']: target_lr = self.opt['train']['lr'] * self.currentStep / ( self.opt['train']['WARMUP']) self.adjust_learning_rate(self.optimizer, target_lr) frames = train_data['frames'] # tensor, [b, t, c, h, w] masks = train_data['masks'] # tensor, [b, t, c, h, w] if self.opt['flow_direction'] == 'for': flows = train_data['forward_flo'] elif self.opt['flow_direction'] == 'back': flows = train_data['backward_flo'] elif self.opt['flow_direction'] == 'bi': raise NotImplementedError('Bidirectory flow mode is not implemented') else: raise ValueError('Unknown flow mode: {}'.format(self.opt['flow_direction'])) frames = frames.to(self.opt['device']) # [b, t, c(3), h, w] masks = masks.to(self.opt['device']) # [b, t, 1, h, w] flows = flows.to(self.opt['device']) # [b, t, c(2), h, w] b, t, c, h, w = flows.shape flows = flows.reshape(b * t, c, h, w) compressed_masks = masks.reshape(b * t, 1, h, w) with torch.no_grad(): flows = self.flow_model(flows, compressed_masks)[0] # filled flows flows = flows.reshape(b, t, c, h, w) flows = self.norm_flows(flows) b, t, c, h, w = frames.shape cm, cf = masks.shape[2], flows.shape[2] masked_frames = frames * (1 - masks) filled_frames = self.model(masked_frames, flows, masks) # filled_frames shape: [b, t, c, h, w] frames = frames.view(b * t, c, h, w) masks = masks.view(b * t, cm, h, w) comp_img = filled_frames * masks + frames * (1 - masks) real_vid_feat = self.dist(frames, t) fake_vid_feat = self.dist(comp_img.detach(), t) dis_real_loss = self.adversarial_loss(real_vid_feat, True, True) dis_fake_loss = self.adversarial_loss(fake_vid_feat, False, True) dis_loss = (dis_real_loss + dis_fake_loss) / 2 self.dist_optim.zero_grad() dis_loss.backward() self.dist_optim.step() # calculate generator loss gen_vid_feat = self.dist(comp_img, t) gan_loss = self.adversarial_loss(gen_vid_feat, True, False) gen_loss = gan_loss * self.opt['adv'] L1Loss_valid = self.validLoss(filled_frames * (1 - masks), frames * (1 - masks)) / torch.mean(1 - masks) L1Loss_masked = self.validLoss(filled_frames * masks, frames * masks) / torch.mean(masks) m_loss_valid = L1Loss_valid * self.opt['L1M'] m_loss_masked = L1Loss_masked * self.opt['L1V'] loss = m_loss_valid + m_loss_masked + gen_loss self.optimizer.zero_grad() loss.backward() self.optimizer.step() if self.opt['use_tb_logger'] and self.rank <= 0 and self.currentStep % 8 == 0: print('Mask: {:.03f}, valid: {:.03f}, dis_fake: {:.03f}, dis_real: {:.03f}, adv: {:.03f}'.format( m_loss_masked.item(), m_loss_valid.item(), dis_fake_loss.item(), dis_real_loss.item(), gen_loss.item() )) if self.opt['use_tb_logger'] and self.rank <= 0 and self.currentStep % 64 == 0: self.tb_logger.add_scalar('{}/recon_mask'.format('train'), m_loss_masked.item(), self.currentStep) self.tb_logger.add_scalar('{}/recon_valid'.format('train'), m_loss_valid.item(), self.currentStep) self.tb_logger.add_scalar('{}/adv'.format('train'), gen_loss.item(), self.currentStep) self.tb_logger.add_scalar('train/dist', dis_loss.item(), self.currentStep) if self.currentStep % self.opt['logger']['PRINT_FREQ'] == 0 and self.rank <= 0: c_frames = comp_img.detach().permute(0, 2, 3, 1).cpu() f_frames = frames.detach().permute(0, 2, 3, 1).cpu() compLog = np.clip(np.array((c_frames + 1) / 2 * 255), 0, 255).astype(np.uint8) framesLog = np.clip(np.array((f_frames + 1) / 2 * 255), 0, 255).astype(np.uint8) logs = calculate_metrics(compLog, framesLog) self._printLog(logs, epoch, loss) def _printLog(self, logs, epoch, loss): if self.countDown % self.opt['record_iter'] == 0: self.total_psnr = 0 self.total_ssim = 0 self.total_l1 = 0 self.total_l2 = 0 self.total_loss = 0 self.countDown = 0 self.countDown += 1 message = '[epoch:{:3d}, iter:{:7d}, lr:('.format(epoch, self.currentStep) for v in self.get_lr(): message += '{:.3e}, '.format(v) message += ')] ' self.total_psnr += logs['psnr'] self.total_ssim += logs['ssim'] self.total_l1 += logs['l1'] self.total_l2 += logs['l2'] self.total_loss += loss.item() mean_psnr = self.total_psnr / self.countDown mean_ssim = self.total_ssim / self.countDown mean_l1 = self.total_l1 / self.countDown mean_l2 = self.total_l2 / self.countDown mean_loss = self.total_loss / self.countDown message += '{:s}: {:.4e} '.format('mean_loss', mean_loss) message += '{:s}: {:} '.format('mean_psnr', mean_psnr) message += '{:s}: {:} '.format('mean_ssim', mean_ssim) message += '{:s}: {:} '.format('mean_l1', mean_l1) message += '{:s}: {:} '.format('mean_l2', mean_l2) if self.opt['use_tb_logger']: self.tb_logger.add_scalar('train/mean_psnr', mean_psnr, self.currentStep) self.tb_logger.add_scalar('train/mean_ssim', mean_ssim, self.currentStep) self.tb_logger.add_scalar('train/mean_l1', mean_l1, self.currentStep) self.tb_logger.add_scalar('train/mean_l2', mean_l2, self.currentStep) self.tb_logger.add_scalar('train/mean_loss', mean_loss, self.currentStep) self.logger.info(message) if self.currentStep % self.opt['logger']['SAVE_CHECKPOINT_FREQ'] == 0: self.save_checkpoint(epoch) def save_checkpoint(self, epoch): if isinstance(self.model, torch.nn.DataParallel) or isinstance(self.model, torch.nn.parallel.DistributedDataParallel): model_state = self.model.module.state_dict() dist_state = self.dist.module.state_dict() else: model_state = self.model.state_dict() dist_state = self.dist.state_dict() gen_state = { 'model_state_dict': model_state } dis_state = { 'dist_state_dict': dist_state } opt_state = { 'epoch': epoch, 'iteration': self.currentStep, 'optimizer_state_dict': self.optimizer.state_dict(), 'dist_optim_state_dict': self.dist_optim.state_dict(), 'scheduler_state_dict': self.scheduler.state_dict(), 'dist_scheduler_state_dict': self.dist_scheduler.state_dict() } gen_name = os.path.join(self.opt['path']['TRAINING_STATE'], 'gen_{}_{}.pth.tar'.format(epoch, self.currentStep)) dist_name = os.path.join(self.opt['path']['TRAINING_STATE'], 'dist_{}_{}.pth.tar'.format(epoch, self.currentStep)) opt_name = os.path.join(self.opt['path']['TRAINING_STATE'], 'opt_{}_{}.pth.tar'.format(epoch, self.currentStep)) torch.save(gen_state, gen_name) torch.save(dis_state, dist_name) torch.save(opt_state, opt_name) def _validate(self, epoch): frame_path = self.valInfo['frame_root'] mask_path = self.valInfo['mask_root'] flow_path = self.valInfo['flow_root'] self.model.eval() test_list = os.listdir(flow_path) if len(test_list) > 10: test_list = test_list[:10] # only valid 10 videos to save test time width, height = self.valInfo['flow_width'], self.valInfo['flow_height'] psnr, ssim, l1, l2 = {}, {}, {}, {} pivot, sequenceLen, ref_length = 20, self.opt['num_frames'], self.opt['ref_length'] for i in range(len(test_list)): videoName = test_list[i] if self.rank <= 0: self.logger.info('Video {} is been processed'.format(videoName)) frame_dir = os.path.join(frame_path, videoName) mask_dir = os.path.join(mask_path, videoName) flow_dir = os.path.join(flow_path, videoName) videoLen = len(glob.glob(os.path.join(mask_dir, '*.png'))) neighbor_ids = [i for i in range(max(0, pivot - sequenceLen // 2), min(videoLen, pivot + sequenceLen // 2))] ref_ids = self.get_ref_index(neighbor_ids, videoLen, ref_length) ref_ids.extend(neighbor_ids) frames = self.read_frames(frame_dir, width, height, ref_ids) masks = self.read_masks(mask_dir, width, height, ref_ids) flows = self.read_flows(flow_dir, width, height, ref_ids, videoLen - 1) if frames == [] or masks == []: if self.rank <= 0: print('Video {} doesn\'t have enough frames'.format(videoName)) continue flows = self.diffusion_flows(flows, masks) if self.rank <= 0: self.logger.info('Frames, masks, and flows have been read') frames = np.stack(frames, axis=0) # [t, h, w, c] masks = np.stack(masks, axis=0) flows = np.stack(flows, axis=0) frames = torch.from_numpy(np.transpose(frames, (0, 3, 1, 2))).unsqueeze(0).float() flows = torch.from_numpy(np.transpose(flows, (0, 3, 1, 2))).unsqueeze(0).float() masks = torch.from_numpy(np.transpose(masks, (0, 3, 1, 2))).unsqueeze(0).float() frames = frames / 127.5 - 1 frames = frames.to(self.opt['device']) masks = masks.to(self.opt['device']) flows = flows.to(self.opt['device']) b, t, c, h, w = flows.shape flows = flows.reshape(b * t, c, h, w) compressed_masks = masks.reshape(b * t, 1, h, w) with torch.no_grad(): flows = self.flow_model(flows, compressed_masks)[0] flows = flows.reshape(b, t, c, h, w) flows = self.norm_flows(flows) b, t, c, h, w = frames.shape cm, cf = masks.shape[2], flows.shape[2] masked_frames = frames * (1 - masks) with torch.no_grad(): filled_frames = self.model(masked_frames, flows, masks) frames = frames.view(b * t, c, h, w) masks = masks.view(b * t, cm, h, w) comp_img = filled_frames * masks + frames * (1 - masks) # [t, c, h, w] # calculate metrics psnr_avg, ssim_avg, l1_avg, l2_avg = self.metrics_calc(comp_img, frames) psnr[videoName] = psnr_avg ssim[videoName] = ssim_avg l1[videoName] = l1_avg l2[videoName] = l2_avg # visualize frames and report the phase performance if self.rank <= 0: if self.opt['use_tb_logger']: self.tb_logger.add_scalar('test/{}/l1'.format(videoName), l1_avg, self.currentStep) self.tb_logger.add_scalar('test/{}/l2'.format(videoName), l2_avg, self.currentStep) self.tb_logger.add_scalar('test/{}/psnr'.format(videoName), psnr_avg, self.currentStep) self.tb_logger.add_scalar('test/{}/ssim'.format(videoName), ssim_avg, self.currentStep) masked_frames = masked_frames.view(b * t, c, h, w) self.vis_frames(comp_img, masked_frames, frames, videoName, epoch) # view the difference between diffused flows and the completed flows mean_psnr = np.mean([psnr[k] for k in psnr.keys()]) mean_ssim = np.mean([ssim[k] for k in ssim.keys()]) mean_l1 = np.mean([l1[k] for k in l1.keys()]) mean_l2 = np.mean([l2[k] for k in l2.keys()]) self.logger.info( '[epoch:{:3d}, vid:{}/{}], mean_l1: {:.4e}, mean_l2: {:.4e}, mean_psnr: {:}, mean_ssim: {:}'.format( epoch, i, len(test_list), mean_l1, mean_l2, mean_psnr, mean_ssim)) # give the overall performance if self.rank <= 0: mean_psnr = np.mean([psnr[k] for k in psnr.keys()]) mean_ssim = np.mean([ssim[k] for k in ssim.keys()]) mean_l1 = np.mean([l1[k] for k in l1.keys()]) mean_l2 = np.mean([l2[k] for k in l2.keys()]) self.logger.info( '[epoch:{:3d}], mean_l1: {:.4e} mean_l2: {:.4e} mean_psnr: {:} mean_ssim: {:}'.format( epoch, mean_l1, mean_l2, mean_psnr, mean_ssim)) self.save_checkpoint(epoch) self.model.train() def get_ref_index(self, neighbor_ids, videoLen, ref_length): ref_indices = [] for i in range(0, videoLen, ref_length): if not i in neighbor_ids: ref_indices.append(i) return ref_indices def metrics_calc(self, results, frames): psnr_avg, ssim_avg, l1_avg, l2_avg = 0, 0, 0, 0 results = np.array(results.permute(0, 2, 3, 1).cpu()) frames = np.array(frames.permute(0, 2, 3, 1).cpu()) result = np.clip((results + 1) / 2 * 255, 0, 255).astype(np.uint8) frames = np.clip((frames + 1) / 2 * 255, 0, 255).astype(np.uint8) logs = calculate_metrics(result, frames) psnr_avg += logs['psnr'] ssim_avg += logs['ssim'] l1_avg += logs['l1'] l2_avg += logs['l2'] return psnr_avg, ssim_avg, l1_avg, l2_avg def read_frames(self, frame_dir, width, height, ref_indices): frame_paths = sorted(glob.glob(os.path.join(frame_dir, '*.jpg'))) frames = [] if len(frame_paths) <= 30: return frames for i in ref_indices: frame_path = os.path.join(frame_dir, '{:05d}.jpg'.format(i)) frame = imageio.imread(frame_path) frame = cv2.resize(frame, (width, height), cv2.INTER_LINEAR) frames.append(frame) return frames def diffusion_flows(self, flows, masks): assert len(flows) == len(masks), 'Length of flow: {}, length of mask: {}'.format(len(flows), len(masks)) ret_flows = [] for i in range(len(flows)): flow, mask = flows[i], masks[i] flow = self.diffusion_flow(flow, mask) ret_flows.append(flow) return ret_flows def diffusion_flow(self, flow, mask): mask = mask[:, :, 0] flow_filled = np.zeros(flow.shape) flow_filled[:, :, 0] = rf.regionfill(flow[:, :, 0] * (1 - mask), mask) flow_filled[:, :, 1] = rf.regionfill(flow[:, :, 1] * (1 - mask), mask) return flow_filled def read_flows(self, flow_dir, width, height, ref_ids, frameMaxIndex): if self.opt['flow_direction'] == 'for': direction = 'forward_flo' shift = 0 elif self.opt['flow_direction'] == 'back': direction = 'backward_flo' shift = -1 elif self.opt['flow_direction'] == 'bi': raise NotImplementedError('Bidirectional flows processing are not implemented') else: raise ValueError('Unknown flow direction: {}'.format(self.opt['flow_direction'])) flows = [] flow_path = os.path.join(flow_dir, direction) for i in ref_ids: i += shift if i >= frameMaxIndex: i = frameMaxIndex - 1 if i < 0: i = 0 flow_p = os.path.join(flow_path, '{:05d}.flo'.format(i)) flow = cvbase.read_flow(flow_p) pre_height, pre_width = flow.shape[:2] flow = cv2.resize(flow, (width, height), cv2.INTER_LINEAR) flow[:, :, 0] = flow[:, :, 0] / pre_width * width flow[:, :, 1] = flow[:, :, 1] / pre_height * height flows.append(flow) return flows def load_edges(self, frames, width, height): edges = [] for i in range(len(frames)): frame = frames[i] frame_gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) edge = canny(frame_gray, sigma=self.valInfo['sigma'], mask=None, low_threshold=self.valInfo['low_threshold'], high_threshold=self.valInfo['high_threshold']).astype(np.float) # [h, w, 1] edge_t = self.to_tensor(edge, width, height, mode='nearest') edges.append(edge_t) return edges def to_tensor(self, frame, width, height, mode='bilinear'): if len(frame.shape) == 2: frame = frame[:, :, np.newaxis] frame_t = torch.from_numpy(frame).unsqueeze(0).permute(0, 3, 1, 2).float() # [b, c, h, w] if width != 0 and height != 0: frame_t = F.interpolate(frame_t, size=(height, width), mode=mode) return frame_t def to_numpy(self, tensor): array = np.array(tensor.permute(1, 2, 0)) return array def read_masks(self, mask_dir, width, height, ref_indices): mask_path = sorted(glob.glob(os.path.join(mask_dir, '*.png'))) masks = [] if len(mask_path) < 30: return masks for i in ref_indices: mask = cv2.imread(mask_path[i], 0) mask = mask / 255. mask = cv2.resize(mask, (width, height), cv2.INTER_NEAREST) mask[mask > 0] = 1 mask = mask[:, :, np.newaxis] masks.append(mask) return masks def vis_frames(self, results, masked_frames, frames, video_name, epoch): out_root = self.opt['path']['VAL_IMAGES'] out_dir = os.path.join(out_root, str(epoch), video_name) if not os.path.exists(out_dir): os.makedirs(out_dir) black_column_pixels = 20 results, masked_frames, frames = results.cpu(), masked_frames.cpu(), frames.cpu() T = results.shape[0] for t in range(T): result, masked_frame, frame = results[t], masked_frames[t], frames[t] result = self.to_numpy(result) masked_frame = self.to_numpy(masked_frame) frame = self.to_numpy(frame) result = np.clip(((result + 1) / 2) * 255, 0, 255) frame = np.clip(((frame + 1) / 2) * 255, 0, 255) masked_frame = np.clip(((masked_frame + 1) / 2) * 255, 0, 255) # normalize to [0~255] height, width = result.shape[:2] canvas = np.zeros((height, width * 3 + black_column_pixels * 2, 3)) canvas[:, 0:width, :] = result canvas[:, width + black_column_pixels: 2 * width + black_column_pixels, :] = frame canvas[:, 2 * (width + black_column_pixels):, :] = masked_frame imageio.imwrite(os.path.join(out_dir, 'result_compare_{:05d}.png'.format(t)), canvas)