addmosaic_options.py 1.7 KB
Newer Older
HypoX64's avatar
preview  
HypoX64 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
import argparse
import os

class AddOptions():
    def __init__(self):
        self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
        self.initialized = False

    def initialize(self):
        self.parser.add_argument('--use_gpu', action='store_true', help='if true, use gpu')
        self.parser.add_argument('--input_dir', type=str, default='./video_or_image',help='put your videos or images here')
        self.parser.add_argument('--result_dir', type=str, default='./result',help='result will be saved here')
        self.parser.add_argument('--model_dir', type=str, default='./pretrained_models/AddMosaic',
                                help='put pre_train model here')
        self.parser.add_argument('--model_name', type=str, default='hands_128.pth',help='name of model use to Add mosaic')
        self.parser.add_argument('--mosaic_mod', type=str, default='squa_avg',help='type of mosaic -> squa_avg | squa_random | squa_avg_circle_edge | rect_avg')
HypoX64's avatar
HypoX64 已提交
17 18
        self.parser.add_argument('--mosaic_size', type=int, default=30,help='mosaic size')
        self.parser.add_argument('--mask_extend', type=int, default=20,help='more mosaic area')
HypoX64's avatar
preview  
HypoX64 已提交
19 20
        self.parser.add_argument('--mask_threshold', type=int, default=64,help='threshold of recognize mosaic position 0~255')
        self.parser.add_argument('--output_size', type=int, default=0,help='size of output file,if 0 -> origin')
HypoX64's avatar
HypoX64 已提交
21
        self.parser.add_argument('--tempimage_type', type=str, default='jpg',help='type of temp image, png | jpg, png is better but occupy more storage space')
HypoX64's avatar
preview  
HypoX64 已提交
22 23 24 25 26 27 28
        self.initialized = True

    def getparse(self):
        if not self.initialized:
            self.initialize()
        self.opt = self.parser.parse_args()
        return self.opt