def download_pretrained():
# https://github.com/tensorflow/models/tree/master/slim
ut.mkdir('../results')
ut.sys_check(
'wget http://download.tensorflow.org/models/vgg_16_2016_08_28.tar.gz '
' -O ../results/vgg_16_2016_08_28.tar.gz')
ut.sys_check('cd ../results; tar -xzf vgg_16_2016_08_28.tar.gz')
def extract_frames((vid_file, time, label, vid_idx, im_dir, prev_free_time)):
examples = []
with ut.TmpDir() as tmp_dir:
free_dur = 0.1
ut.sys_check('ffmpeg -loglevel warning -ss %f -i "%s" -vf scale=%d:%d -t %f -r %d "%s/%%07d.png"' % \
(prev_free_time, vid_file, full_dim, full_dim, free_dur, sample_fps, tmp_dir))
#fname = sorted(ut.glob(pj(tmp_dir, '*.png')))[0]
assert len(ut.glob(pj(tmp_dir,
'*.png'))), 'no frames for prev_free_time'
fname = random.choice(sorted(ut.glob(pj(tmp_dir, '*.png'))))
prev_file = pj(im_dir, 'prev_%s_%05d_%d.png' % (vid_idx, 0, label))
ut.sys_check('cp %s %s' % (fname, prev_file))
with ut.TmpDir() as tmp_dir:
# ut.sys_check('ffmpeg -i "%s" -vf scale=%d:%d -ss %f -t %f -r %d "%s/%%07d.png"' % \
# (vid_file, full_dim, full_dim, time,
# sample_dur_secs, sample_fps, tmp_dir))
ut.sys_check('ffmpeg -loglevel warning -ss %f -i "%s" -vf scale=%d:%d -t %f -r %d "%s/%%07d.png"' % \
(time, vid_file, full_dim, full_dim,
sample_dur_secs, sample_fps, tmp_dir))
for frame_idx, fname in enumerate(sorted(ut.glob(pj(tmp_dir,
'*.png')))):
im_file = pj(im_dir,
'%s_%05d_%d.png' % (vid_idx, frame_idx, label))
ut.sys_check('cp %s %s' % (fname, im_file))
examples.append((im_file, prev_file, label, vid_file))
return examples
def run(vid_file,
start_time,
dur,
pr,
gpu,
buf=0.05,
mask=None,
arg=None,
net=None):
print pr
dur = dur + buf
with ut.TmpDir() as vid_path:
height_s = '-vf "scale=-2:\'min(%d,ih)\'"' % arg.max_full_height if arg.max_full_height > 0 else ''
ut.sys_check(
ut.frm(
'ffmpeg -loglevel error -ss %(start_time)s -i "%(vid_file)s" -safe 0 '
'-t %(dur)s -r %(pr.fps)s -vf scale=256:256 "%(vid_path)s/small_%%04d.png"'
))
ut.sys_check(
ut.frm(
'ffmpeg -loglevel error -ss %(start_time)s -i "%(vid_file)s" -safe 0 '
'-t %(dur)s -r %(pr.fps)s %(height_s)s "%(vid_path)s/full_%%04d.png"'
))
ut.sys_check(
ut.frm(
'ffmpeg -loglevel error -ss %(start_time)s -i "%(vid_file)s" -safe 0 '
'-t %(dur)s -ar %(pr.samp_sr)s -ac 2 "%(vid_path)s/sound.wav"')
)
if arg.fullres:
fulls = map(
ig.load,
sorted(ut.glob(vid_path, 'full_*.png'))[:pr.sampled_frames])
fulls = np.array(fulls)
snd = sound.load_sound(pj(vid_path, 'sound.wav'))
samples_orig = snd.normalized().samples
samples_orig = samples_orig[:pr.num_samples]
samples_src = samples_orig.copy()
if samples_src.shape[0] < pr.num_samples:
return None
ims = map(ig.load, sorted(ut.glob(vid_path, 'small_*.png')))
ims = np.array(ims)
d = 224
y = x = ims.shape[1] / 2 - d / 2
ims = ims[:, y:y + d, x:x + d]
ims = ims[:pr.sampled_frames]
if mask == 'l':
ims[:, :, :ims.shape[2] / 2] = 128
if arg.fullres:
fulls[:, :, :fulls.shape[2] / 2] = 128
elif mask == 'r':
ims[:, :, ims.shape[2] / 2:] = 128
if arg.fullres:
fulls[:, :, fulls.shape[2] / 2:] = 128
elif mask is None:
pass
else:
raise RuntimeError()
samples_src = mu.normalize_rms_np(samples_src[None], pr.input_rms)[0]
net.init()
ret = net.predict(ims[None], samples_src[None])
samples_pred_fg = ret['samples_pred_fg'][0][:, None]
samples_pred_bg = ret['samples_pred_bg'][0][:, None]
spec_pred_fg = ret['spec_pred_fg'][0]
spec_pred_bg = ret['spec_pred_bg'][0]
print spec_pred_bg.shape
spec_mix = ret['spec_mix'][0]
if arg.cam:
cam, vis = find_cam(fulls, samples_orig, arg)
else:
if arg.fullres:
vis = fulls
else:
vis = ims
return dict(ims=vis,
samples_pred_fg=samples_pred_fg,
samples_pred_bg=samples_pred_bg,
samples_mix=ret['samples_mix'][0],
samples_src=samples_src,
spec_pred_fg=spec_pred_fg,
spec_pred_bg=spec_pred_bg,
spec_mix=spec_mix)
def write_data(out_dir,
rebalance_data=True,
train_frac=0.75,
val_frac=0.0,
n=None,
seed=0):
#def write_data(out_dir, rebalance_data = True, train_frac = 0.75, val_frac = 0.0, n = 10):
assert not os.path.exists(out_dir)
ut.mkdir(out_dir)
base_data = '../data/grasp/'
ut.sys_check('find -L %s -name "*.hdf5" > %s/all_db_files.txt' %
(base_data, out_dir))
all_db_files = map(os.path.abspath,
ut.read_lines(pj(out_dir, 'all_db_files.txt'))[:n])
all_db_files = ut.shuffled_with_seed(all_db_files, seed)
all_db_files = filter(db_ok, all_db_files)
ut.write_lines(pj(out_dir, 'db_files.txt'), all_db_files)
by_name = ut.accum_dict((name_from_file(x), x) for x in all_db_files)
names = ut.shuffled_with_seed(sorted(by_name.keys()), seed)
num_names = len(names)
num_train = int(train_frac * num_names)
num_val = int(val_frac * num_names)
i = 0
train_names = names[i:num_train]
i += num_train
val_names = names[i:i + num_val]
i += num_val
test_names = names[i:]
print num_train, num_val, len(test_names)
splits = [('train', train_names), ('val', val_names), ('test', test_names)]
print 'Number of objects in each split:'
for s, o in splits:
print s, '->', len(o)
#press_clf = press.NetClf(press_model_file, gpu = write_data_gpu)
press_clf = None #press.NetClf(press_model_file, gpu = write_data_gpu)
for dset_name, names in splits:
ut.write_lines(pj(out_dir, '%s_objects.txt' % dset_name), names)
tf_file = pj(out_dir, '%s.tf' % dset_name)
pk_file = pj(out_dir, '%s.pk' % dset_name)
full_pk_file = pj(out_dir, 'full_%s.pk' % dset_name)
if os.path.exists(tf_file):
os.remove(tf_file)
writer = tf.python_io.TFRecordWriter(tf_file)
split_db_files = ut.flatten(by_name[name] for name in names)
split_db_files = ut.shuffled_with_seed(split_db_files, dset_name)
data = []
for db_file in ut.time_est(split_db_files):
with h5py.File(db_file, 'r') as db:
#print 'keys =', db.keys()
def im(x, crop=False, compress=True):
x = ig.uncompress(x)
x = np.array(x)
if crop:
x = crop_kinect(x)
#ig.show(x)
x = ig.scale(x, (256, 256), 1)
if compress:
x = ig.compress(x)
return x
def depth(x):
x = np.array(x).astype('float32')
x = ig.scale(x, (256, 256), 1)
return x
def parse_ee(x):
names = [
'angle_of_EE_at_grasping', 'location_of_EE_at_grasping'
]
vs = [x[name].value for name in names]
ee = np.concatenate([np.array(v).flatten()
for v in vs]).astype('float32')
return ee
label_file = pj(
label_path,
db_file.split('/')[-1].replace('.hdf5', '.txt'))
if os.path.exists(label_file):
print 'Reading label from file'
is_gripping = bool(ut.read_file(label_file))
else:
is_gripping = int(np.array(db['is_gripping']))
pre, mid, _ = milestone_frames(db)
# Estimate the probability that the robot is initially gripping the object
if 0:
press_a = press_clf.predict(
im(db['/GelSightA_image'].value[mid], compress=False),
im(db['/GelSightA_image'].value[pre], compress=False))
press_b = press_clf.predict(
im(db['/GelSightB_image'].value[mid], compress=False),
im(db['/GelSightB_image'].value[pre], compress=False))
initial_press_prob = 0.5 * (press_a + press_b)
else:
initial_press_prob = np.float32(-1.)
#print initial_press_prob, ig.show(im(db['/GelSightA_image'].value[mid], compress = False))
d = dict(
gel0_pre=im(db['/GelSightA_image'].value[pre]),
gel1_pre=im(db['/GelSightB_image'].value[pre]),
gel0_post=im(db['/GelSightA_image'].value[mid]),
gel1_post=im(db['/GelSightB_image'].value[mid]),
im0_pre=im(db['/color_image_KinectA'].value[pre],
crop=True),
im0_post=im(db['/color_image_KinectA'].value[mid],
crop=True),
im1_pre=im(db['/color_image_KinectB'].value[pre],
crop=True),
im1_post=im(db['/color_image_KinectB'].value[mid],
crop=True),
depth0_pre=depth(
crop_kinect(db['/depth_image_KinectA'].value[pre])),
depth0_post=depth(
crop_kinect(db['/depth_image_KinectA'].value[mid])),
initial_press_prob=initial_press_prob,
is_gripping=int(is_gripping),
end_effector=parse_ee(db),
object_name=str(np.array(db['object_name'].value)[0]),
db_file=db_file)
data.append(d)
# for db files
ut.save(full_pk_file, data)
# rebalance data?
if rebalance_data:
by_label = [[], []]
for x in ut.shuffled_with_seed(data, 'rebalance1'):
by_label[x['is_gripping']].append(x)
n = min(map(len, by_label))
print len(data), 'before rebalance'
data = ut.shuffled_with_seed(by_label[0][:n] + by_label[1][:n],
'rebalance2')
print len(data), 'after rebalance'
writer = tf.python_io.TFRecordWriter(tf_file)
for d in data:
fbl = lambda x: tf.train.Feature(bytes_list=tf.train.BytesList(
value=[x]))
fl = lambda x: tf.train.Feature(float_list=tf.train.FloatList(
value=map(float, x.flatten())))
il = lambda x: tf.train.Feature(int64_list=tf.train.Int64List(value
=x))
feat = {
'gel0_pre': fbl(d['gel0_pre']),
'gel1_pre': fbl(d['gel1_pre']),
'gel0_post': fbl(d['gel0_post']),
'gel1_post': fbl(d['gel1_post']),
'im0_pre': fbl(d['im0_pre']),
'im0_post': fbl(d['im0_post']),
'im1_pre': fbl(d['im1_pre']),
'im1_post': fbl(d['im1_post']),
'depth0_pre': fl(d['depth0_pre']),
'depth0_post': fl(d['depth0_post']),
'end_effector': fl(d['end_effector']),
'initial_press_prob': fl(d['initial_press_prob']),
'is_gripping': il([d['is_gripping']])
}
ex = tf.train.Example(features=tf.train.Features(feature=feat))
writer.write(ex.SerializeToString())
writer.close()
ut.save(pk_file, data)
print dset_name, '->', len(data), 'examples'
def write_data(out_dir, train_frac=0.75, val_frac=0.05):
ut.mkdir(out_dir)
base_data = '../data/grasp/'
ut.sys_check('find %s -name "*.hdf5" > %s/db_files.txt' %
(base_data, out_dir))
all_db_files = ut.read_lines(pj(out_dir, 'db_files.txt'))
all_db_files = ut.shuffled_with_seed(all_db_files)
name_from_file = lambda x: '_'.join(x.split('/')[-1].split('_')[2:])
by_name = ut.accum_dict((name_from_file(x), x) for x in all_db_files)
names = ut.shuffled_with_seed(sorted(by_name.keys()))
num_names = len(all_db_files)
num_train = int(train_frac * num_names)
num_val = int(val_frac * num_names)
i = 0
train_names = names[i:num_train]
i += num_train
val_names = names[i:i + num_val]
i += num_val
test_names = names[i:]
for dset_name, names in [('train', train_names), ('val', val_names),
('test', test_names)]:
ut.write_lines(pj(out_dir, '%s_objects.txt' % dset_name), names)
tf_file = pj(out_dir, '%s.tf' % dset_name)
pk_file = pj(out_dir, '%s.pk' % dset_name)
if os.path.exists(tf_file):
os.remove(tf_file)
writer = tf.python_io.TFRecordWriter(tf_file)
data = []
for name in names:
for db_file in by_name[name]:
with h5py.File(db_file, 'r') as db:
def im(x):
x = np.array(x)
x = ig.scale(x, (256, 256), 1)
return ig.compress(x)
if 'is_gripping' in db:
label = int(np.array(db['is_gripping']))
elif 'Is gripping?' in db:
label = int(np.array(db['Is gripping?']))
else:
print 'Skipping: %s. Missing is_gripping' % db_file
print 'Keys:', ' '.join(db.keys())
continue
data.append({
'gel0_pre':
im(db['GelSightA_image_pre_gripping']),
'gel1_pre':
im(db['GelSightB_image_pre_gripping']),
'gel0_post':
im(db['GelSightA_image_post_gripping']),
'gel1_post':
im(db['GelSightB_image_post_gripping']),
'is_gripping':
label
})
fbl = lambda x: tf.train.Feature(bytes_list=tf.train.
BytesList(value=[x]))
feat = {
'gel0_pre':
fbl(im(db['GelSightA_image_pre_gripping'])),
'gel1_pre':
fbl(im(db['GelSightB_image_pre_gripping'])),
'gel0_post':
fbl(im(db['GelSightA_image_post_gripping'])),
'gel1_post':
fbl(im(db['GelSightB_image_post_gripping'])),
'is_gripping':
tf.train.Feature(int64_list=tf.train.Int64List(
value=[label]))
}
ex = tf.train.Example(features=tf.train.Features(
feature=feat))
writer.write(ex.SerializeToString())
writer.close()
ut.save(pk_file, data)
print dset_name, '->', len(data), 'examples'