def get_data(name,
data_dir,
height,
width,
batch_size,
workers,
trainset=False,
flip=False):
dataset = datasets.create(name, data_dir)
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_set = sorted(dataset.train) if trainset else list(
set(dataset.query) | set(dataset.gallery))
test_loader = DataLoader(Preprocessor(test_set,
root=dataset.images_dir,
transform=test_transformer,
flip=flip),
batch_size=batch_size,
num_workers=workers,
shuffle=False,
pin_memory=True)
return dataset, test_loader
def get_test_loader(args,
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
def get_train_loader(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])
transformers = [T.Resize((height, width), interpolation=3),
T.RandomHorizontalFlip(p=0.5),
T.Pad(10),
T.RandomCrop((height, width)),
T.ToTensor(),
normalizer]
train_transformer = T.Compose(transformers)
train_set = dataset.train if trainset is None else trainset
rmgs_flag = num_instances > 0
if rmgs_flag:
sampler = RandomMultipleGallerySampler(train_set, num_instances)
train_loader = IterLoader(
DataLoader(Preprocessor(train_set, root=dataset.images_dir,
transform=train_transformer),
num_workers=workers, pin_memory=True,
batch_sampler=ShuffleBatchSampler(sampler, batch_size, True)), length=iters)
else:
train_loader = IterLoader(
DataLoader(Preprocessor(train_set, root=dataset.images_dir,
transform=train_transformer),
batch_size=batch_size, num_workers=workers, sampler=None,
shuffle=True, pin_memory=True, drop_last=True), length=iters)
return train_loader
def get_train_loader(args,
dataset,
height,
width,
batch_size,
workers,
num_instances,
iters,
epochs,
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,
mutual=True),
batch_size=batch_size,
num_workers=workers,
sampler=sampler,
shuffle=not rmgs_flag,
pin_memory=True,
drop_last=True),
length=iters)
return train_loader
