File size: 16,231 Bytes
7fab858
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
# 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