示例#1
0
def get_data(name,
             data_dir,
             height,
             width,
             batch_size,
             workers,
             num_instances,
             iters=200):
    root = osp.join(data_dir)

    dataset = datasets.create(name, root)

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

    train_set = dataset.train
    num_classes = dataset.num_train_pids

    train_transformer = T.Compose([
        T.Resize((height, width), interpolation=3),
        T.RandomHorizontalFlip(p=0.5),
        T.Pad(10),
        T.RandomCrop((height, width)),
        # T.AugMix(),
        T.ToTensor(),
        normalizer
    ])

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

    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)

    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, num_classes, train_loader, test_loader
def get_train_loader(dataset,
                     height,
                     width,
                     choice_c,
                     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.596, 0.558, 0.497])
    ])

    train_set = trainset  #dataset.train if trainset is None else trainset
    rmgs_flag = num_instances > 0
    if rmgs_flag:
        sampler = RandomMultipleGallerySampler(train_set, num_instances,
                                               choice_c)
    else:
        sampler = None

    train_loader = IterLoader(DataLoader(Preprocessor(
        train_set,
        root=dataset.images_dir,
        transform=train_transformer,
        mutual=True),
                                         batch_size=batch_size,
                                         num_workers=workers,
                                         sampler=sampler,
                                         shuffle=not rmgs_flag,
                                         pin_memory=True,
                                         drop_last=True),
                              length=iters)

    # train_loader = IterLoader(
    #     DataLoader(UnsupervisedCamStylePreprocessor(train_set, root=dataset.images_dir, transform=train_transformer,
    #                                                 num_cam=dataset.num_cam,camstyle_dir=dataset.camstyle_dir, mutual=True),
    #                batch_size=batch_size, num_workers=0, sampler=sampler,#workers
    #                shuffle=not rmgs_flag, pin_memory=True, drop_last=True), length=iters)

    return train_loader
def get_test_loader(dataset, height, width, batch_size, workers, testset=None):

    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