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
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])