def __init__(self, pretrain=False): super(ResDown, self).__init__() self.features = resnet50(layer3=True, layer4=False) 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, 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()) 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)))
def __init__(self): resume = '/home/saad/Root/vision/Computer_Vision/Tracking-systems/SiamMask_DAVIS.pth' config = '/home/saad/Root/vision/Computer_Vision/Tracking-systems/config_davis.json' self.cfg = load_config(config=config) self.siammask = Custom(anchors=self.cfg['anchors']) self.siammask = load_pretrain(self.siammask, resume) self.active_boxes = [] self.frames_generator = get_detection_output_as_frames_generator() self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.siammask.eval().to(self.device) self.tracker_initialized = False self.state = None self.times = np_.asarray([])
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': model = Custom(anchors=cfg['anchors']) elif args.arch == 'Custom_Sky': model = Custom_Sky(anchors=cfg['anchors']) else: exit() logger.info(model) if args.pretrained: model = load_pretrain(model, args.pretrained) # print(summary(model=model, input_size=(3, 511, 511), batch_size=1)) 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: print(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)