Я пытаюсь создать загрузчик данных, который будет принимать изображения и позы. Изображения сохраняются в виде файлов .jpg, а позы - в виде файлов .npy. Изображения и позы находятся в разных папках, но имеют одинаковую структуру и имя подпапки. Подпапки имеют вид классов, т. Е. Каждый класс имеет соответствующую папку. Я хочу применить преобразования изображений, а затем вернуть изображения (для которых я использую torchvision datasets.ImageFolder). Для поз я использую torchvision datasets.DatasetFolder. Как мне объединить эти два набора данных, чтобы получить одновременно позу и изображение с одним и тем же именем?
class ReIDFolder_images(datasets.ImageFolder):
def __init__(self, root, transform):
super().__init__(root, transform)
targets = np.asarray([s[1] for s in self.samples])
self.targets = targets
self.img_num = len(self.samples)
print(self.img_num)
def _get_cam_id(self, path):
camera_id = []
filename = os.path.basename(path)
camera_id = filename.split('c')[1][0]
return int(camera_id)-1
def _get_pos_sample(self, target, index, path):
pos_index = np.argwhere(self.targets == target)
pos_index = pos_index.flatten()
pos_index = np.setdiff1d(pos_index, index)
if len(pos_index)==0: # in the query set, only one sample
return path
else:
rand = random.randint(0,len(pos_index)-1)
return self.samples[pos_index[rand]][0]
def _get_neg_sample(self, target):
neg_index = np.argwhere(self.targets != target)
neg_index = neg_index.flatten()
rand = random.randint(0,len(neg_index)-1)
return self.samples[neg_index[rand]]
def __getitem__(self, index):
path, target = self.samples[index]
sample = self.loader(path)
pos_path = self._get_pos_sample(target, index, path)
pos = self.loader(pos_path)
if self.transform is not None:
sample = self.transform(sample)
pos = self.transform(pos)
if self.target_transform is not None:
target = self.target_transform(target)
return sample, target, pos
class ReIDFolder_poses(datasets.DatasetFolder):
def __init__(self, root):
super().__init__(root, loader=self.npy_loader, extensions='.npy')
targets = np.asarray([s[1] for s in self.samples])
self.targets = targets
self.img_num = len(self.samples)
print(self.img_num)
def npy_loader(self, path):
sample = torch.Tensor(np.load(path))
return sample
def _get_cam_id(self, path):
camera_id = []
filename = os.path.basename(path)
camera_id = filename.split('c')[1][0]
return int(camera_id)-1
def _get_pos_sample(self, target, index, path):
pos_index = np.argwhere(self.targets == target)
pos_index = pos_index.flatten()
pos_index = np.setdiff1d(pos_index, index)
if len(pos_index)==0: # in the query set, only one sample
return path
else:
rand = random.randint(0,len(pos_index)-1)
return self.samples[pos_index[rand]][0]
def _get_neg_sample(self, target):
neg_index = np.argwhere(self.targets != target)
neg_index = neg_index.flatten()
rand = random.randint(0,len(neg_index)-1)
return self.samples[neg_index[rand]]
def __getitem__(self, index):
path, target = self.samples[index]
sample = self.loader(path)
pos_path = self._get_pos_sample(target, index, path)
pos = self.loader(pos_path)
return sample, target, pos