Beispiel #1
0
def get_test_loader(dataset, height, width, batch_size, workers, testset=None):
    normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])

    test_transformer = T.Compose(
        [T.Resize((height, width), interpolation=3),
         T.ToTensor(), normalizer])

    if (testset is None):
        testset = list(set(dataset.query) | set(dataset.gallery))

    test_loader = DataLoader(Preprocessor(testset,
                                          root=dataset.images_dir,
                                          transform=test_transformer),
                             batch_size=batch_size,
                             num_workers=workers,
                             shuffle=False,
                             pin_memory=True)

    return test_loader
Beispiel #2
0
def get_data(name, data_dir, height, width, batch_size, workers):
    root = osp.join(data_dir, name)

    dataset = datasets.create(name, root)

    normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])

    test_transformer = T.Compose(
        [T.Resize((height, width), interpolation=3),
         T.ToTensor(), normalizer])

    test_loader = DataLoader(Preprocessor(
        list(set(dataset.query) | set(dataset.gallery)),
        root=dataset.images_dir,
        transform=test_transformer),
                             batch_size=batch_size,
                             num_workers=workers,
                             shuffle=False,
                             pin_memory=True)

    return dataset, test_loader
Beispiel #3
0
def get_train_loader(args,
                     dataset,
                     height,
                     width,
                     batch_size,
                     workers,
                     num_instances,
                     iters,
                     trainset=None):

    normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    train_transformer = T.Compose([
        T.Resize((height, width), interpolation=3),
        T.RandomHorizontalFlip(p=0.5),
        T.Pad(10),
        T.RandomCrop((height, width)),
        T.ToTensor(), normalizer,
        T.RandomErasing(probability=0.5, mean=[0.485, 0.456, 0.406])
    ])

    train_set = sorted(dataset.train) if trainset is None else sorted(trainset)
    rmgs_flag = num_instances > 0
    if rmgs_flag:
        sampler = RandomMultipleGallerySampler(train_set, num_instances)
    else:
        sampler = None
    train_loader = IterLoader(DataLoader(Preprocessor(
        train_set, root=dataset.images_dir, transform=train_transformer),
                                         batch_size=batch_size,
                                         num_workers=workers,
                                         sampler=sampler,
                                         shuffle=not rmgs_flag,
                                         pin_memory=True,
                                         drop_last=True),
                              length=iters)

    return train_loader