def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: (image, target) where target is class_index of the target class.
"""
path, target, cam = self.imgs[index]
img = self.loader(path)
if self.true_pair:
img = self.transform(img)
random_index = list(range(self.len))
random.shuffle(random_index)
for i in random_index:
tpath, ttarget, tcam = self.imgs[i]
if ttarget == target:
timg = self.loader(tpath)
timg = self.transform(timg)
return img, target, path, cam, timg, tcam
if self.pseudo_pair == 0:
img = self.transform(img)
else:
img_list = [self.transform(img)]
generator = transforms.Compose(
[transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
transforms.RandomRotation(10),
my_transforms.RandomCrop(range=(0.70,0.95)),
])
for i in range(self.pseudo_pair-1):
img_list.append(self.transform(generator(img)))
img = torch.stack(img_list, dim=0)
if self.require_path:
return img, target, path, cam
return img, target
args = parser.parse_args()
image_dir = args.dataset_dir
data_transform = transforms.Compose([
transforms.Resize((args.img_h, args.img_w)),
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
transforms.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])
])
data_transform2 = transforms.Compose([
transforms.Resize((args.img_bi_h, args.img_bi_w)),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
transforms.RandomRotation(10),
my_transforms.RandomCrop(range=(0.70, 0.95)),
#transforms.RandomCrop(size=(384, 128)),
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
transforms.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]),
#RandomErasing(probability = 1.0, mean=[0.0, 0.0, 0.0])
])
image_datasets = {}
image_datasets['train'] = DatasetTri(image_dir, data_transform, data_transform2)
dataloaders = torch.utils.data.DataLoader(image_datasets['train'], batch_size=args.batch_size, shuffle=True, num_workers=8)
dataset_sizes = len(image_datasets['train'])
triplet_loss = nn.TripletMarginLoss(margin=3.0)
