def get_state_snapshots(state_info):
rope_pr2_viz = RopePR2Viz()
seg_info, save_dir = state_info
num_segs = len(seg_info)
snapshots = {}
composite = None
for i in xrange(num_segs):
num_mini_segs = len(seg_info[i])
for j in xrange(num_mini_segs):
robot_tfm, robot_dofs, rope_nodes = seg_info[i][j]
rope_pr2_viz.set_robot_pose(robot_dofs, robot_tfm)
if i==0 and j==0:
rope_pr2_viz.update_rope(rope_nodes)
else:
rope_pr2_viz.update_rope_links(rope_nodes)
env_img = rope_pr2_viz.get_env_snapshot()
snapshots['%d_%d.jpg'%(i,j)] = env_img
composite = cpu.tile_images(snapshots, int(math.ceil(len(snapshots)/4.0)), 4, max_width=2500)
return (composite, snapshots, save_dir)
def main(demo_type, n_clusters, n_segs=None, use_clouds=False):
demofile = h5py.File(osp.join(demo_files_dir, demo_type, demo_type+'.h5'), 'r')
iden = ''
if n_segs is not None:
iden = str(n_segs)
if use_clouds:
iden += '_clouds'
smfile = "sim_matrix"+iden+".cp"
try:
with open(smfile, 'r') as f:
sm = pickle.load(f)
seg_num = 0
keys = {}
for demo_name in demofile:
if demo_name != "ar_demo":
for seg_name in demofile[demo_name]:
if seg_name != 'done':
keys[seg_num] = (demo_name, seg_name)
print demo_name, seg_name
seg_num += 1
if n_segs is not None and seg_num >= n_segs: break
except IOError:
if use_clouds:
keys, clouds = extract_clouds(demofile)
sm = similarity_matrix(clouds,n_segs)
else:
keys, segs = extract_segs(demofile)
sm = similarity_matrix_segs(segs,n_segs)
with open(smfile, 'wa') as f:
pickle.dump(sm, f)
print sm
labels = spectral_clustering(sm, n_clusters = n_clusters, eigen_solver='arpack',assign_labels='discretize')
names = {i:[] for i in xrange(args.num_clusters)}
images = {i:[] for i in xrange(args.num_clusters)}
print keys
for i in xrange(len(labels)):
label = labels[i]
names[label].append(keys[i])
images[label].append(np.asarray(demofile[keys[i][0]][keys[i][1]]["rgb"]))
rows = []
i = 0
print "Press q to exit, left/right arrow keys to navigate"
while True:
print "Label %i"%(i+1)
print names[i]
import math
ncols = 7
nrows = int(math.ceil(1.0*len(images[i])/ncols))
row = cpu.tile_images(images[i], nrows, ncols)
rows.append(np.asarray(row))
cv2.imshow("clustering result", row)
kb = cv2.waitKey()
if kb == 1113939:
i = min(i+1,args.num_clusters-1)
elif kb == 1113937:
i = max(i-1,0)
elif kb == 1048689:
break
return
bigimg = cpu.tile_images(rows, len(rows), 50)
cv2.imshow("clustering result", bigimg)
print "press any key to continue"
cv2.waitKey()
names = dict(enumerate(demo_names))
images = {i:[] for i in xrange(num_demos)}
print "Getting images."
for i,dname in names.items():
print "Processing image %i."%(i+1)
images[i] = [np.asarray(demofile[dname][sname]["rgb"]) for sname in demofile[dname]]
print "Done images."
i = 0
inc = True
print "Press q to exit, left/right arrow keys to navigate"
while True:
print
print names[i]
row = cpu.tile_images(images[i], 1, len(images[i]))
cv2.imshow("segments", row)
kb = cv2.waitKey()
if kb == 1113939 or kb == 65363:
i = min(i+1,num_demos-1)
inc = True
elif kb == 1113937 or kb == 65361:
i = max(i-1,0)
inc = False
elif kb == 1048689 or kb == 113:
break
else:
images = [np.asarray(demofile[args.demo_name][sname]["rgb"]) for sname in demofile[args.demo_name]]
row = cpu.tile_images(images, 1, len(images))
def main(demo_type, n_clusters, num_seg=None):
demofile = h5py.File(osp.join(demo_files_dir, demo_type, demo_type+'.h5'), 'r')
print "Loaded file."
iden = ''
if num_seg is not None:
iden = str(num_seg)
cost_file = osp.join(similarity_costs_dir, demo_type)+iden+'.costs'
costs = get_costs(cost_file)
print "Got costs."
seg_num = 0
keys = {}
done = False
for demo_name in demofile:
if demo_name != "ar_demo":
for seg_name in demofile[demo_name]:
if seg_name != 'done':
keys[seg_num] = (demo_name, seg_name)
seg_num += 1
if num_seg is not None and seg_num >= num_seg:
done = True
break
if done:
break
ts = time.time()
mat = generate_sim_matrix(costs, keys)
print 'Time taken to generate sim matrix: %f'%(time.time() - ts)
print mat
ts = time.time()
labels = spectral_clustering(mat, n_clusters = n_clusters, eigen_solver='arpack',assign_labels='discretize')
print 'Time taken to cluster: %f'%(time.time() - ts)
names = {i:[] for i in xrange(args.num_clusters)}
images = {i:[] for i in xrange(args.num_clusters)}
for i in xrange(len(labels)):
label = labels[i]
names[label].append(get_name(keys[i]))
images[label].append(np.asarray(demofile[keys[i][0]][keys[i][1]]["rgb"]))
rows = []
i = 0
inc = True
print "Press q to exit, left/right arrow keys to navigate"
while True:
if len(images[i]) == 0:
if i == n_clusters-1: inc = False
elif i == 0: inc = True
if inc: i = min(i+1,n_clusters-1)
else: i = max(i-1,0)
continue
print "Label %i"%(i+1)
print names[i]
import math
ncols = 7
nrows = int(math.ceil(1.0*len(images[i])/ncols))
row = cpu.tile_images(images[i], nrows, ncols)
rows.append(np.asarray(row))
cv2.imshow("clustering result", row)
kb = cv2.waitKey()
if kb == 1113939 or kb == 65363:
i = min(i+1,args.num_clusters-1)
inc = True
elif kb == 1113937 or kb == 65361:
i = max(i-1,0)
inc = False
elif kb == 1048689 or kb == 113:
break
def cluster_demos (demofile, costs, n_clusters, visualize=False, save_images=False):
seg_num = 0
keys = {}
for demo_name in demofile:
if demo_name != "ar_demo":
for seg_name in demofile[demo_name]:
if seg_name != 'done':
keys[seg_num] = (str(demo_name), str(seg_name))
seg_num += 1
print "Generating sim matrix."
sm = generate_sim_matrix(costs, keys)
print "Getting the cluster rankings"
cdata = cluster_and_rank_demos(sm, n_clusters)
if visualize:
names = {i:[] for i in xrange(args.num_clusters)}
images = {i:[] for i in xrange(args.num_clusters)}
for i in cdata:
names[i] = [get_name(keys[j]) for j in cdata[i]]
images[i] = [np.asarray(demofile[keys[j][0]][keys[j][1]]["rgb"]) for j in cdata[i]]
i = 0
inc = True
print "Press q to exit, left/right arrow keys to navigate"
while True:
if len(images[i]) == 0:
if i == n_clusters-1: inc = False
elif i == 0: inc = True
if inc: i = min(i+1,n_clusters-1)
else: i = max(i-1,0)
continue
print "Label %i"%(i+1)
print names[i]
import math
ncols = 7
nrows = int(math.ceil(1.0*len(images[i])/ncols))
row = cpu.tile_images(images[i], nrows, ncols)
cv2.imshow("clustering result", row)
kb = cv2.waitKey()
if kb == 1113939 or kb == 65363:
i = min(i+1,args.num_clusters-1)
inc = True
elif kb == 1113937 or kb == 65361:
i = max(i-1,0)
inc = False
elif kb == 1048689 or kb == 113:
break
if save_images and yes_or_no("Do you want to save images?"):
cluster_img_dir = '/home/sibi/cluster_image_dir'
if not osp.exists(cluster_img_dir):
os.mkdir(cluster_img_dir)
idx = 1
for i in range(n_clusters):
if len(images[i]) == 0:
continue
print "Label %i"%(i+1)
print names[i]
ncols = 7
nrows = int(math.ceil(1.0*len(images[i])/ncols))
row = cpu.tile_images(images[i], nrows, ncols)
cv2.imwrite(osp.join(cluster_img_dir,'image%02i'%idx), row)
idx += 1
return cdata, keys
def main(demo_type, n_base, n_perts, load_sm = False):
demofile = h5py.File(osp.join(demo_files_dir, demo_type, demo_type+'.h5'), 'r')
if load_sm:
sm_file = osp.join(hd_data_dir, similarity_costs_dir, matrix_file%demo_type)
with open(sm_file, 'r') as f: sm = pickle.load(f)
else:
sm = extract_init(demo_type)
seg_num = 0
keys = {}
for demo_name in demofile:
if demo_name != "ar_demo":
for seg_name in demofile[demo_name]:
if seg_name == 'seg00':
keys[seg_num] = (demo_name, seg_name)
#print demo_name, seg_name
seg_num += 1
n_clusters = n_base - len(BASIC_DEMOS)
labels = spectral_clustering(sm, n_clusters = n_clusters, eigen_solver='arpack',assign_labels='discretize')
names = {i:[] for i in xrange(n_clusters)}
images = {i:[] for i in xrange(n_clusters)}
demos = {i:[] for i in xrange(n_clusters)}
for i in xrange(len(labels)):
label = labels[i]
names[label].append(keys[i])
images[label].append(np.asarray(demofile[keys[i][0]][keys[i][1]]["rgb"]))
demos[label].append(i)
rows = []
i = 0
inc = True
print "Press q to exit, left/right arrow keys to navigate"
while True:
if len(images[i]) == 0:
if i == n_clusters-1: inc = False
elif i == 0: inc = True
if inc: i = min(i+1,n_clusters-1)
else: i = max(i-1,0)
continue
print "Label %i"%(i+1)
print names[i]
import math
ncols = 7
nrows = int(math.ceil(1.0*len(images[i])/ncols))
row = cpu.tile_images(images[i], nrows, ncols)
rows.append(np.asarray(row))
cv2.imshow("clustering result", row)
kb = cv2.waitKey()
if kb == 1113939 or kb == 65363:
i = min(i+1,n_clusters-1)
inc = True
elif kb == 1113937 or kb == 65361:
i = max(i-1,0)
inc = False
elif kb == 1048689 or kb == 113:
break
bests = find_best(demos, sm)
for i in BASIC_DEMOS:
if i in bests:
bests.remove(i)
# add basic demos
bests = BASIC_DEMOS+bests
print bests
best_xyz = []
best_images = {i:None for i in bests}
for best in bests:
xyz = np.asarray(demofile[keys[best][0]][keys[best][1]]["cloud_xyz"])
best_xyz.append(xyz)
best_images[best] = np.asarray(demofile[keys[best][0]][keys[best][1]]["rgb"])
print"Found %i clouds."%len(bests)
print "These are the demos being saved."
ncols = 10
nrows = int(math.ceil(1.0*len(bests)/ncols))
row = cpu.tile_images(best_images.values(), nrows, ncols)
cv2.imshow("best", row)
kb = cv2.waitKey()
if not yes_or_no("Do these look fine to you?"):
return
remaining = n_base-len(bests)
print "remaining:", remaining
while remaining > 0:
fig = pylab.figure()
xyz = best_xyz[remaining-1]
while True:
fig.clf()
perturbed_xyz = sample_random_rope(xyz, True)
ax = fig.gca(projection='3d')
ax.set_autoscale_on(False)
ax.plot(perturbed_xyz[:,0], perturbed_xyz[:,1], perturbed_xyz[:,2], 'o')
fig.show()
cv2.imshow("pert", best_images[bests[remaining-1]])
kb = cv2.waitKey()
if yes_or_no("Does this pert. look fine to you?"):
best_xyz.append(perturbed_xyz)
remaining -= 1
break
if n_perts != 0:
fig = pylab.figure()
fig2 = pylab.figure()
for i in xrange(len(best_xyz)):
xyz = best_xyz[i]
n_p = n_perts
while n_p > 0:
perturbed_xyz = sample_random_rope(xyz, True)
ax = fig.gca(projection='3d')
ax.set_autoscale_on(False)
ax.plot(perturbed_xyz[:,0], perturbed_xyz[:,1], perturbed_xyz[:,2], 'o')
fig.show()
if i < len(bests):
cv2.imshow("pert", best_images[bests[i]])
kb = cv2.waitKey()
else:
ax = fig2.gca(projection='3d')
ax.set_autoscale_on(False)
ax.plot(xyz[:,0], xyz[:,1], xyz[:,2], 'o')
fig2.show()
if not yes_or_no("Does this pert. look fine to you?"):
best_xyz.append(perturbed_xyz)
n_p -= 1
pickle.dump(best_xyz, open(osp.join(demo_files_dir, perturbation_file), 'wa'))
return
