File size: 4,906 Bytes
89ce6b3 |
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 |
import argparse
import tensorflow as tf
import os
from utils import *
from AUGAN import AUGAN
from ops import *
import time
parser = argparse.ArgumentParser(description="")
parser.add_argument(
"--dataset_dir", dest="dataset_dir", default="bdd100k", help="path of the dataset"
)
parser.add_argument(
"--experiment_name",
dest="experiment_name",
type=str,
default="bdd_exp",
help="name of experiment",
)
parser.add_argument("--epoch", dest="epoch", type=int, default=20, help="# of epoch")
parser.add_argument(
"--epoch_step",
dest="epoch_step",
type=int,
default=10,
help="# of epoch to decay lr",
)
parser.add_argument(
"--batch_size", dest="batch_size", type=int, default=1, help="# images in batch"
)
parser.add_argument(
"--train_size",
dest="train_size",
type=int,
default=1e8,
help="# images used to train",
)
parser.add_argument(
"--load_size",
dest="load_size",
type=int,
default=286,
help="scale images to this size",
)
parser.add_argument(
"--fine_size",
dest="fine_size",
type=int,
default=256,
help="then crop to this size",
)
parser.add_argument(
"--ngf",
dest="ngf",
type=int,
default=64,
help="# of gen filters in first conv layer",
)
parser.add_argument(
"--ndf",
dest="ndf",
type=int,
default=64,
help="# of discri filters in first conv layer",
)
parser.add_argument(
"--n_d", dest="n_d", type=int, default=2, help="# of discriminators"
)
parser.add_argument(
"--n_scale", dest="n_scale", type=int, default=2, help="# of scales"
)
parser.add_argument(
"--gpu", dest="gpu", type=int, default=0, help="# index of gpu device"
)
parser.add_argument(
"--input_nc", dest="input_nc", type=int, default=3, help="# of input image channels"
)
parser.add_argument(
"--output_nc",
dest="output_nc",
type=int,
default=3,
help="# of output image channels",
)
parser.add_argument(
"--lr", dest="lr", type=float, default=0.0002, help="initial learning rate for adam"
)
parser.add_argument(
"--beta1", dest="beta1", type=float, default=0.5, help="momentum term of adam"
)
parser.add_argument(
"--which_direction", dest="which_direction", default="AtoB", help="AtoB or BtoA "
)
parser.add_argument("--phase", dest="phase", default="test", help="train, test")
parser.add_argument(
"--save_freq",
dest="save_freq",
type=int,
default=1000,
help="save a model every save_freq iterations",
)
parser.add_argument(
"--print_freq",
dest="print_freq",
type=int,
default=100,
help="print the debug information every print_freq iterations",
)
parser.add_argument(
"--L1_lambda",
dest="L1_lambda",
type=float,
default=10.0,
help="weight on L1 term in objective",
)
parser.add_argument(
"--conf_lambda",
dest="conf_lambda",
type=float,
default=1.0,
help="weight on L1 term in objective",
)
parser.add_argument(
"--use_resnet",
dest="use_resnet",
type=bool,
default=True,
help="generation network using reidule block",
)
parser.add_argument(
"--use_lsgan",
dest="use_lsgan",
type=bool,
default=True,
help="gan loss defined in lsgan",
)
parser.add_argument(
"--use_uncertainty",
dest="use_uncertainty",
type=bool,
default=True,
help="max size of image pool, 0 means do not use image pool",
)
parser.add_argument(
"--max_size",
dest="max_size",
type=int,
default=50,
help="max size of image pool, 0 means do not use image pool",
)
parser.add_argument(
"--continue_train",
dest="continue_train",
type=bool,
default=False,
help="if continue training, load the latest model: 1: true, 0: false",
)
parser.add_argument(
"--save_conf",
dest="save_conf",
type=bool,
default=False,
help="save conf map in test phase",
)
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
def main(_):
set_path(args, args.experiment_name)
tfconfig = tf.compat.v1.ConfigProto(allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
with tf.compat.v1.Session(config=tfconfig) as sess:
model = AUGAN(sess, args)
# show_all_variables()
# model.train(args) if args.phase == 'train' \
# else model.test(args)
if args.phase == "train":
model.train(args)
elif args.phase == "test":
model.test(args)
elif args.phase == "convert":
model.convert_image(args, "inf_data/b1ca2e5d-84cf9134.jpg", "out")
else:
raise Exception("Give a phase")
if __name__ == "__main__":
tf.compat.v1.app.run()
|