Rocky1/SadTalker
0
1"""This script defines the visualizer for Deep3DFaceRecon_pytorch2"""3 4import numpy as np5import os6import sys7import ntpath8import time9from . import util, html10from subprocess import Popen, PIPE11from torch.utils.tensorboard import SummaryWriter12 13def save_images(webpage, visuals, image_path, aspect_ratio=1.0, width=256):14 """Save images to the disk.15 16 Parameters:17 webpage (the HTML class) -- the HTML webpage class that stores these imaegs (see html.py for more details)18 visuals (OrderedDict) -- an ordered dictionary that stores (name, images (either tensor or numpy) ) pairs19 image_path (str) -- the string is used to create image paths20 aspect_ratio (float) -- the aspect ratio of saved images21 width (int) -- the images will be resized to width x width22 23 This function will save images stored in 'visuals' to the HTML file specified by 'webpage'.24 """25 image_dir = webpage.get_image_dir()26 short_path = ntpath.basename(image_path[0])27 name = os.path.splitext(short_path)[0]28 29 webpage.add_header(name)30 ims, txts, links = [], [], []31 32 for label, im_data in visuals.items():33 im = util.tensor2im(im_data)34 image_name = '%s/%s.png' % (label, name)35 os.makedirs(os.path.join(image_dir, label), exist_ok=True)36 save_path = os.path.join(image_dir, image_name)37 util.save_image(im, save_path, aspect_ratio=aspect_ratio)38 ims.append(image_name)39 txts.append(label)40 links.append(image_name)41 webpage.add_images(ims, txts, links, width=width)42 43 44class Visualizer():45 """This class includes several functions that can display/save images and print/save logging information.46 47 It uses a Python library tensprboardX for display, and a Python library 'dominate' (wrapped in 'HTML') for creating HTML files with images.48 """49 50 def __init__(self, opt):51 """Initialize the Visualizer class52 53 Parameters:54 opt -- stores all the experiment flags; needs to be a subclass of BaseOptions55 Step 1: Cache the training/test options56 Step 2: create a tensorboard writer57 Step 3: create an HTML object for saveing HTML filters58 Step 4: create a logging file to store training losses59 """60 self.opt = opt # cache the option61 self.use_html = opt.isTrain and not opt.no_html62 self.writer = SummaryWriter(os.path.join(opt.checkpoints_dir, 'logs', opt.name))63 self.win_size = opt.display_winsize64 self.name = opt.name65 self.saved = False66 if self.use_html: # create an HTML object at <checkpoints_dir>/web/; images will be saved under <checkpoints_dir>/web/images/67 self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')68 self.img_dir = os.path.join(self.web_dir, 'images')69 print('create web directory %s...' % self.web_dir)70 util.mkdirs([self.web_dir, self.img_dir])71 # create a logging file to store training losses72 self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')73 with open(self.log_name, "a") as log_file:74 now = time.strftime("%c")75 log_file.write('================ Training Loss (%s) ================\n' % now)76 77 def reset(self):78 """Reset the self.saved status"""79 self.saved = False80 81 82 def display_current_results(self, visuals, total_iters, epoch, save_result):83 """Display current results on tensorboad; save current results to an HTML file.84 85 Parameters:86 visuals (OrderedDict) - - dictionary of images to display or save87 total_iters (int) -- total iterations88 epoch (int) - - the current epoch89 save_result (bool) - - if save the current results to an HTML file90 """91 for label, image in visuals.items():92 self.writer.add_image(label, util.tensor2im(image), total_iters, dataformats='HWC')93 94 if self.use_html and (save_result or not self.saved): # save images to an HTML file if they haven't been saved.95 self.saved = True96 # save images to the disk97 for label, image in visuals.items():98 image_numpy = util.tensor2im(image)99 img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.png' % (epoch, label))100 util.save_image(image_numpy, img_path)101 102 # update website103 webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, refresh=0)104 for n in range(epoch, 0, -1):105 webpage.add_header('epoch [%d]' % n)106 ims, txts, links = [], [], []107 108 for label, image_numpy in visuals.items():109 image_numpy = util.tensor2im(image)110 img_path = 'epoch%.3d_%s.png' % (n, label)111 ims.append(img_path)112 txts.append(label)113 links.append(img_path)114 webpage.add_images(ims, txts, links, width=self.win_size)115 webpage.save()116 117 def plot_current_losses(self, total_iters, losses):118 # G_loss_collection = {}119 # D_loss_collection = {}120 # for name, value in losses.items():121 # if 'G' in name or 'NCE' in name or 'idt' in name:122 # G_loss_collection[name] = value123 # else:124 # D_loss_collection[name] = value125 # self.writer.add_scalars('G_collec', G_loss_collection, total_iters)126 # self.writer.add_scalars('D_collec', D_loss_collection, total_iters)127 for name, value in losses.items():128 self.writer.add_scalar(name, value, total_iters)129 130 # losses: same format as |losses| of plot_current_losses131 def print_current_losses(self, epoch, iters, losses, t_comp, t_data):132 """print current losses on console; also save the losses to the disk133 134 Parameters:135 epoch (int) -- current epoch136 iters (int) -- current training iteration during this epoch (reset to 0 at the end of every epoch)137 losses (OrderedDict) -- training losses stored in the format of (name, float) pairs138 t_comp (float) -- computational time per data point (normalized by batch_size)139 t_data (float) -- data loading time per data point (normalized by batch_size)140 """141 message = '(epoch: %d, iters: %d, time: %.3f, data: %.3f) ' % (epoch, iters, t_comp, t_data)142 for k, v in losses.items():143 message += '%s: %.3f ' % (k, v)144 145 print(message) # print the message146 with open(self.log_name, "a") as log_file:147 log_file.write('%s\n' % message) # save the message148 149 150class MyVisualizer:151 def __init__(self, opt):152 """Initialize the Visualizer class153 154 Parameters:155 opt -- stores all the experiment flags; needs to be a subclass of BaseOptions156 Step 1: Cache the training/test options157 Step 2: create a tensorboard writer158 Step 3: create an HTML object for saveing HTML filters159 Step 4: create a logging file to store training losses160 """161 self.opt = opt # cache the optio162 self.name = opt.name163 self.img_dir = os.path.join(opt.checkpoints_dir, opt.name, 'results')164 165 if opt.phase != 'test':166 self.writer = SummaryWriter(os.path.join(opt.checkpoints_dir, opt.name, 'logs'))167 # create a logging file to store training losses168 self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')169 with open(self.log_name, "a") as log_file:170 now = time.strftime("%c")171 log_file.write('================ Training Loss (%s) ================\n' % now)172 173 174 def display_current_results(self, visuals, total_iters, epoch, dataset='train', save_results=False, count=0, name=None,175 add_image=True):176 """Display current results on tensorboad; save current results to an HTML file.177 178 Parameters:179 visuals (OrderedDict) - - dictionary of images to display or save180 total_iters (int) -- total iterations181 epoch (int) - - the current epoch182 dataset (str) - - 'train' or 'val' or 'test'183 """184 # if (not add_image) and (not save_results): return185 186 for label, image in visuals.items():187 for i in range(image.shape[0]):188 image_numpy = util.tensor2im(image[i])189 if add_image:190 self.writer.add_image(label + '%s_%02d'%(dataset, i + count),191 image_numpy, total_iters, dataformats='HWC')192 193 if save_results:194 save_path = os.path.join(self.img_dir, dataset, 'epoch_%s_%06d'%(epoch, total_iters))195 if not os.path.isdir(save_path):196 os.makedirs(save_path)197 198 if name is not None:199 img_path = os.path.join(save_path, '%s.png' % name)200 else:201 img_path = os.path.join(save_path, '%s_%03d.png' % (label, i + count))202 util.save_image(image_numpy, img_path)203 204 205 def plot_current_losses(self, total_iters, losses, dataset='train'):206 for name, value in losses.items():207 self.writer.add_scalar(name + '/%s'%dataset, value, total_iters)208 209 # losses: same format as |losses| of plot_current_losses210 def print_current_losses(self, epoch, iters, losses, t_comp, t_data, dataset='train'):211 """print current losses on console; also save the losses to the disk212 213 Parameters:214 epoch (int) -- current epoch215 iters (int) -- current training iteration during this epoch (reset to 0 at the end of every epoch)216 losses (OrderedDict) -- training losses stored in the format of (name, float) pairs217 t_comp (float) -- computational time per data point (normalized by batch_size)218 t_data (float) -- data loading time per data point (normalized by batch_size)219 """220 message = '(dataset: %s, epoch: %d, iters: %d, time: %.3f, data: %.3f) ' % (221 dataset, epoch, iters, t_comp, t_data)222 for k, v in losses.items():223 message += '%s: %.3f ' % (k, v)224 225 print(message) # print the message226 with open(self.log_name, "a") as log_file:227 log_file.write('%s\n' % message) # save the message228