Spaces:
Runtime error
Runtime error
| from __future__ import absolute_import | |
| from __future__ import division | |
| from __future__ import print_function | |
| import argparse | |
| import os | |
| import sys | |
| import json | |
| class opts(object): | |
| def __init__(self): | |
| self.parser = argparse.ArgumentParser() | |
| # basic experiment setting | |
| self.parser.add_argument('task', default='', | |
| help='ctdet | ddd | multi_pose ' | |
| '| tracking or combined with ,') | |
| self.parser.add_argument('--dataset', default='coco', | |
| help='see lib/dataset/dataset_facotry for ' + | |
| 'available datasets') | |
| self.parser.add_argument('--test_dataset', default='', | |
| help='coco | kitti | coco_hp | pascal') | |
| self.parser.add_argument('--exp_id', default='default') | |
| self.parser.add_argument('--test', action='store_true') | |
| self.parser.add_argument('--debug', type=int, default=0, | |
| help='level of visualization.' | |
| '1: only show the final detection results' | |
| '2: show the network output features' | |
| '3: use matplot to display' # useful when lunching training with ipython notebook | |
| '4: save all visualizations to disk') | |
| self.parser.add_argument('--no_pause', action='store_true') | |
| self.parser.add_argument('--demo', default='', | |
| help='path to image/ image folders/ video. ' | |
| 'or "webcam"') | |
| self.parser.add_argument('--load_model', default='', | |
| help='path to pretrained model') | |
| self.parser.add_argument('--resume', action='store_true', | |
| help='resume an experiment. ' | |
| 'Reloaded the optimizer parameter and ' | |
| 'set load_model to model_last.pth ' | |
| 'in the exp dir if load_model is empty.') | |
| # system | |
| self.parser.add_argument('--gpus', default='0', | |
| help='-1 for CPU, use comma for multiple gpus') | |
| self.parser.add_argument('--num_workers', type=int, default=4, | |
| help='dataloader threads. 0 for single-thread.') | |
| self.parser.add_argument('--not_cuda_benchmark', action='store_true', | |
| help='disable when the input size is not fixed.') | |
| self.parser.add_argument('--seed', type=int, default=317, | |
| help='random seed') # from CornerNet | |
| self.parser.add_argument('--not_set_cuda_env', action='store_true', | |
| help='used when training in slurm clusters.') | |
| # log | |
| self.parser.add_argument('--print_iter', type=int, default=0, | |
| help='disable progress bar and print to screen.') | |
| self.parser.add_argument('--save_all', action='store_true', | |
| help='save model to disk every 5 epochs.') | |
| self.parser.add_argument('--vis_thresh', type=float, default=0.3, | |
| help='visualization threshold.') | |
| self.parser.add_argument('--debugger_theme', default='white', | |
| choices=['white', 'black']) | |
| self.parser.add_argument('--eval_val', action='store_true') | |
| self.parser.add_argument('--save_imgs', default='', help='') | |
| self.parser.add_argument('--save_img_suffix', default='', help='') | |
| self.parser.add_argument('--skip_first', type=int, default=-1, help='') | |
| self.parser.add_argument('--save_video', action='store_true') | |
| self.parser.add_argument('--save_framerate', type=int, default=30) | |
| self.parser.add_argument('--resize_video', action='store_true') | |
| self.parser.add_argument('--video_h', type=int, default=512, help='') | |
| self.parser.add_argument('--video_w', type=int, default=512, help='') | |
| self.parser.add_argument('--transpose_video', action='store_true') | |
| self.parser.add_argument('--show_track_color', action='store_true') | |
| self.parser.add_argument('--not_show_bbox', action='store_true') | |
| self.parser.add_argument('--not_show_number', action='store_true') | |
| self.parser.add_argument('--qualitative', action='store_true') | |
| self.parser.add_argument('--tango_color', action='store_true') | |
| # model | |
| self.parser.add_argument('--arch', default='dla_34', | |
| help='model architecture. Currently tested' | |
| 'res_18 | res_101 | resdcn_18 | resdcn_101 |' | |
| 'dlav0_34 | dla_34 | hourglass') | |
| self.parser.add_argument('--dla_node', default='dcn') | |
| self.parser.add_argument('--head_conv', type=int, default=-1, | |
| help='conv layer channels for output head' | |
| '0 for no conv layer' | |
| '-1 for default setting: ' | |
| '64 for resnets and 256 for dla.') | |
| self.parser.add_argument('--num_head_conv', type=int, default=1) | |
| self.parser.add_argument('--head_kernel', type=int, default=3, help='') | |
| self.parser.add_argument('--down_ratio', type=int, default=4, | |
| help='output stride. Currently only supports 4.') | |
| self.parser.add_argument('--not_idaup', action='store_true') | |
| self.parser.add_argument('--num_classes', type=int, default=-1) | |
| self.parser.add_argument('--num_layers', type=int, default=101) | |
| self.parser.add_argument('--backbone', default='dla34') | |
| self.parser.add_argument('--neck', default='dlaup') | |
| self.parser.add_argument('--msra_outchannel', type=int, default=256) | |
| self.parser.add_argument('--efficient_level', type=int, default=0) | |
| self.parser.add_argument('--prior_bias', type=float, default=-4.6) # -2.19 | |
| self.parser.add_argument('--embedding', action='store_true') | |
| self.parser.add_argument('--box_nms', type=float, default=-1) | |
| self.parser.add_argument('--inference', action='store_true') | |
| self.parser.add_argument('--clip_len', type=int, default=1, help='number of images used in trades' | |
| 'including the current image') | |
| self.parser.add_argument('--no_repeat', action='store_true', default=True) | |
| self.parser.add_argument('--seg', action='store_true', default=False) | |
| self.parser.add_argument('--seg_feat_channel', default=8, type=int, help='.') | |
| self.parser.add_argument('--deform_kernel_size', type=int, default=3) | |
| self.parser.add_argument('--trades', action='store_true', help='Track to Detect and Segment:' | |
| 'An Online Multi Object Tracker') | |
| # input | |
| self.parser.add_argument('--input_res', type=int, default=-1, | |
| help='input height and width. -1 for default from ' | |
| 'dataset. Will be overriden by input_h | input_w') | |
| self.parser.add_argument('--input_h', type=int, default=-1, | |
| help='input height. -1 for default from dataset.') | |
| self.parser.add_argument('--input_w', type=int, default=-1, | |
| help='input width. -1 for default from dataset.') | |
| self.parser.add_argument('--dataset_version', default='') | |
| # train | |
| self.parser.add_argument('--optim', default='adam') | |
| self.parser.add_argument('--lr', type=float, default=1.25e-4, | |
| help='learning rate for batch size 32.') | |
| self.parser.add_argument('--lr_step', type=str, default='60', | |
| help='drop learning rate by 10.') | |
| self.parser.add_argument('--save_point', type=str, default='90', | |
| help='when to save the model to disk.') | |
| self.parser.add_argument('--num_epochs', type=int, default=70, | |
| help='total training epochs.') | |
| self.parser.add_argument('--batch_size', type=int, default=32, | |
| help='batch size') | |
| self.parser.add_argument('--master_batch_size', type=int, default=-1, | |
| help='batch size on the master gpu.') | |
| self.parser.add_argument('--num_iters', type=int, default=-1, | |
| help='default: #samples / batch_size.') | |
| self.parser.add_argument('--val_intervals', type=int, default=10000, | |
| help='number of epochs to run validation.') | |
| self.parser.add_argument('--trainval', action='store_true', | |
| help='include validation in training and ' | |
| 'test on test set') | |
| self.parser.add_argument('--ltrb', action='store_true', | |
| help='') | |
| self.parser.add_argument('--ltrb_weight', type=float, default=0.1, | |
| help='') | |
| self.parser.add_argument('--reset_hm', action='store_true') | |
| self.parser.add_argument('--reuse_hm', action='store_true') | |
| self.parser.add_argument('--use_kpt_center', action='store_true') | |
| self.parser.add_argument('--add_05', action='store_true') | |
| self.parser.add_argument('--dense_reg', type=int, default=1, help='') | |
| # test | |
| self.parser.add_argument('--flip_test', action='store_true', | |
| help='flip data augmentation.') | |
| self.parser.add_argument('--test_scales', type=str, default='1', | |
| help='multi scale test augmentation.') | |
| self.parser.add_argument('--nms', action='store_true', | |
| help='run nms in testing.') | |
| self.parser.add_argument('--K', type=int, default=100, | |
| help='max number of output objects.') | |
| self.parser.add_argument('--not_prefetch_test', action='store_true', | |
| help='not use parallal data pre-processing.') | |
| self.parser.add_argument('--fix_short', type=int, default=-1) | |
| self.parser.add_argument('--keep_res', action='store_true', | |
| help='keep the original resolution' | |
| ' during validation.') | |
| self.parser.add_argument('--map_argoverse_id', action='store_true', | |
| help='if trained on nuscenes and eval on kitti') | |
| self.parser.add_argument('--out_thresh', type=float, default=-1, | |
| help='') | |
| self.parser.add_argument('--depth_scale', type=float, default=1, | |
| help='') | |
| self.parser.add_argument('--save_results', action='store_true') | |
| self.parser.add_argument('--load_results', default='') | |
| self.parser.add_argument('--use_loaded_results', action='store_true') | |
| self.parser.add_argument('--ignore_loaded_cats', default='') | |
| self.parser.add_argument('--model_output_list', action='store_true', | |
| help='Used when convert to onnx') | |
| self.parser.add_argument('--non_block_test', action='store_true') | |
| self.parser.add_argument('--vis_gt_bev', default='', help='') | |
| self.parser.add_argument('--kitti_split', default='3dop', | |
| help='different validation split for kitti: ' | |
| '3dop | subcnn') | |
| self.parser.add_argument('--test_focal_length', type=int, default=-1) | |
| # dataset | |
| self.parser.add_argument('--not_rand_crop', action='store_true', | |
| help='not use the random crop data augmentation' | |
| 'from CornerNet.') | |
| self.parser.add_argument('--not_max_crop', action='store_true', | |
| help='used when the training dataset has' | |
| 'inbalanced aspect ratios.') | |
| self.parser.add_argument('--shift', type=float, default=0, | |
| help='when not using random crop, 0.1' | |
| 'apply shift augmentation.') | |
| self.parser.add_argument('--scale', type=float, default=0, | |
| help='when not using random crop, 0.4' | |
| 'apply scale augmentation.') | |
| self.parser.add_argument('--aug_rot', type=float, default=0, | |
| help='probability of applying ' | |
| 'rotation augmentation.') | |
| self.parser.add_argument('--rotate', type=float, default=0, | |
| help='when not using random crop' | |
| 'apply rotation augmentation.') | |
| self.parser.add_argument('--flip', type=float, default=0.5, | |
| help='probability of applying flip augmentation.') | |
| self.parser.add_argument('--no_color_aug', action='store_true', | |
| help='not use the color augmenation ' | |
| 'from CornerNet') | |
| # Tracking | |
| self.parser.add_argument('--tracking', action='store_true') | |
| self.parser.add_argument('--pre_hm', action='store_true') | |
| self.parser.add_argument('--same_aug_pre', action='store_true') | |
| self.parser.add_argument('--zero_pre_hm', action='store_true') | |
| self.parser.add_argument('--hm_disturb', type=float, default=0) | |
| self.parser.add_argument('--lost_disturb', type=float, default=0) | |
| self.parser.add_argument('--fp_disturb', type=float, default=0) | |
| self.parser.add_argument('--pre_thresh', type=float, default=-1) | |
| self.parser.add_argument('--track_thresh', type=float, default=0.3) | |
| self.parser.add_argument('--match_thresh', type=float, default=0.8) | |
| self.parser.add_argument('--track_buffer', type=int, default=30) | |
| self.parser.add_argument('--new_thresh', type=float, default=0.0) | |
| self.parser.add_argument('--max_frame_dist', type=int, default=3) | |
| self.parser.add_argument('--ltrb_amodal', action='store_true') | |
| self.parser.add_argument('--ltrb_amodal_weight', type=float, default=0.1) | |
| self.parser.add_argument('--window_size', type=int, default=20) | |
| self.parser.add_argument('--public_det', action='store_true') | |
| self.parser.add_argument('--no_pre_img', action='store_true') | |
| self.parser.add_argument('--zero_tracking', action='store_true') | |
| self.parser.add_argument('--hungarian', action='store_true') | |
| self.parser.add_argument('--max_age', type=int, default=-1) | |
| # loss | |
| self.parser.add_argument('--tracking_weight', type=float, default=1) | |
| self.parser.add_argument('--reg_loss', default='l1', | |
| help='regression loss: sl1 | l1 | l2') | |
| self.parser.add_argument('--hm_weight', type=float, default=1, | |
| help='loss weight for keypoint heatmaps.') | |
| self.parser.add_argument('--off_weight', type=float, default=1, | |
| help='loss weight for keypoint local offsets.') | |
| self.parser.add_argument('--wh_weight', type=float, default=0.1, | |
| help='loss weight for bounding box size.') | |
| self.parser.add_argument('--hp_weight', type=float, default=1, | |
| help='loss weight for human pose offset.') | |
| self.parser.add_argument('--hm_hp_weight', type=float, default=1, | |
| help='loss weight for human keypoint heatmap.') | |
| self.parser.add_argument('--amodel_offset_weight', type=float, default=1, | |
| help='Please forgive the typo.') | |
| self.parser.add_argument('--dep_weight', type=float, default=1, | |
| help='loss weight for depth.') | |
| self.parser.add_argument('--dim_weight', type=float, default=1, | |
| help='loss weight for 3d bounding box size.') | |
| self.parser.add_argument('--rot_weight', type=float, default=1, | |
| help='loss weight for orientation.') | |
| self.parser.add_argument('--nuscenes_att', action='store_true') | |
| self.parser.add_argument('--nuscenes_att_weight', type=float, default=1) | |
| self.parser.add_argument('--velocity', action='store_true') | |
| self.parser.add_argument('--velocity_weight', type=float, default=1) | |
| self.parser.add_argument('--nID', type=int, default=-1) | |
| # custom dataset | |
| self.parser.add_argument('--custom_dataset_img_path', default='') | |
| self.parser.add_argument('--custom_dataset_ann_path', default='') | |
| def parse(self, args=''): | |
| if args == '': | |
| opt = self.parser.parse_args() | |
| else: | |
| opt = self.parser.parse_args(args) | |
| if opt.test_dataset == '': | |
| opt.test_dataset = opt.dataset | |
| opt.gpus_str = opt.gpus | |
| opt.gpus = [int(gpu) for gpu in opt.gpus.split(',')] | |
| opt.gpus = [i for i in range(len(opt.gpus))] if opt.gpus[0] >=0 else [-1] | |
| opt.lr_step = [int(i) for i in opt.lr_step.split(',')] | |
| opt.save_point = [int(i) for i in opt.save_point.split(',')] | |
| opt.test_scales = [float(i) for i in opt.test_scales.split(',')] | |
| opt.save_imgs = [i for i in opt.save_imgs.split(',')] \ | |
| if opt.save_imgs != '' else [] | |
| opt.ignore_loaded_cats = \ | |
| [int(i) for i in opt.ignore_loaded_cats.split(',')] \ | |
| if opt.ignore_loaded_cats != '' else [] | |
| opt.num_workers = max(opt.num_workers, 2 * len(opt.gpus)) | |
| opt.pre_img = False | |
| if 'tracking' in opt.task: | |
| print('Running tracking') | |
| opt.tracking = True | |
| # opt.out_thresh = max(opt.track_thresh, opt.out_thresh) | |
| # opt.pre_thresh = max(opt.track_thresh, opt.pre_thresh) | |
| # opt.new_thresh = max(opt.track_thresh, opt.new_thresh) | |
| opt.pre_img = not opt.no_pre_img | |
| print('Using tracking threshold for out threshold!', opt.track_thresh) | |
| # if 'ddd' in opt.task: | |
| opt.show_track_color = True | |
| if opt.dataset in ['mot', 'mots', 'youtube_vis']: | |
| opt.overlap_thresh = 0.05 | |
| elif opt.dataset == 'nuscenes': | |
| opt.window_size = 7 | |
| opt.overlap_thresh = -1 | |
| else: | |
| opt.overlap_thresh = 0.05 | |
| opt.fix_res = not opt.keep_res | |
| print('Fix size testing.' if opt.fix_res else 'Keep resolution testing.') | |
| if opt.head_conv == -1: # init default head_conv | |
| opt.head_conv = 256 if 'dla' in opt.arch else 64 | |
| opt.pad = 127 if 'hourglass' in opt.arch else 31 | |
| opt.num_stacks = 2 if opt.arch == 'hourglass' else 1 | |
| if opt.master_batch_size == -1: | |
| opt.master_batch_size = opt.batch_size // len(opt.gpus) | |
| rest_batch_size = (opt.batch_size - opt.master_batch_size) | |
| opt.chunk_sizes = [opt.master_batch_size] | |
| for i in range(len(opt.gpus) - 1): | |
| slave_chunk_size = rest_batch_size // (len(opt.gpus) - 1) | |
| if i < rest_batch_size % (len(opt.gpus) - 1): | |
| slave_chunk_size += 1 | |
| opt.chunk_sizes.append(slave_chunk_size) | |
| print('training chunk_sizes:', opt.chunk_sizes) | |
| if opt.debug > 0: | |
| opt.num_workers = 0 | |
| opt.batch_size = 1 | |
| opt.gpus = [opt.gpus[0]] | |
| opt.master_batch_size = -1 | |
| # log dirs | |
| opt.root_dir = os.path.join(os.path.dirname(__file__), '..', '..') | |
| opt.data_dir = os.path.join(opt.root_dir, 'data') | |
| opt.exp_dir = os.path.join(opt.root_dir, 'exp', opt.task) | |
| opt.save_dir = os.path.join(opt.exp_dir, opt.exp_id) | |
| opt.debug_dir = os.path.join(opt.save_dir, 'debug') | |
| if opt.resume and opt.load_model == '': | |
| opt.load_model = os.path.join(opt.save_dir, 'model_last.pth') | |
| return opt | |
| def update_dataset_info_and_set_heads(self, opt, dataset): | |
| opt.num_classes = dataset.num_categories \ | |
| if opt.num_classes < 0 else opt.num_classes | |
| # input_h(w): opt.input_h overrides opt.input_res overrides dataset default | |
| input_h, input_w = dataset.default_resolution | |
| input_h = opt.input_res if opt.input_res > 0 else input_h | |
| input_w = opt.input_res if opt.input_res > 0 else input_w | |
| opt.input_h = opt.input_h if opt.input_h > 0 else input_h | |
| opt.input_w = opt.input_w if opt.input_w > 0 else input_w | |
| opt.output_h = opt.input_h // opt.down_ratio | |
| opt.output_w = opt.input_w // opt.down_ratio | |
| opt.input_res = max(opt.input_h, opt.input_w) | |
| opt.output_res = max(opt.output_h, opt.output_w) | |
| opt.heads = {'hm': opt.num_classes, 'reg': 2, 'wh': 2} | |
| if not opt.trades: | |
| if 'tracking' in opt.task: | |
| opt.heads.update({'tracking': 2}) | |
| if 'ddd' in opt.task: | |
| opt.heads.update({'dep': 1, 'rot': 8, 'dim': 3, 'amodel_offset': 2}) | |
| if 'multi_pose' in opt.task: | |
| opt.heads.update({ | |
| 'hps': dataset.num_joints * 2, 'hm_hp': dataset.num_joints, | |
| 'hp_offset': 2}) | |
| if opt.ltrb: | |
| opt.heads.update({'ltrb': 4}) | |
| if opt.ltrb_amodal: | |
| opt.heads.update({'ltrb_amodal': 4}) | |
| if opt.nuscenes_att: | |
| opt.heads.update({'nuscenes_att': 8}) | |
| if opt.velocity: | |
| opt.heads.update({'velocity': 3}) | |
| if opt.embedding: | |
| opt.heads.update({'embedding': 128}) | |
| if opt.seg: | |
| opt.heads.update({'conv_weight': 2*opt.seg_feat_channel**2 + 5*opt.seg_feat_channel + 1}) | |
| opt.heads.update({'seg_feat': opt.seg_feat_channel}) | |
| weight_dict = {'hm': opt.hm_weight, 'wh': opt.wh_weight, | |
| 'reg': opt.off_weight, 'hps': opt.hp_weight, | |
| 'hm_hp': opt.hm_hp_weight, 'hp_offset': opt.off_weight, | |
| 'dep': opt.dep_weight, 'rot': opt.rot_weight, | |
| 'dim': opt.dim_weight, | |
| 'amodel_offset': opt.amodel_offset_weight, | |
| 'ltrb': opt.ltrb_weight, | |
| 'tracking': opt.tracking_weight, | |
| 'ltrb_amodal': opt.ltrb_amodal_weight, | |
| 'nuscenes_att': opt.nuscenes_att_weight, | |
| 'velocity': opt.velocity_weight, | |
| 'embedding': 1.0, | |
| 'conv_weight': 1.0, | |
| 'seg_feat':1.0} | |
| opt.weights = {head: weight_dict[head] for head in opt.heads} | |
| if opt.trades: | |
| opt.weights['cost_volume'] = 1.0 | |
| if opt.seg: | |
| opt.weights['mask_loss'] = 1.0 | |
| for head in opt.weights: | |
| if opt.weights[head] == 0: | |
| del opt.heads[head] | |
| opt.head_conv = {head: [opt.head_conv \ | |
| for i in range(opt.num_head_conv if head != 'reg' else 1)] for head in opt.heads} | |
| print('input h w:', opt.input_h, opt.input_w) | |
| print('heads', opt.heads) | |
| print('weights', opt.weights) | |
| print('head conv', opt.head_conv) | |
| return opt | |
| def init(self, args=''): | |
| # only used in demo | |
| default_dataset_info = { | |
| 'ctdet': 'coco', 'multi_pose': 'coco_hp', 'ddd': 'nuscenes', | |
| 'tracking,ctdet': 'coco', 'tracking,multi_pose': 'coco_hp', | |
| 'tracking,ddd': 'nuscenes' | |
| } | |
| opt = self.parse() | |
| from dataset.dataset_factory import dataset_factory | |
| train_dataset = default_dataset_info[opt.task] \ | |
| if opt.task in default_dataset_info else 'coco' | |
| if opt.dataset != 'coco': | |
| dataset = dataset_factory[opt.dataset] | |
| else: | |
| dataset = dataset_factory[train_dataset] | |
| opt = self.update_dataset_info_and_set_heads(opt, dataset) | |
| return opt | |