Beispiel #1
0
def create_dataloader(opt, path, imgsz, stride, hyp=None, pad=0.0, prefix='',
                      void_classes=[], valid_classes=[], cache=False, rect=False, rank=-1):
    cfg = get_cfg_defaults()
    cfg.merge_from_file('../dataloaders/configs/sunrgbd.yaml')
    cfg.merge_from_list(['DATASET.ANNOTATION_TYPE', 'bbox',
                         'DATASET.NO_TRANSFORMS', True,
                         'TRAIN.BATCH_SIZE', 1])

    train_loader, val_loader, test_loader, num_class = make_data_loader(cfg)
    dataloader = train_loader
    return dataloader, dataloader.dataset
Beispiel #2
0
    from torch.utils.data import DataLoader
    import matplotlib.pyplot as plt
    import argparse

    parser = argparse.ArgumentParser(description="Test sbd Loader")
    parser.add_argument('config_file', help='config file path')
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )

    args = parser.parse_args()

    cfg = get_cfg_defaults()
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze()
    print(cfg)

    sbd_train = SBDSegmentation(cfg, split='train')
    dataloader = DataLoader(sbd_train, batch_size=2, shuffle=True, num_workers=2)

    for ii, sample in enumerate(dataloader):
        for jj in range(sample["image"].size()[0]):
            img = sample['image'].numpy()
            gt = sample['label'].numpy()
            tmp = np.array(gt[jj]).astype(np.uint8)
            segmap = decode_segmap(tmp, dataset='pascal')
            img_tmp = np.transpose(img[jj], axes=[1, 2, 0])