def test_icdar2015(): dataset = build_dataset('ICDAR2015', '../../data/icdar2015') print(len(dataset)) average_time = [] for i in range(10): tic = time.time() data = dataset[i] toc = time.time() canvas = draw_polygons(data['image'], data['polygons']) average_time.append(toc - tic) plt.figure("img") plt.figure(figsize=(12, 12)) plt.imshow(np.array(canvas, dtype=np.uint8)) plt.show() print(np.mean(np.array(average_time)))
def build_train_data_loader(cfg): names = cfg.DATA.NAMES dirs = cfg.DATA.DIRS dataset = [] for name, dataset_dir in zip(names, dirs): dataset.append(build_dataset(name, dataset_dir)) # 读取数据集 dataset = ConcatDataset(dataset) # 数据增强 dataset = Augment(dataset, cfg) # encoder dataset = TargetEncoder(dataset, cfg) sampler = TrainingSampler(len(dataset)) data_loader = torch.utils.data.DataLoader(dataset, cfg.SOLVER.BATCH_SIZE, num_workers=cfg.SOLVER.NUM_WORKS, pin_memory=cfg.SOLVER.PIN_MEMORY, sampler=sampler) return data_loader