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 = 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_bin = np.argmax(np.abs(np.dot(kmeans_dict, tmpy)))
ydata_res = tmpy - kmeans_dict[ydata_bin, :]
ydata_bin = ydata_bin * torch.ones(1).long()
ydata_res = torch.from_numpy(ydata_res).float()
ydata = torch.from_numpy(tmpy).float()
label = label * torch.ones(1).long()
sample = {
'xdata': xdata,
'ydata': ydata,
'label': label,
'ydata_bin': ydata_bin,
'ydata_res': ydata_res
}
return sample
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 __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
num_clusters = int(sys.argv[1])
print('num_clusters: ', num_clusters)
kmeans_file = 'data/kmeans_' + str(num_clusters) + '_.pkl'
# setup data loader to access the images
#train_data = ImagesAll(image_path, 'render')
#image_names = np.concatenate(train_data.list_image_names)
image_names = os.listdir(image_path + '/aeroplane')
# get pose targets from training data
bar = progressbar.ProgressBar()
ydata = []
for i in bar(range(len(image_names))):
image_name = image_names[i]
#print(image_name)
_, _, az, el, ct, _ = parse_name(image_name)
#print(az, el, ct)
R = rotation_matrix(az, el, -ct)
y = get_y(R)
ydata.append(y)
ydata = np.stack(ydata)
print('\nData size: ', ydata.shape)
# run kmeans
kmeans = KMeans(num_clusters, verbose=1, n_jobs=10)
kmeans.fit(ydata)
print(kmeans.cluster_centers_)
# save output
fid = open(kmeans_file, 'wb')
pickle.dump(kmeans, fid)