options.py

import argparse
import os
import time
import numpy as np
from . import util,dsp,plot

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

    def initialize(self):
        # ------------------------Base------------------------
        self.parser.add_argument('--gpu_id', type=int, default=0,help='choose which gpu want to use, 0 | 1 | 2 ...')        
        self.parser.add_argument('--no_cudnn', action='store_true', help='if specified, do not use cudnn')
        self.parser.add_argument('--label', type=str, default='auto',help='number of labels')
        self.parser.add_argument('--input_nc', type=str, default='auto', help='of input channels')
        self.parser.add_argument('--loadsize', type=str, default='auto', help='load data in this size')
        self.parser.add_argument('--finesize', type=str, default='auto', help='crop your data into this size')
        self.parser.add_argument('--label_name', type=str, default='auto',help='name of labels,example:"a,b,c,d,e,f"')
        
        # ------------------------Dataset------------------------
        self.parser.add_argument('--dataset_dir', type=str, default='./datasets/simple_test',help='your dataset path')
        self.parser.add_argument('--save_dir', type=str, default='./checkpoints/',help='save checkpoints')
        self.parser.add_argument('--load_thread', type=int, default=8,help='how many threads when load data')  
        self.parser.add_argument('--normliaze', type=str, default='5_95', help='mode of normliaze, 5_95 | maxmin | None')      
        self.parser.add_argument('--k_fold', type=int, default=0,help='fold_num of k-fold.If 0 or 1, no k-fold and cut 0.8 to train and other to eval')
        """--fold_index
        5-fold:
        Cut dataset into sub-set using index , and then run k-fold with sub-set
        If input 'auto', it will shuffle dataset and then cut dataset equally
        If input: [2,4,6,7]
        when len(dataset) == 10
        sub-set: dataset[0:2],dataset[2:4],dataset[4:6],dataset[6:7],dataset[7:]
        ---------------------------------------------------------------
        No-fold:
        If input 'auto', it will shuffle dataset and then cut 80% dataset to train and other to eval
        If input: [5]
        when len(dataset) == 10
        train-set : dataset[0:5]  eval-set : dataset[5:]
        """
        self.parser.add_argument('--fold_index', type=str, default='auto',
            help='where to fold, eg. when 5-fold and input: [2,4,6,7] -> sub-set: dataset[0:2],dataset[2:4],dataset[4:6],dataset[6:7],dataset[7:]')
        self.parser.add_argument('--mergelabel', type=str, default='None',
            help='merge some labels to one label and give the result, example:"[[0,1,4],[2,3,5]]" -> label(0,1,4) regard as 0,label(2,3,5) regard as 1')
        self.parser.add_argument('--mergelabel_name', type=str, default='None',help='name of labels,example:"a,b,c,d,e,f"')
        
        # ------------------------Network------------------------
        """Available Network
        1d: lstm, cnn_1d, resnet18_1d, resnet34_1d, multi_scale_resnet_1d,
            micro_multi_scale_resnet_1d,autoencoder,mlp
        2d: mobilenet, dfcnn, multi_scale_resnet, resnet18, resnet50, resnet101,
            densenet121, densenet201, squeezenet
        """
        self.parser.add_argument('--model_name', type=str, default='micro_multi_scale_resnet_1d',help='Choose model  lstm...')
        self.parser.add_argument('--model_type', type=str, default='auto',help='1d | 2d')
        # For lstm 
        self.parser.add_argument('--lstm_inputsize', type=str, default='auto',help='lstm_inputsize of LSTM')
        self.parser.add_argument('--lstm_timestep', type=int, default=100,help='time_step of LSTM')
        # For autoecoder
        self.parser.add_argument('--feature', type=int, default=3, help='number of encoder features')
        # For 2d network(stft spectrum)
        # Please cheek ./save_dir/spectrum_eg.jpg to change the following parameters
        self.parser.add_argument('--stft_size', type=int, default=512, help='length of each fft segment')
        self.parser.add_argument('--stft_stride', type=int, default=128, help='stride of each fft segment')
        self.parser.add_argument('--stft_n_downsample', type=int, default=1, help='downsample befor stft')
        self.parser.add_argument('--stft_no_log', action='store_true', help='if specified, do not log1p spectrum')
        self.parser.add_argument('--stft_shape', type=str, default='auto', help='shape of stft. It depend on \
            stft_size,stft_stride,stft_n_downsample. Do not input this parameter.')

        # ------------------------Training Matters------------------------
        self.parser.add_argument('--pretrained', type=str, default='',help='pretrained model path. If not specified, fo not use pretrained model')
        self.parser.add_argument('--continue_train', action='store_true', help='if specified, continue train')
        self.parser.add_argument('--lr', type=float, default=0.001,help='learning rate') 
        self.parser.add_argument('--batchsize', type=int, default=64,help='batchsize')
        self.parser.add_argument('--weight_mod', type=str, default='auto',help='Choose weight mode: auto | normal')
        self.parser.add_argument('--epochs', type=int, default=20,help='end epoch')
        self.parser.add_argument('--network_save_freq', type=int, default=5,help='the freq to save network')

        self.initialized = True

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

        if self.opt.gpu_id != -1:
            os.environ["CUDA_VISIBLE_DEVICES"] = str(self.opt.gpu_id)

        if self.opt.label != 'auto':
            self.opt.label = int(self.opt.label)
        if self.opt.input_nc !='auto':
            self.opt.input_nc = int(self.opt.input_nc)
        if self.opt.loadsize !='auto':
            self.opt.loadsize = int(self.opt.loadsize)
        if self.opt.finesize !='auto':
            self.opt.finesize = int(self.opt.finesize)
        if self.opt.lstm_inputsize != 'auto':
            self.opt.lstm_inputsize = int(self.opt.lstm_inputsize)

        if self.opt.model_type == 'auto':
            if self.opt.model_name in ['lstm', 'cnn_1d', 'resnet18_1d', 'resnet34_1d', 
                'multi_scale_resnet_1d','micro_multi_scale_resnet_1d','autoencoder','mlp']:
                self.opt.model_type = '1d'
            elif self.opt.model_name in ['dfcnn', 'multi_scale_resnet', 'resnet18', 'resnet50',
                'resnet101','densenet121', 'densenet201', 'squeezenet', 'mobilenet']:
                self.opt.model_type = '2d'
            else:
                print('\033[1;31m'+'Error: do not support this network '+self.opt.model_name+'\033[0m')
                exit(0)

        if self.opt.k_fold == 0 :
            self.opt.k_fold = 1

        if self.opt.fold_index != 'auto':
            self.opt.fold_index = eval(self.opt.fold_index)

        if os.path.isfile(os.path.join(self.opt.dataset_dir,'index.npy')):
            self.opt.fold_index = (np.load(os.path.join(self.opt.dataset_dir,'index.npy'))).tolist()

        self.opt.mergelabel = eval(self.opt.mergelabel)
        if self.opt.mergelabel_name != 'None':
            self.opt.mergelabel_name = self.opt.mergelabel_name.replace(" ", "").split(",")

        """Print and save options
        It will print both current options and default values(if different).
        It will save options into a text file / [checkpoints_dir] / opt.txt
        """
        message = ''
        message += '----------------- Options ---------------\n'
        for k, v in sorted(vars(self.opt).items()):
            comment = ''
            default = self.parser.get_default(k)
            if v != default:
                comment = '\t[default: %s]' % str(default)
            message += '{:>20}: {:<30}{}\n'.format(str(k), str(v), comment)
        message += '----------------- End -------------------'
        localtime = time.asctime(time.localtime(time.time()))
        util.makedirs(self.opt.save_dir)
        util.writelog(str(localtime)+'\n'+message, self.opt,True)

        return self.opt

