def __getitem__(self, idx):
# return sample with xdata, ydata, label
xdata, ydata, label = [], [], []
for i in range(1):
image_name = self.image_names[i]
label.append(i*torch.ones(1).long())
# read image
#print(self.image_names)
img_pil = Image.open(os.path.join(self.db_path, self.classes[i], image_name))
xdata.append(self.preprocess(img_pil))
# parse image name to get correponding target
#pdb.set_trace()
_, _, az, el, ct, _ = parse_name(image_name)
if self.db_type == 'real':
R = rotation_matrix(az, el, ct)
elif self.db_type == 'render':
R = rotation_matrix(az, el, -ct)
else:
raise NameError('Unknown db_type passed')
if self.ydata_type == 'axis_angle':
tmpy = get_y(R)
elif self.ydata_type == 'quaternion':
tmpy = get_quaternion(R)
else:
raise NameError('Uknown ydata_type passed')
ydata.append(torch.from_numpy(tmpy).float())
xdata = torch.stack(xdata)
ydata = torch.stack(ydata)
label = torch.stack(label)
sample = {'xdata': xdata, 'ydata': ydata, 'label': label}
return sample
def get_residuals(ydata, key_rotations):
ydata_res = np.zeros((ydata.shape[0], len(key_rotations), 3))
for i in range(ydata.shape[0]):
for j in range(len(key_rotations)):
ydata_res[i,
j, :] = get_y(np.dot(key_rotations[j].T,
get_R(ydata[i])))
return ydata_res
def __getitem__(self, idx): # run the item handler of the renderedImages dataset sample = super().__getitem__(idx) # update the ydata target using kmeans dictionary ydata = sample['ydata'].numpy() # rotation matrix ydata_rot = np.stack([get_R(ydata[i]) for i in range(ydata.shape[0])]) sample['ydata_rot'] = torch.from_numpy(ydata_rot).float() # bin part ydata_bin = self.kmeans.predict(ydata) sample['ydata_bin'] = torch.from_numpy(ydata_bin).long() # residual part ydata_res = np.stack([get_y(np.dot(self.rotations_dict[ydata_bin[i]].T, ydata_rot[i])) for i in range(ydata.shape[0])]) sample['ydata_res'] = torch.from_numpy(ydata_res).float() return sample
def __getitem__(self, idx):
# return sample with xdata, ydata, label
image_name = self.image_names[idx]
label = self.labels[idx]
# read image
img_pil = Image.open(
os.path.join(self.db_path, self.classes[label],
image_name + '.png'))
xdata = self.preprocess(img_pil)
# parse image name to get correponding target
_, _, az, el, ct, _ = parse_name(image_name)
R = rotation_matrix(az, el, ct)
if self.ydata_type == 'axis_angle':
tmpy = get_y(R)
elif self.ydata_type == 'quaternion':
tmpy = get_quaternion(R)
else:
raise NameError('Uknown ydata_type passed')
ydata = torch.from_numpy(tmpy).float()
label = label * torch.ones(1).long()
sample = {'xdata': xdata, 'ydata': ydata, 'label': label}
return sample
def __getitem__(self, idx):
# return sample with xdata, ydata, label
xdata, ydata, label = [], [], []
for i in range(self.num_classes):
image_name = self.image_names[i][idx % self.num_images[i]]
label.append(i*torch.ones(1).long())
# read image
img_pil = Image.open(os.path.join(self.db_path, self.classes[i], image_name + '.png'))
xdata.append(preprocess(img_pil))
# parse image name to get correponding target
_, _, az, el, ct, _ = parse_name(image_name)
R = rotation_matrix(az, el, ct)
tmpy = get_y(R)
ydata.append(torch.from_numpy(tmpy).float())
xdata = torch.stack(xdata)
ydata = torch.stack(ydata)
ydata_bin = self.kmeans.predict(ydata.numpy())
ydata_res = ydata.numpy() - self.kmeans.cluster_centers_[ydata_bin, :]
ydata_bin = torch.from_numpy(ydata_bin).long()
ydata_res = torch.from_numpy(ydata_res).float()
label = torch.stack(label)
sample = {'xdata': xdata, 'ydata': ydata, 'label': label, 'ydata_bin': ydata_bin, 'ydata_res': ydata_res}
return sample