def __init__(self,
rand_erasing=True,
dataset='Market1501',
data_path='./datasets/market_1501',
batchid=16,
batchimage=4,
batchtest=16,
num_workers=0,
query_image=None,
mode='train',
input_height=384,
input_width=128,
**kwargs):
train_transforms = [
transforms.Resize((input_height, input_width), interpolation=3),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]
if rand_erasing:
train_transforms.append(
RandomErasing(probability=0.5, mean=[0.0, 0.0, 0.0]))
train_transform = transforms.Compose(train_transforms)
test_transform = transforms.Compose([
transforms.Resize((input_height, input_width), interpolation=3),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
if dataset == 'Market1501':
self.trainset = Market1501(train_transform, 'train', data_path)
self.testset = Market1501(test_transform, 'test', data_path)
self.queryset = Market1501(test_transform, 'query', data_path)
elif dataset == 'DeepFashion':
self.trainset = DeepFashion(train_transform, 'train', data_path)
self.testset = DeepFashion(test_transform, 'test', data_path)
self.queryset = DeepFashion(test_transform, 'query', data_path)
else:
raise Exception('Dataset not implemented: %s' % dataset)
self.train_loader = dataloader.DataLoader(
self.trainset,
sampler=RandomSampler(self.trainset,
batch_id=batchid,
batch_image=batchimage),
batch_size=batchid * batchimage,
num_workers=num_workers,
pin_memory=True)
self.test_loader = dataloader.DataLoader(self.testset,
batch_size=batchtest,
num_workers=num_workers,
pin_memory=True)
self.query_loader = dataloader.DataLoader(self.queryset,
batch_size=batchtest,
num_workers=num_workers,
pin_memory=True)
if mode == 'vis':
self.query_image = test_transform(default_loader(query_image))
def __init__(self, dataset="prcc", test=None):
rgb_transform = transforms.Compose(
[transforms.Resize((384, 128), interpolation=3)])
gray_transform = transforms.Compose([
transforms.Resize((384, 128), interpolation=3),
transforms.Grayscale(3)
])
test_transform = transforms.Compose([
transforms.Resize((384, 128), interpolation=3),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
process_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
RandomErasing(probability=0.5, mean=[0.0, 0.0, 0.0])
])
woEr_process_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
if dataset == "prcc":
self.trainset = prcc(data_path=opt.data_path,
dtype="train",
rgb_trans=rgb_transform,
gray_trans=gray_transform,
process_trans=process_transform)
self.trainset_woEr = prcc(data_path=opt.data_path,
dtype="train",
rgb_trans=rgb_transform,
gray_trans=gray_transform,
process_trans=woEr_process_transform)
if test is None:
self.testset = prcc(data_path=opt.data_path,
dtype="test",
test_trans=test_transform)
else:
self.testset = prcc(data_path=opt.data_path + '/' + str(test),
dtype="test",
test_trans=test_transform)
self.queryset = prcc(data_path=opt.data_path,
dtype="query",
test_trans=test_transform)
elif dataset == "ltcc":
self.trainset = ltcc(data_path=opt.data_path,
dtype="train",
rgb_trans=rgb_transform,
gray_trans=gray_transform,
process_trans=process_transform)
self.trainset_woEr = ltcc(data_path=opt.data_path,
dtype="train",
rgb_trans=rgb_transform,
gray_trans=gray_transform,
process_trans=woEr_process_transform)
self.testset = ltcc(data_path=opt.data_path,
dtype="test",
test_trans=test_transform)
self.queryset = ltcc(data_path=opt.data_path,
dtype="query",
test_trans=test_transform)
self.train_loader = dataloader.DataLoader(
self.trainset,
sampler=RandomSampler(self.trainset,
batch_id=opt.batchid,
batch_image=opt.batchimage),
batch_size=opt.batchid * opt.batchimage,
num_workers=0,
pin_memory=True)
self.train_loader_woEr = dataloader.DataLoader(
self.trainset_woEr,
sampler=RandomSampler(self.trainset_woEr,
batch_id=opt.batchid,
batch_image=opt.batchimage),
batch_size=opt.batchid * opt.batchimage,
num_workers=0,
pin_memory=True)
self.test_loader = dataloader.DataLoader(self.testset,
batch_size=opt.batchtest,
num_workers=0,
pin_memory=True)
self.query_loader = dataloader.DataLoader(self.queryset,
batch_size=opt.batchtest,
num_workers=0,
pin_memory=True)