def set_tracking_opt(): Tracker_opt = opts().init() return Tracker_opt
def main(opt): torch.manual_seed(opt.seed) torch.backends.cudnn.benchmark = not opt.not_cuda_benchmark and not opt.test print('Setting up data...') Dataset = get_dataset(opt.dataset, opt.task) f = open(opt.data_cfg) data_config = json.load(f) trainset_paths = data_config['train'] dataset_root = data_config['root'] f.close() transforms = T.Compose([T.ToTensor()]) dataset = Dataset(opt, dataset_root, trainset_paths, (1088, 608), augment=True, transforms=transforms) opt = opts().update_dataset_info_and_set_heads(opt, dataset) print(opt) logger = Logger(opt) os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str opt.device = torch.device('cuda' if opt.gpus[0] >= 0 else 'cpu') print('Creating model...') model = create_model(opt.arch, opt.heads, opt.head_conv) optimizer = torch.optim.Adam(model.parameters(), opt.lr) start_epoch = 0 if opt.load_model != '': model, optimizer, start_epoch = load_model( model, opt.load_model, optimizer, opt.resume, opt.lr, opt.lr_step) # Get dataloader train_loader = torch.utils.data.DataLoader( dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers, pin_memory=True, drop_last=True ) print('Starting training...') Trainer = train_factory[opt.task] trainer = Trainer(opt, model, optimizer) trainer.set_device(opt.gpus, opt.chunk_sizes, opt.device) best = 1e10 for epoch in range(start_epoch + 1, opt.num_epochs + 1): mark = epoch if opt.save_all else 'last' log_dict_train, _ = trainer.train(epoch, train_loader) logger.write('epoch: {} |'.format(epoch)) for k, v in log_dict_train.items(): logger.scalar_summary('train_{}'.format(k), v, epoch) logger.write('{} {:8f} | '.format(k, v)) if opt.val_intervals > 0 and epoch % opt.val_intervals == 0: save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(mark)), epoch, model, optimizer) else: save_model(os.path.join(opt.save_dir, 'model_last.pth'), epoch, model, optimizer) logger.write('\n') if epoch in opt.lr_step: save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(epoch)), epoch, model, optimizer) lr = opt.lr * (0.1 ** (opt.lr_step.index(epoch) + 1)) print('Drop LR to', lr) for param_group in optimizer.param_groups: param_group['lr'] = lr if epoch % 5 == 0: save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(epoch)), epoch, model, optimizer) logger.close()
def demo(opt): result_root = opt.output_root if opt.output_root != '' else '.' mkdir_if_missing(result_root) logger.info('Starting tracking...') dataloader = datasets.LoadVideo(opt.input_video, opt.img_size) dataloader.cap.set(cv2.CAP_PROP_POS_MSEC, round(1000 * 120)) result_filename = os.path.join(result_root, 'results.txt') frame_rate = dataloader.frame_rate frame_dir = None if opt.output_format == 'text' else osp.join( result_root, 'frame') eval_seq(opt, dataloader, 'mot', result_filename, save_dir=frame_dir, show_image=False, frame_rate=frame_rate) if opt.output_format == 'video': output_video_path = osp.join(result_root, 'result.mp4') cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -b 5000k -c:v mpeg4 {}'.format( osp.join(result_root, 'frame'), output_video_path) os.system(cmd_str) if __name__ == '__main__': opt = opts().init() demo(opt)
log_dict_train, _ = trainer.train(epoch, train_loader) logger.write('epoch: {} |'.format(epoch)) for k, v in log_dict_train.items(): logger.scalar_summary('train_{}'.format(k), v, epoch) logger.write('{} {:8f} | '.format(k, v)) if opt.val_intervals > 0 and epoch % opt.val_intervals == 0: save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(mark)), epoch, model, optimizer) else: save_model(os.path.join(opt.save_dir, 'model_last.pth'), epoch, model, optimizer) logger.write('\n') if epoch in opt.lr_step: save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(epoch)), epoch, model, optimizer) lr = opt.lr * (0.1 ** (opt.lr_step.index(epoch) + 1)) print('Drop LR to', lr) for param_group in optimizer.param_groups: param_group['lr'] = lr if epoch % 5 == 0: save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(epoch)), epoch, model, optimizer) logger.close() if __name__ == '__main__': os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1' opt = opts().parse() main(opt)
sub_imgs_detection = [] for img_name_detection in sub_imgs_names_detection: sub_img_detection = cv2.imread( os.path.join(intermediate_resutls_path_detection, img_name_detection)) sub_imgs_detection.append(sub_img_detection) return action_index, sub_imgs_tracking, ReID_features_tracking, sub_imgs_detection, ReID_features_detection if __name__ == "__main__": from opt import opt from FairMot.lib.opts import opts from CalibrateTransfer.img_operation import ScreenSHot detector_opt = opts().init() queueSize = 1000 Tracker_output_queue = Queue(1000) dir_name = opt.dir_name root_path = os.path.join(opt.data_root, '{}'.format(dir_name)) file_name = opt.file_name Videoparameters, \ setting_parameter, \ action_datas, \ channel_list, \ parameter = read_data_from_json_file(root_path, file_name, opt) vis_path = os.path.join(root_path, 'vis') os.makedirs(vis_path, exist_ok=True) multi = 10