<reponame>nsfzyzz/dispersion-score<gh_stars>1-10
import argparse
import auxiliary.my_utils as my_utils
import os
import datetime
import json
import sys
import logging
from termcolor import colored
from easydict import EasyDict
import tensorboardX
import yaml
from os.path import exists, join



def parser():
    parser = argparse.ArgumentParser()

    #Training parameters
    parser.add_argument("--views_search", action="store_true", help="increment views to train")
    parser.add_argument("--mode", type=str, default="object", choices=['viewer', 'object'])
    parser.add_argument("--nviews_train", type=int, default=1, help='num of view per shape for training')
    parser.add_argument("--nviews_test", type=int, default=1, help='num of view per shape for test')
    parser.add_argument("--no_learning", action="store_true", help="Learning mode (batchnorms...)")
    parser.add_argument("--train_only_encoder", action="store_true", help="only train the encoder")
    parser.add_argument('--batch_size', type=int, default=32, help='input batch size')
    parser.add_argument('--batch_size_test', type=int, default=32, help='input batch size')
    parser.add_argument('--parallel', action="store_true", help='use data parallel')
    parser.add_argument('--workers', type=int, help='number of data loading workers', default=8)
    parser.add_argument('--nepoch', type=int, default=150, help='number of epochs to train for')
    parser.add_argument('--start_epoch', type=int, default=0, help='number of epochs to train for')
    #parser.add_argument("--manual_seed", type=int, default=1, help="if seed > 0, then fixed, if < 0 then random")
    parser.add_argument('--lrate', type=float, default=0.001, help='learning rate')
    parser.add_argument('--lr_decay_1', type=int, default=120, help='learning rate decay 1')
    parser.add_argument('--lr_decay_2', type=int, default=140, help='learning rate decay 2')
    parser.add_argument('--lr_decay_3', type=int, default=145, help='learning rate decay 2')
    parser.add_argument("--run_single_eval", action="store_true", help="evaluate a trained network")
    parser.add_argument("--demo", action="store_true", help="run demo autoencoder or single-view")

    # Data
    parser.add_argument('--normalization', type=str, default="UnitBall",
                        choices=['UnitBall', 'BoundingBox', 'Identity'])
    parser.add_argument("--shapenet13", action="store_true", help="Load 13 usual shapenet categories")
    parser.add_argument("--SVR", action="store_true", help="Single_view Reconstruction")
    parser.add_argument("--sample", action="store_false", help="Sample the input pointclouds")
    parser.add_argument('--class_choice', nargs='+', default=["airplane"], type=str)
    parser.add_argument('--number_points', type=int, default=2500, help='Number of point sampled on the object during training, and generated by atlasnet')
    parser.add_argument('--number_points_eval', type=int, default=2500,
                        help='Number of points generated by atlasnet (rounded to the nearest squared number) ')

    parser.add_argument("--img_aug", action="store_true", help="apply image augmentation like random crop")
    parser.add_argument("--img_aug_type", type=str, default='rgb', choices=['grayscale', 'rgb', 'binary', 'color_aug', 'color_aug_random', 'autoaugment', 'geometricaugment', 'geometricaugmentMag', 'coloraugment'])
    parser.add_argument("--autoaug_type", type=str, default='RGB', choices=['ImageNet', 'CIFAR10', 'Seq', 'RGB', 'SVHN', 'None'])
    parser.add_argument("--color_aug_factor", nargs='+', type=float, default=[1.0, 1.0, 1.0, 1.0], help='brightness, contrast, saturation, hue')
    parser.add_argument("--magnitude", default=0, type=int, help="magnitude for coloraug")
    parser.add_argument("--mag_idx", default=0, type=int, help="magnitude level index")
    parser.add_argument("--n_op", default=0, type=int, help="number of operations")
    parser.add_argument("--prob", default=0, type=float, help="prob")

    parser.add_argument("--test_augment", action="store_true", default=False, help='test augment')
    parser.add_argument("--random_rotation", action="store_true", help="apply data augmentation : random rotation")
    parser.add_argument("--data_augmentation_axis_rotation", action="store_true",
                        help="apply data augmentation : axial rotation ")
    parser.add_argument("--data_augmentation_random_flips", action="store_true",
                        help="apply data augmentation : random flips")
    parser.add_argument("--random_translation", action="store_true",
                        help="apply data augmentation :  random translation ")
    parser.add_argument("--anisotropic_scaling", action="store_true",
                        help="apply data augmentation : anisotropic scaling")

    # Save dirs and reload
    parser.add_argument('--id', type=str, default="0", help='training name')
    parser.add_argument('--description', type=str, default="", help='descript this training')
    parser.add_argument('--tbd_dir', type=str, default="log/tensorboard", help='name of the tensorboard folder.')
    parser.add_argument('--train_endsignal', type=str, default="train_end.txt", help='ending signal for scheduler.')
    
    parser.add_argument('--dir_name', type=str, default="log/oneExp", help='name of the log folder.')
    parser.add_argument('--save_pred', action="store_true", help="save prediction points")
    parser.add_argument('--demo_input_path', type=str, default="./doc/pictures/plane_input_demo.png", help='dirname')
    parser.add_argument('--reload_options_path', type=str, default=None, help='options path')
    parser.add_argument('--reload_decoder_path', type=str, default="", help='dirname')
    parser.add_argument('--reload_model_path', type=str, default='', help='optional reload model path')
    parser.add_argument('--reload_optimizer_path', type=str, default='', help='optional reload model path')

    # Network
    parser.add_argument("--encoder", type=str, default="resnet18", help="[resnet18, resnet50]")
    parser.add_argument("--network", type=str, default="atlasnet", help="[psgn, atlasnet]")
    parser.add_argument('--num_layers', type=int, default=2, help='number of hidden MLP Layer')
    parser.add_argument('--hidden_neurons', type=int, default=512, help='number of neurons in each hidden layer')
    parser.add_argument('--loop_per_epoch', type=int, default=1, help='number of data loop per epoch')
    parser.add_argument('--nb_primitives', type=int, default=25, help='number of primitives')
    parser.add_argument('--template_type', type=str, default="SPHERE", choices=["SPHERE", "SQUARE"],
                        help='dim_out_patch')
    parser.add_argument('--multi_gpu', nargs='+', type=int, default=[0], help='Use multiple gpus')
    parser.add_argument("--remove_all_batchNorms", action="store_true", help="Replace all batchnorms by identity")
    parser.add_argument('--bottleneck_size', type=int, default=1024, help='dim_out_patch')
    parser.add_argument('--activation', type=str, default='relu',
                        choices=["relu", "sigmoid", "softplus", "logsigmoid", "softsign", "tanh"], help='dim_out_patch')

    # Loss
    parser.add_argument("--no_metro", action="store_false", help="Compute metro distance")
    parser.add_argument('--no_compile_chamfer', action="store_true", help="compile c++ version chamfer distance")
    
    
    #Metric
    parser.add_argument("--c_method", nargs='+', type=str, default=[])      #['KMedoids', 'AP', 'OPTICS']
    parser.add_argument("--e_method", nargs='+', type=str, default=[])    #' Inertia'
    parser.add_argument("--cluster_k", nargs='+', type=int, default=[])
    parser.add_argument("--perf_pc_list", nargs='+', type=int, default=[], help='list of pref pc')

    
    parser.add_argument("--type", type=str, choices=["points", 'image', "pred", "category"])
    parser.add_argument("--perceptual", action="store_true", default=False)
    parser.add_argument("--perceptual-max-pooling", action="store_true", default=False)
    parser.add_argument("--split", type=str, choices=["train", "test", "pred"])
    parser.add_argument("--metric", type=str, choices=["chamfer", "mse", "perceptual"])
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--seed_list", nargs='+', type=int, default=[], help='list of seeds')
    parser.add_argument("--rsample", type=float, default=0.05)    #0.05
    parser.add_argument("--nsample", type=int,  default=-1)       #0.1
    parser.add_argument("--perceptual_batch_size", type=int, default=32)
    parser.add_argument("--pred_batch_size", type=int, default=64)
    parser.add_argument("--dismat_batch_size", type=int, default=32)
    parser.add_argument("--num_worker", type=int, default=16)
    parser.add_argument("--content_lr_idx", type=int, default=2)
    parser.add_argument("--perf_pc", type=int,  default=50)
    parser.add_argument("--st_ind_row_ratio", type=float, default=0.75)
    parser.add_argument("--ed_ind_row_ratio", type=float, default=1)
    parser.add_argument("--res_folder", type=str, default="dm_part_sscore")
    parser.add_argument("--data_base_dir", type=str, default="../data/cubesphere_1000/")
    parser.add_argument('--train_json_file', type=str, default='cluster_image_aug_2by10_cltsize01.json', help='toy dataset json file for train')
    parser.add_argument('--test_json_file', type=str, default='test_interp_1000.json', help='toy dataset json file for test')
    parser.add_argument('--val_json_file', type=str, default='val_interp_1000.json', help='toy dataset json file for val')
    #parser.add_argument("--prediction_path", type=str, default="prediction path for clustering and Oracle NN method")

    parser.add_argument("--trained_exp_dir", type=str, default="")
    parser.add_argument('--anglelimit_search', action='store_true', default=False, help='use limited yaw rendering')
    parser.add_argument('--rendering_root_dir', type=str, default="",
                                                            help='')
    parser.add_argument('--yawrange', type=int, default=0, help='')
    
    opt = parser.parse_args()
    opt.date = str(datetime.datetime.now())
    now = datetime.datetime.now()
    opt = EasyDict(opt.__dict__)


    # Set-up output directories

    #net_desc = opt.description
    #opt.dir_name = os.path.join(opt.dir_name, net_desc)
    my_utils.check_exist_or_mkdirs(opt.dir_name)


    with open(os.path.join(opt.dir_name, 'opts.yaml'), 'w') as file:
        yaml.dump(dict(opt), file, default_flow_style=False, sort_keys=False)

    # Hard code dimension of the template.
    dim_template_dict = {
        "SQUARE": 2,
        "SPHERE": 3,
    }

    opt.dim_template = dim_template_dict[opt.template_type]


    if opt.run_single_eval and opt.save_pred:
        opt.results_dir = os.path.join(opt.dir_name, 'results')
        my_utils.check_exist_or_mkdirs(opt.results_dir)


    return opt
