def __init__(self, pretrain=False):
super(ResDown, self).__init__()
self.features = resnet50(layer3=True, layer4=False)
self.features_127 = self.features_255 = self.features
if pretrain:
load_pretrain(self.features, 'resnet.model')
self.downsample = ResDownS(1024, 256)
self.layers = [self.downsample, self.features.layer2, self.features.layer3]
self.train_nums = [1, 3]
self.change_point = [0, 0.5]
self.unfix(0.0)
def main():
global args, best_acc, tb_writer, logger
args = parser.parse_args()
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global')
logger.info("\n" + collect_env_info())
logger.info(args)
cfg = load_config(args)
logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
if args.log_dir:
tb_writer = SummaryWriter(args.log_dir)
else:
tb_writer = Dummy()
# build dataset
train_loader, val_loader = build_data_loader(cfg)
if args.arch == 'Custom':
from custom import Custom
model = Custom(pretrain=True, anchors=cfg['anchors'])
else:
exit()
logger.info(model)
if args.pretrained:
model = load_pretrain(model, args.pretrained)
model = model.cuda()
dist_model = torch.nn.DataParallel(model,
list(range(
torch.cuda.device_count()))).cuda()
if args.resume and args.start_epoch != 0:
model.features.unfix((args.start_epoch - 1) / args.epochs)
optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch)
# optionally resume from a checkpoint
if args.resume:
assert os.path.isfile(args.resume), '{} is not a valid file'.format(
args.resume)
model, optimizer, args.start_epoch, best_acc, arch = restore_from(
model, optimizer, args.resume)
dist_model = torch.nn.DataParallel(
model, list(range(torch.cuda.device_count()))).cuda()
logger.info(lr_scheduler)
logger.info('model prepare done')
train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch,
cfg)
def main():
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
params = {
'penalty_k': args.penalty_k,
'window_influence': args.window_influence,
'lr': args.lr,
'instance_size': args.search_region
}
num_search = len(params['penalty_k']) * len(params['window_influence']) * \
len(params['lr']) * len(params['instance_size'])
print(params)
print(num_search)
cfg = load_config(args)
if args.arch == 'Custom':
from custom import Custom
model = Custom(anchors=cfg['anchors'])
else:
model = models.__dict__[args.arch](anchors=cfg['anchors'])
if args.resume:
assert isfile(args.resume), '{} is not a valid file'.format(
args.resume)
model = load_pretrain(model, args.resume)
model.eval()
model = model.cuda()
default_hp = cfg.get('hp', {})
p = dict()
p['network'] = model
p['network_name'] = args.arch + '_mask_' + args.resume.split(
'/')[-1].split('.')[0]
p['dataset'] = args.dataset
global ims, gt, annos, image_files, anno_files
dataset_info = load_dataset(args.dataset)
videos = list(dataset_info.keys())
np.random.shuffle(videos)
for video in videos:
print(video)
p['video'] = video
ims = None
annos = None
image_files = dataset_info[video]['image_files']
anno_files = dataset_info[video]['anno_files']
gt = dataset_info[video]['gt']
np.random.shuffle(params['penalty_k'])
np.random.shuffle(params['window_influence'])
np.random.shuffle(params['lr'])
for penalty_k in params['penalty_k']:
for window_influence in params['window_influence']:
for lr in params['lr']:
for instance_size in params['instance_size']:
p['hp'] = default_hp.copy()
p['hp'].update({
'penalty_k': penalty_k,
'window_influence': window_influence,
'lr': lr,
'instance_size': instance_size,
})
tune(p)
def main():
global args, logger, v_id
args = parser.parse_args()
cfg = load_config(args)
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global')
logger.info(args)
# setup model
if args.arch == 'Custom':
from custom import Custom
model = Custom(anchors=cfg['anchors'])
else:
parser.error('invalid architecture: {}'.format(args.arch))
if args.resume:
assert isfile(args.resume), '{} is not a valid file'.format(
args.resume)
model = load_pretrain(model, args.resume)
model.eval()
device = torch.device('cuda' if (
torch.cuda.is_available() and not args.cpu) else 'cpu')
model = model.to(device)
# setup dataset
dataset = load_dataset(args.dataset)
# VOS or VOT?
if args.dataset in ['DAVIS2016', 'DAVIS2017', 'ytb_vos'] and args.mask:
vos_enable = True # enable Mask output
else:
vos_enable = False
total_lost = 0 # VOT
iou_lists = [] # VOS
speed_list = []
for v_id, video in enumerate(dataset.keys(), start=1):
if args.video != '' and video != args.video:
continue
if vos_enable:
iou_list, speed = track_vos(
model,
dataset[video],
cfg['hp'] if 'hp' in cfg.keys() else None,
args.mask,
args.refine,
args.dataset in ['DAVIS2017', 'ytb_vos'],
device=device)
iou_lists.append(iou_list)
else:
lost, speed = track_vot(model,
dataset[video],
cfg['hp'] if 'hp' in cfg.keys() else None,
args.mask,
args.refine,
device=device)
total_lost += lost
speed_list.append(speed)
# report final result
if vos_enable:
for thr, iou in zip(thrs, np.mean(np.concatenate(iou_lists), axis=0)):
logger.info('Segmentation Threshold {:.2f} mIoU: {:.3f}'.format(
thr, iou))
else:
logger.info('Total Lost: {:d}'.format(total_lost))
logger.info('Mean Speed: {:.2f} FPS'.format(np.mean(speed_list)))