File size: 1,822 Bytes
89cf463 |
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 |
import argparse
class HifiTestOptions:
def __init__(self):
self.parser = argparse.ArgumentParser()
self.initialized = False
def initialize(self):
self.parser.add_argument('--model_name', type=str, default='er8_bs1', help='er8_bs1')
self.parser.add_argument('--input_size', type=int, default='256')
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('--src_img_path', type=str, default='data/source/elon-musk1.jpg')
self.parser.add_argument('--video_path', type=str, default='data/source/demo.mp4')
self.parser.add_argument('--video_to_1080p', action="store_true", help='change video resolution to 1080p')
self.parser.add_argument('--mode', type=str, default='default', help='default merge')
self.parser.add_argument('--align_method', type=str, default='68', help='face align method:68 5class')
self.parser.add_argument('--use_gfpgan', action="store_true", help='use gfpgan for sr or not')
self.parser.add_argument('--sr_weight', type=float, default=1.0)
self.parser.add_argument('--use_color_trans', action="store_true", help='use color transfer or not')
self.parser.add_argument('--color_trans_mode', type=str, default='rct', help='rct lct mkl idt sot')
self.parser.add_argument('--output_dir', type=str, default='data/output')
def parse(self, save=True):
if not self.initialized:
self.initialize()
self.opt = self.parser.parse_args()
str_ids = self.opt.gpu_ids.split(',')
self.opt.gpu_ids = []
for str_id in str_ids:
id = int(str_id)
if id >= 0:
self.opt.gpu_ids.append(id)
return self.opt
|