Spaces:
Build error
Build error
File size: 3,240 Bytes
6d314be |
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 |
import argparse
from trainer import train
from tester import test
def add_data_options(parser):
group = parser.add_argument_group("dataset")
group.add_argument("--image_dir", type=str, default=None, help="the directory of image")
group.add_argument("--annot_dir", type=str, default=None, help="the directory of annot")
def add_base_options(parser):
group = parser.add_argument_group("base")
group.add_argument("--mode", type=str, default="train", help="train or test")
group.add_argument("--config_name", type=str, default="alignment", help="set configure file name")
group.add_argument('--device_ids', type=str, default="0,1,2,3",
help="set device ids, -1 means use cpu device, >= 0 means use gpu device")
group.add_argument('--data_definition', type=str, default='WFLW', help="COFW, 300W, WFLW")
group.add_argument('--learn_rate', type=float, default=0.001, help='learning rate')
group.add_argument("--batch_size", type=int, default=128, help="the batch size in train process")
group.add_argument('--width', type=int, default=256, help='the width of input image')
group.add_argument('--height', type=int, default=256, help='the height of input image')
def add_train_options(parser):
group = parser.add_argument_group('train')
group.add_argument("--train_num_workers", type=int, default=None, help="the num of workers in train process")
group.add_argument('--loss_func', type=str, default='STARLoss_v2', help="loss function")
group.add_argument("--val_batch_size", type=int, default=None, help="the batch size in val process")
group.add_argument("--val_num_workers", type=int, default=None, help="the num of workers in val process")
def add_eval_options(parser):
group = parser.add_argument_group("eval")
group.add_argument("--pretrained_weight", type=str, default=None,
help="set pretrained model file name, if ignored then train the network without pretrain model")
group.add_argument('--norm_type', type=str, default='default', help='default, ocular, pupil')
group.add_argument('--test_file', type=str, default="test.tsv", help='for wflw, test.tsv/test_xx_metadata.tsv')
def add_starloss_options(parser):
group = parser.add_argument_group('starloss')
group.add_argument('--star_w', type=float, default=1, help="regular loss ratio")
group.add_argument('--star_dist', type=str, default='smoothl1', help='STARLoss distance function')
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Entry Function")
add_base_options(parser)
add_data_options(parser)
add_train_options(parser)
add_eval_options(parser)
add_starloss_options(parser)
args = parser.parse_args()
print(args)
print(
"mode is %s, config_name is %s, pretrained_weight is %s, image_dir is %s, annot_dir is %s, device_ids is %s" % (
args.mode, args.config_name, args.pretrained_weight, args.image_dir, args.annot_dir, args.device_ids))
args.device_ids = list(map(int, args.device_ids.split(",")))
if args.mode == "train":
train(args)
elif args.mode == "test":
test(args)
else:
print("unknown running mode")
|