Spaces:
Runtime error
Runtime error
| # Copyright (c) Microsoft Corporation. | |
| # Licensed under the MIT License. | |
| import argparse | |
| import os | |
| from util import util | |
| import torch | |
| class BaseOptions: | |
| def __init__(self): | |
| self.parser = argparse.ArgumentParser() | |
| self.initialized = False | |
| def initialize(self): | |
| # experiment specifics | |
| self.parser.add_argument( | |
| "--name", | |
| type=str, | |
| default="label2city", | |
| help="name of the experiment. It decides where to store samples and models", | |
| ) | |
| self.parser.add_argument( | |
| "--gpu_ids", type=str, default="0", help="gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU" | |
| ) | |
| self.parser.add_argument( | |
| "--checkpoints_dir", type=str, default="./checkpoints", help="models are saved here" | |
| ) ## note: to add this param when using philly | |
| # self.parser.add_argument('--project_dir', type=str, default='./', help='the project is saved here') ################### This is necessary for philly | |
| self.parser.add_argument( | |
| "--outputs_dir", type=str, default="./outputs", help="models are saved here" | |
| ) ## note: to add this param when using philly Please end with '/' | |
| self.parser.add_argument("--model", type=str, default="pix2pixHD", help="which model to use") | |
| self.parser.add_argument( | |
| "--norm", type=str, default="instance", help="instance normalization or batch normalization" | |
| ) | |
| self.parser.add_argument("--use_dropout", action="store_true", help="use dropout for the generator") | |
| self.parser.add_argument( | |
| "--data_type", | |
| default=32, | |
| type=int, | |
| choices=[8, 16, 32], | |
| help="Supported data type i.e. 8, 16, 32 bit", | |
| ) | |
| self.parser.add_argument("--verbose", action="store_true", default=False, help="toggles verbose") | |
| # input/output sizes | |
| self.parser.add_argument("--batchSize", type=int, default=1, help="input batch size") | |
| self.parser.add_argument("--loadSize", type=int, default=1024, help="scale images to this size") | |
| self.parser.add_argument("--fineSize", type=int, default=512, help="then crop to this size") | |
| self.parser.add_argument("--label_nc", type=int, default=35, help="# of input label channels") | |
| self.parser.add_argument("--input_nc", type=int, default=3, help="# of input image channels") | |
| self.parser.add_argument("--output_nc", type=int, default=3, help="# of output image channels") | |
| # for setting inputs | |
| self.parser.add_argument("--dataroot", type=str, default="./datasets/cityscapes/") | |
| self.parser.add_argument( | |
| "--resize_or_crop", | |
| type=str, | |
| default="scale_width", | |
| help="scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop]", | |
| ) | |
| self.parser.add_argument( | |
| "--serial_batches", | |
| action="store_true", | |
| help="if true, takes images in order to make batches, otherwise takes them randomly", | |
| ) | |
| self.parser.add_argument( | |
| "--no_flip", | |
| action="store_true", | |
| help="if specified, do not flip the images for data argumentation", | |
| ) | |
| self.parser.add_argument("--nThreads", default=2, type=int, help="# threads for loading data") | |
| self.parser.add_argument( | |
| "--max_dataset_size", | |
| type=int, | |
| default=float("inf"), | |
| help="Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.", | |
| ) | |
| # for displays | |
| self.parser.add_argument("--display_winsize", type=int, default=512, help="display window size") | |
| self.parser.add_argument( | |
| "--tf_log", | |
| action="store_true", | |
| help="if specified, use tensorboard logging. Requires tensorflow installed", | |
| ) | |
| # for generator | |
| self.parser.add_argument("--netG", type=str, default="global", help="selects model to use for netG") | |
| self.parser.add_argument("--ngf", type=int, default=64, help="# of gen filters in first conv layer") | |
| self.parser.add_argument("--k_size", type=int, default=3, help="# kernel size conv layer") | |
| self.parser.add_argument("--use_v2", action="store_true", help="use DCDCv2") | |
| self.parser.add_argument("--mc", type=int, default=1024, help="# max channel") | |
| self.parser.add_argument("--start_r", type=int, default=3, help="start layer to use resblock") | |
| self.parser.add_argument( | |
| "--n_downsample_global", type=int, default=4, help="number of downsampling layers in netG" | |
| ) | |
| self.parser.add_argument( | |
| "--n_blocks_global", | |
| type=int, | |
| default=9, | |
| help="number of residual blocks in the global generator network", | |
| ) | |
| self.parser.add_argument( | |
| "--n_blocks_local", | |
| type=int, | |
| default=3, | |
| help="number of residual blocks in the local enhancer network", | |
| ) | |
| self.parser.add_argument( | |
| "--n_local_enhancers", type=int, default=1, help="number of local enhancers to use" | |
| ) | |
| self.parser.add_argument( | |
| "--niter_fix_global", | |
| type=int, | |
| default=0, | |
| help="number of epochs that we only train the outmost local enhancer", | |
| ) | |
| self.parser.add_argument( | |
| "--load_pretrain", | |
| type=str, | |
| default="", | |
| help="load the pretrained model from the specified location", | |
| ) | |
| # for instance-wise features | |
| self.parser.add_argument( | |
| "--no_instance", action="store_true", help="if specified, do *not* add instance map as input" | |
| ) | |
| self.parser.add_argument( | |
| "--instance_feat", | |
| action="store_true", | |
| help="if specified, add encoded instance features as input", | |
| ) | |
| self.parser.add_argument( | |
| "--label_feat", action="store_true", help="if specified, add encoded label features as input" | |
| ) | |
| self.parser.add_argument("--feat_num", type=int, default=3, help="vector length for encoded features") | |
| self.parser.add_argument( | |
| "--load_features", action="store_true", help="if specified, load precomputed feature maps" | |
| ) | |
| self.parser.add_argument( | |
| "--n_downsample_E", type=int, default=4, help="# of downsampling layers in encoder" | |
| ) | |
| self.parser.add_argument( | |
| "--nef", type=int, default=16, help="# of encoder filters in the first conv layer" | |
| ) | |
| self.parser.add_argument("--n_clusters", type=int, default=10, help="number of clusters for features") | |
| # diy | |
| self.parser.add_argument("--self_gen", action="store_true", help="self generate") | |
| self.parser.add_argument( | |
| "--mapping_n_block", type=int, default=3, help="number of resblock in mapping" | |
| ) | |
| self.parser.add_argument("--map_mc", type=int, default=64, help="max channel of mapping") | |
| self.parser.add_argument("--kl", type=float, default=0, help="KL Loss") | |
| self.parser.add_argument( | |
| "--load_pretrainA", | |
| type=str, | |
| default="", | |
| help="load the pretrained model from the specified location", | |
| ) | |
| self.parser.add_argument( | |
| "--load_pretrainB", | |
| type=str, | |
| default="", | |
| help="load the pretrained model from the specified location", | |
| ) | |
| self.parser.add_argument("--feat_gan", action="store_true") | |
| self.parser.add_argument("--no_cgan", action="store_true") | |
| self.parser.add_argument("--map_unet", action="store_true") | |
| self.parser.add_argument("--map_densenet", action="store_true") | |
| self.parser.add_argument("--fcn", action="store_true") | |
| self.parser.add_argument("--is_image", action="store_true", help="train image recon only pair data") | |
| self.parser.add_argument("--label_unpair", action="store_true") | |
| self.parser.add_argument("--mapping_unpair", action="store_true") | |
| self.parser.add_argument("--unpair_w", type=float, default=1.0) | |
| self.parser.add_argument("--pair_num", type=int, default=-1) | |
| self.parser.add_argument("--Gan_w", type=float, default=1) | |
| self.parser.add_argument("--feat_dim", type=int, default=-1) | |
| self.parser.add_argument("--abalation_vae_len", type=int, default=-1) | |
| ######################### useless, just to cooperate with docker | |
| self.parser.add_argument("--gpu", type=str) | |
| self.parser.add_argument("--dataDir", type=str) | |
| self.parser.add_argument("--modelDir", type=str) | |
| self.parser.add_argument("--logDir", type=str) | |
| self.parser.add_argument("--data_dir", type=str) | |
| self.parser.add_argument("--use_skip_model", action="store_true") | |
| self.parser.add_argument("--use_segmentation_model", action="store_true") | |
| self.parser.add_argument("--spatio_size", type=int, default=64) | |
| self.parser.add_argument("--test_random_crop", action="store_true") | |
| ########################## | |
| self.parser.add_argument("--contain_scratch_L", action="store_true") | |
| self.parser.add_argument( | |
| "--mask_dilation", type=int, default=0 | |
| ) ## Don't change the input, only dilation the mask | |
| self.parser.add_argument( | |
| "--irregular_mask", type=str, default="", help="This is the root of the mask" | |
| ) | |
| self.parser.add_argument( | |
| "--mapping_net_dilation", | |
| type=int, | |
| default=1, | |
| help="This parameter is the dilation size of the translation net", | |
| ) | |
| self.parser.add_argument( | |
| "--VOC", type=str, default="VOC_RGB_JPEGImages.bigfile", help="The root of VOC dataset" | |
| ) | |
| self.parser.add_argument("--non_local", type=str, default="", help="which non_local setting") | |
| self.parser.add_argument( | |
| "--NL_fusion_method", | |
| type=str, | |
| default="add", | |
| help="how to fuse the origin feature and nl feature", | |
| ) | |
| self.parser.add_argument( | |
| "--NL_use_mask", action="store_true", help="If use mask while using Non-local mapping model" | |
| ) | |
| self.parser.add_argument( | |
| "--correlation_renormalize", | |
| action="store_true", | |
| help="Since after mask out the correlation matrix(which is softmaxed), the sum is not 1 any more, enable this param to re-weight", | |
| ) | |
| self.parser.add_argument("--Smooth_L1", action="store_true", help="Use L1 Loss in image level") | |
| self.parser.add_argument( | |
| "--face_restore_setting", type=int, default=1, help="This is for the aligned face restoration" | |
| ) | |
| self.parser.add_argument("--face_clean_url", type=str, default="") | |
| self.parser.add_argument("--syn_input_url", type=str, default="") | |
| self.parser.add_argument("--syn_gt_url", type=str, default="") | |
| self.parser.add_argument( | |
| "--test_on_synthetic", | |
| action="store_true", | |
| help="If you want to test on the synthetic data, enable this parameter", | |
| ) | |
| self.parser.add_argument("--use_SN", action="store_true", help="Add SN to every parametric layer") | |
| self.parser.add_argument( | |
| "--use_two_stage_mapping", action="store_true", help="choose the model which uses two stage" | |
| ) | |
| self.parser.add_argument("--L1_weight", type=float, default=10.0) | |
| self.parser.add_argument("--softmax_temperature", type=float, default=1.0) | |
| self.parser.add_argument( | |
| "--patch_similarity", | |
| action="store_true", | |
| help="Enable this denotes using 3*3 patch to calculate similarity", | |
| ) | |
| self.parser.add_argument( | |
| "--use_self", | |
| action="store_true", | |
| help="Enable this denotes that while constructing the new feature maps, using original feature (diagonal == 1)", | |
| ) | |
| self.parser.add_argument("--use_own_dataset", action="store_true") | |
| self.parser.add_argument( | |
| "--test_hole_two_folders", | |
| action="store_true", | |
| help="Enable this parameter means test the restoration with inpainting given twp folders which are mask and old respectively", | |
| ) | |
| self.parser.add_argument( | |
| "--no_hole", | |
| action="store_true", | |
| help="While test the full_model on non_scratch data, do not add random mask into the real old photos", | |
| ) ## Only for testing | |
| self.parser.add_argument( | |
| "--random_hole", | |
| action="store_true", | |
| help="While training the full model, 50% probability add hole", | |
| ) | |
| self.parser.add_argument("--NL_res", action="store_true", help="NL+Resdual Block") | |
| self.parser.add_argument("--image_L1", action="store_true", help="Image level loss: L1") | |
| self.parser.add_argument( | |
| "--hole_image_no_mask", | |
| action="store_true", | |
| help="while testing, give hole image but not give the mask", | |
| ) | |
| self.parser.add_argument( | |
| "--down_sample_degradation", | |
| action="store_true", | |
| help="down_sample the image only, corresponds to [down_sample_face]", | |
| ) | |
| self.parser.add_argument( | |
| "--norm_G", type=str, default="spectralinstance", help="The norm type of Generator" | |
| ) | |
| self.parser.add_argument( | |
| "--init_G", | |
| type=str, | |
| default="xavier", | |
| help="normal|xavier|xavier_uniform|kaiming|orthogonal|none", | |
| ) | |
| self.parser.add_argument("--use_new_G", action="store_true") | |
| self.parser.add_argument("--use_new_D", action="store_true") | |
| self.parser.add_argument( | |
| "--only_voc", action="store_true", help="test the trianed celebA face model using VOC face" | |
| ) | |
| self.parser.add_argument( | |
| "--cosin_similarity", | |
| action="store_true", | |
| help="For non-local, using cosin to calculate the similarity", | |
| ) | |
| self.parser.add_argument( | |
| "--downsample_mode", | |
| type=str, | |
| default="nearest", | |
| help="For partial non-local, choose how to downsample the mask", | |
| ) | |
| self.parser.add_argument("--mapping_exp",type=int,default=0,help='Default 0: original PNL|1: Multi-Scale Patch Attention') | |
| self.parser.add_argument("--inference_optimize",action='store_true',help='optimize the memory cost') | |
| self.initialized = True | |
| def parse(self, save=True): | |
| if not self.initialized: | |
| self.initialize() | |
| self.opt = self.parser.parse_args() | |
| self.opt.isTrain = self.isTrain # train or test | |
| str_ids = self.opt.gpu_ids.split(",") | |
| self.opt.gpu_ids = [] | |
| for str_id in str_ids: | |
| int_id = int(str_id) | |
| if int_id >= 0: | |
| self.opt.gpu_ids.append(int_id) | |
| # set gpu ids | |
| if len(self.opt.gpu_ids) > 0: | |
| # pass | |
| torch.cuda.set_device(self.opt.gpu_ids[0]) | |
| args = vars(self.opt) | |
| # print('------------ Options -------------') | |
| # for k, v in sorted(args.items()): | |
| # print('%s: %s' % (str(k), str(v))) | |
| # print('-------------- End ----------------') | |
| # save to the disk | |
| expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name) | |
| util.mkdirs(expr_dir) | |
| if save and not self.opt.continue_train: | |
| file_name = os.path.join(expr_dir, "opt.txt") | |
| with open(file_name, "wt") as opt_file: | |
| opt_file.write("------------ Options -------------\n") | |
| for k, v in sorted(args.items()): | |
| opt_file.write("%s: %s\n" % (str(k), str(v))) | |
| opt_file.write("-------------- End ----------------\n") | |
| return self.opt | |