def get_auto_options(opt,label_cnt_per,label_num,signals):
    
    shape = signals.shape
    if opt.label =='auto':
        opt.label = label_num
    if opt.input_nc =='auto':
        opt.input_nc = shape[1]
    if opt.loadsize =='auto':
        opt.loadsize = shape[2]
    if opt.finesize =='auto':
        opt.finesize = int(shape[2]*0.9)
    if opt.lstm_inputsize =='auto':
        opt.lstm_inputsize = opt.finesize//opt.lstm_timestep

    # weight
    opt.weight = np.ones(opt.label)
    if opt.weight_mod == 'auto':
        opt.weight = 1/label_cnt_per
        opt.weight = opt.weight/np.min(opt.weight)
    util.writelog('Loss_weight:'+str(opt.weight),opt,True)
    import torch
    opt.weight = torch.from_numpy(opt.weight).float()
    if opt.gpu_id != -1:      
        opt.weight = opt.weight.cuda()

    # label name
    if opt.label_name == 'auto':
        names = []
        for i in range(opt.label):
            names.append(str(i))
        opt.label_name = names
    elif not isinstance(opt.label_name,list):
        opt.label_name = opt.label_name.replace(" ", "").split(",")


    # check stft spectrum
    if opt.model_type =='2d':
        spectrums = []
        data = signals[np.random.randint(0,shape[0]-1)]
        for i in range(shape[1]):
            spectrums.append(dsp.signal2spectrum(data[i],opt.stft_size, opt.stft_stride, opt.stft_n_downsample, not opt.stft_no_log))
        plot.draw_spectrums(spectrums,opt)
        opt.stft_shape = spectrums[0].shape
        h,w = opt.stft_shape
        print('Shape of stft spectrum h,w:',opt.stft_shape)
        print('\033[1;37m'+'Please cheek ./save_dir/spectrum_eg.jpg to change parameters'+'\033[0m')
        
        if h<64 or w<64:
            print('\033[1;33m'+'Warning: spectrum is too small'+'\033[0m') 
        if h>512 or w>512:
            print('\033[1;33m'+'Warning: spectrum is too large'+'\033[0m')

    return opt