RabbitRUI/ruispace

1"""This script contains base options for Deep3DFaceRecon_pytorch2"""3 4import argparse5import os6from util import util7import numpy as np8import torch9import face3d.models as models10import face3d.data as data11 12 13class BaseOptions():14    """This class defines options used during both training and test time.15 16    It also implements several helper functions such as parsing, printing, and saving the options.17    It also gathers additional options defined in <modify_commandline_options> functions in both dataset class and model class.18    """19 20    def __init__(self, cmd_line=None):21        """Reset the class; indicates the class hasn't been initailized"""22        self.initialized = False23        self.cmd_line = None24        if cmd_line is not None:25            self.cmd_line = cmd_line.split()26 27    def initialize(self, parser):28        """Define the common options that are used in both training and test."""29        # basic parameters30        parser.add_argument('--name', type=str, default='face_recon', help='name of the experiment. It decides where to store samples and models')31        parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')32        parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')33        parser.add_argument('--vis_batch_nums', type=float, default=1, help='batch nums of images for visulization')34        parser.add_argument('--eval_batch_nums', type=float, default=float('inf'), help='batch nums of images for evaluation')35        parser.add_argument('--use_ddp', type=util.str2bool, nargs='?', const=True, default=True, help='whether use distributed data parallel')36        parser.add_argument('--ddp_port', type=str, default='12355', help='ddp port')37        parser.add_argument('--display_per_batch', type=util.str2bool, nargs='?', const=True, default=True, help='whether use batch to show losses')38        parser.add_argument('--add_image', type=util.str2bool, nargs='?', const=True, default=True, help='whether add image to tensorboard')39        parser.add_argument('--world_size', type=int, default=1, help='batch nums of images for evaluation')40 41        # model parameters42        parser.add_argument('--model', type=str, default='facerecon', help='chooses which model to use.')43 44        # additional parameters45        parser.add_argument('--epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')46        parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information')47        parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{load_size}')48 49        self.initialized = True50        return parser51 52    def gather_options(self):53        """Initialize our parser with basic options(only once).54        Add additional model-specific and dataset-specific options.55        These options are defined in the <modify_commandline_options> function56        in model and dataset classes.57        """58        if not self.initialized:  # check if it has been initialized59            parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)60            parser = self.initialize(parser)61 62        # get the basic options63        if self.cmd_line is None:64            opt, _ = parser.parse_known_args()65        else:66            opt, _ = parser.parse_known_args(self.cmd_line)67 68        # set cuda visible devices69        os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_ids70 71        # modify model-related parser options72        model_name = opt.model73        model_option_setter = models.get_option_setter(model_name)74        parser = model_option_setter(parser, self.isTrain)75        if self.cmd_line is None:76            opt, _ = parser.parse_known_args()  # parse again with new defaults77        else:78            opt, _ = parser.parse_known_args(self.cmd_line)  # parse again with new defaults79 80        # modify dataset-related parser options81        if opt.dataset_mode:82            dataset_name = opt.dataset_mode83            dataset_option_setter = data.get_option_setter(dataset_name)84            parser = dataset_option_setter(parser, self.isTrain)85 86        # save and return the parser87        self.parser = parser88        if self.cmd_line is None:89            return parser.parse_args()90        else:91            return parser.parse_args(self.cmd_line)92 93    def print_options(self, opt):94        """Print and save options95 96        It will print both current options and default values(if different).97        It will save options into a text file / [checkpoints_dir] / opt.txt98        """99        message = ''100        message += '----------------- Options ---------------\n'101        for k, v in sorted(vars(opt).items()):102            comment = ''103            default = self.parser.get_default(k)104            if v != default:105                comment = '\t[default: %s]' % str(default)106            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)107        message += '----------------- End -------------------'108        print(message)109 110        # save to the disk111        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)112        util.mkdirs(expr_dir)113        file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))114        try:115            with open(file_name, 'wt') as opt_file:116                opt_file.write(message)117                opt_file.write('\n')118        except PermissionError as error:119            print("permission error {}".format(error))120            pass121 122    def parse(self):123        """Parse our options, create checkpoints directory suffix, and set up gpu device."""124        opt = self.gather_options()125        opt.isTrain = self.isTrain   # train or test126 127        # process opt.suffix128        if opt.suffix:129            suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''130            opt.name = opt.name + suffix131 132 133        # set gpu ids134        str_ids = opt.gpu_ids.split(',')135        gpu_ids = []136        for str_id in str_ids:137            id = int(str_id)138            if id >= 0:139                gpu_ids.append(id)140        opt.world_size = len(gpu_ids)141        # if len(opt.gpu_ids) > 0:142        #     torch.cuda.set_device(gpu_ids[0])143        if opt.world_size == 1:144            opt.use_ddp = False145 146        if opt.phase != 'test':147            # set continue_train automatically148            if opt.pretrained_name is None:149                model_dir = os.path.join(opt.checkpoints_dir, opt.name)150            else:151                model_dir = os.path.join(opt.checkpoints_dir, opt.pretrained_name)152            if os.path.isdir(model_dir):153                model_pths = [i for i in os.listdir(model_dir) if i.endswith('pth')]154                if os.path.isdir(model_dir) and len(model_pths) != 0:155                    opt.continue_train= True156        157            # update the latest epoch count158            if opt.continue_train:159                if opt.epoch == 'latest':160                    epoch_counts = [int(i.split('.')[0].split('_')[-1]) for i in model_pths if 'latest' not in i]161                    if len(epoch_counts) != 0:162                        opt.epoch_count = max(epoch_counts) + 1163                else:164                    opt.epoch_count = int(opt.epoch) + 1165                    166 167        self.print_options(opt)168        self.opt = opt169        return self.opt170