1# Licensed to the Apache Software Foundation (ASF) under one 2# or more contributor license agreements. See the NOTICE file 3# distributed with this work for additional information 4# regarding copyright ownership. The ASF licenses this file 5# to you under the Apache License, Version 2.0 (the 6# "License"); you may not use this file except in compliance 7# with the License. You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, 12# software distributed under the License is distributed on an 13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14# KIND, either express or implied. See the License for the 15# specific language governing permissions and limitations 16# under the License. 17 18import argparse 19import os 20 21class Options(): 22 def __init__(self): 23 self.parser = argparse.ArgumentParser(description="parser for MXNet-Gluon-Style-Transfer") 24 subparsers = self.parser.add_subparsers(title="subcommands", dest="subcommand") 25 26 # training args 27 train_arg = subparsers.add_parser("train", 28 help="parser for training arguments") 29 train_arg.add_argument("--ngf", type=int, default=128, 30 help="number of generator filter channels, default 128") 31 train_arg.add_argument("--epochs", type=int, default=4, 32 help="number of training epochs, default is 2") 33 train_arg.add_argument("--batch-size", type=int, default=4, 34 help="batch size for training, default is 4") 35 train_arg.add_argument("--dataset", type=str, default="dataset/", 36 help="path to training dataset, the path should point to a folder " 37 "containing another folder with all the training images") 38 train_arg.add_argument("--style-folder", type=str, default="images/styles/", 39 help="path to style-folder") 40 train_arg.add_argument("--save-model-dir", type=str, default="models/", 41 help="path to folder where trained model will be saved.") 42 train_arg.add_argument("--image-size", type=int, default=256, 43 help="size of training images, default is 256 X 256") 44 train_arg.add_argument("--style-size", type=int, default=512, 45 help="size of style-image, default is the original size of style image") 46 train_arg.add_argument("--cuda", type=int, default=1, 47 help="set it to 1 for running on GPU, 0 for CPU") 48 train_arg.add_argument("--seed", type=int, default=42, 49 help="random seed for training") 50 train_arg.add_argument("--content-weight", type=float, default=1.0, 51 help="weight for content-loss, default is 1.0") 52 train_arg.add_argument("--style-weight", type=float, default=5.0, 53 help="weight for style-loss, default is 5.0") 54 train_arg.add_argument("--lr", type=float, default=1e-3, 55 help="learning rate, default is 0.001") 56 train_arg.add_argument("--log-interval", type=int, default=500, 57 help="number of images after which the training loss is logged, default is 500") 58 train_arg.add_argument("--resume", type=str, default=None, 59 help="resume if needed") 60 61 # optim args (Gatys CVPR 2016) 62 optim_arg = subparsers.add_parser("optim", 63 help="parser for optimization arguments") 64 optim_arg.add_argument("--iters", type=int, default=500, 65 help="number of training iterations, default is 500") 66 optim_arg.add_argument("--content-image", type=str, default="images/content/venice-boat.jpg", 67 help="path to content image you want to stylize") 68 optim_arg.add_argument("--style-image", type=str, default="images/9styles/candy.jpg", 69 help="path to style-image") 70 optim_arg.add_argument("--content-size", type=int, default=512, 71 help="factor for scaling down the content image") 72 optim_arg.add_argument("--style-size", type=int, default=512, 73 help="size of style-image, default is the original size of style image") 74 optim_arg.add_argument("--output-image", type=str, default="output.jpg", 75 help="path for saving the output image") 76 optim_arg.add_argument("--cuda", type=int, default=1, 77 help="set it to 1 for running on GPU, 0 for CPU") 78 optim_arg.add_argument("--content-weight", type=float, default=1.0, 79 help="weight for content-loss, default is 1.0") 80 optim_arg.add_argument("--style-weight", type=float, default=5.0, 81 help="weight for style-loss, default is 5.0") 82 optim_arg.add_argument("--lr", type=float, default=1e1, 83 help="learning rate, default is 0.001") 84 optim_arg.add_argument("--log-interval", type=int, default=50, 85 help="number of images after which the training loss is logged, default is 50") 86 87 # evaluation args 88 eval_arg = subparsers.add_parser("eval", help="parser for evaluation/stylizing arguments") 89 eval_arg.add_argument("--ngf", type=int, default=128, 90 help="number of generator filter channels, default 128") 91 eval_arg.add_argument("--content-image", type=str, required=True, 92 help="path to content image you want to stylize") 93 eval_arg.add_argument("--style-image", type=str, default="images/9styles/candy.jpg", 94 help="path to style-image") 95 eval_arg.add_argument("--content-size", type=int, default=512, 96 help="factor for scaling down the content image") 97 eval_arg.add_argument("--style-size", type=int, default=512, 98 help="size of style-image, default is the original size of style image") 99 eval_arg.add_argument("--style-folder", type=str, default="images/9styles/", 100 help="path to style-folder") 101 eval_arg.add_argument("--output-image", type=str, default="output.jpg", 102 help="path for saving the output image") 103 eval_arg.add_argument("--model", type=str, required=True, 104 help="saved model to be used for stylizing the image") 105 eval_arg.add_argument("--cuda", type=int, default=1, 106 help="set it to 1 for running on GPU, 0 for CPU") 107 108 def parse(self): 109 return self.parser.parse_args() 110