def heatmap(frames, cam, lo_frac=0.5, adapt=True, max_val=35):
""" Set heatmap threshold adaptively, to deal with large variation in possible input videos. """
frames = np.asarray(frames)
max_prob = 0.35
if adapt:
max_val = np.percentile(cam, 97)
same = np.max(cam) - np.min(cam) <= 0.001
if same:
return frames
outs = []
for i in xrange(frames.shape[0]):
lo = lo_frac * max_val
hi = max_val + 0.001
im = frames[i]
f = cam.shape[0] * float(i) / frames.shape[0]
l = int(f)
r = min(1 + l, cam.shape[0] - 1)
p = f - l
frame_cam = ((1 - p) * cam[l]) + (p * cam[r])
frame_cam = ig.scale(frame_cam, im.shape[:2], 1)
#vis = ut.cmap_im(pylab.cm.hot, np.minimum(frame_cam, hi), lo = lo, hi = hi)
vis = ut.cmap_im(pylab.cm.jet, frame_cam, lo=lo, hi=hi)
#p = np.clip((frame_cam - lo)/float(hi - lo), 0, 1.)
p = np.clip((frame_cam - lo) / float(hi - lo), 0, max_prob)
p = p[..., None]
im = np.array(im, 'd')
vis = np.array(vis, 'd')
outs.append(np.uint8(im * (1 - p) + vis * p))
return np.array(outs)
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 vis_example(db_file):
with h5py.File(db_file, 'r') as db:
pre, mid, post = milestone_frames(db)
sc = lambda x: ig.scale(x, (600, None))
im_mid = sc(crop_kinect(ig.uncompress(db['color_image_KinectA'][mid])))
im_post = sc(
crop_kinect(ig.uncompress(db['color_image_KinectA'][post])))
depth = sc(color_depth(crop_kinect(db['depth_image_KinectA'][mid])))
gel_a_0 = sc(ig.uncompress(db['GelSightA_image'][pre]))
gel_b_0 = sc(ig.uncompress(db['GelSightB_image'][pre]))
gel_a_1 = sc(ig.uncompress(db['GelSightA_image'][mid]))
gel_b_1 = sc(ig.uncompress(db['GelSightB_image'][mid]))
row = [
'Color:', im_mid, 'Depth:', depth, 'Gel_A_1:', gel_a_1, 'Gel_B_1:',
gel_b_1, 'Gel_A_0:', gel_a_0, 'Gel_B_0:', gel_b_0, 'Im after:',
im_post, 'Name:',
str(np.array(db['object_name'].value[0])), 'Path:',
db_file.split('/')[-1]
]
return row
def examples_from_db((db_file, im_dir)):
examples = []
try:
with h5py.File(db_file, 'r') as db:
#print db.keys()
sc = lambda x: ig.scale(x, (full_dim, full_dim))
for x in ['A', 'B']:
im_file = ut.make_temp('.png', dir=im_dir)
prev_file = ut.make_temp('.png', dir=im_dir)
ig.save(im_file, sc(db['GelSight%s_image_post_gripping' % x]))
ig.save(prev_file, sc(db['GelSight%s_image_pre_gripping' % x]))
if 'is_gripping' in db:
label = int(np.array(db['is_gripping'])[0])
elif 'Is gripping?' in db:
label = int(np.array(db['Is gripping?'])[0])
else:
raise RuntimeError('No label!')
examples.append((im_file, prev_file, label, db_file))
except:
print 'Failed to open:', db_file
return examples
def crop_from_cam(ims, cam, pr):
cam = np.array([ig.blur(x, 2.) for x in cam])
cam = np.abs(cam)
cam = cam.mean(0)
ims = np.asarray(ims)
y, x = np.nonzero(cam >= cam.max() - 1e-8)
y, x = y[0], x[0]
y = int(round((y + 0.5) * ims.shape[1] / float(cam.shape[0])))
x = int(round((x + 0.5) * ims.shape[2] / float(cam.shape[1])))
d = np.mean(ims.shape[1:3])
h = int(max(224, d // 2.5))
w = int(max(224, d // 2.5))
y0 = int(np.clip(y - h / 2, 0, ims.shape[1] - h))
x0 = int(np.clip(x - w / 2, 0, ims.shape[2] - w))
crop = ims[:, y0:y0 + h, x0:x0 + w]
crop = np.array(
[ig.scale(im, (pr.crop_im_dim, pr.crop_im_dim)) for im in crop])
return crop
def predict_cam(self, ims, samples, n=5, num_times=3):
#num_times = 1
if 1:
f = min(ims.shape[1:3])
ims = np.array([ig.scale(im, (f, f)) for im in ims])
d = int(224. / 256 * ims.shape[1])
print 'd =', d, ims.shape
full = None
count = None
if n == 1:
ys = [ims.shape[1] / 2]
xs = [ims.shape[2] / 2]
else:
ys = np.linspace(0, ims.shape[1] - d, n).astype('int64')
xs = np.linspace(0, ims.shape[2] - d, n).astype('int64')
if num_times == 1:
print 'Using one time'
ts = [0.]
else:
ts = np.linspace(-2, 2., n)
for y in ys:
for x in xs:
crop = ims[:, y:y + d, x:x + d]
crop = resize_nd(
crop,
(crop.shape[0], pr.crop_im_dim, pr.crop_im_dim, 3),
order=1)
for shift in ts:
print x, y, t
snd = sound.Sound(samples, self.pr.samp_sr)
s0 = int(shift * snd.rate)
s1 = s0 + snd.samples.shape[0]
shifted = snd.pad_slice(s0, s1)
assert shifted.samples.shape[0] == snd.samples.shape[0]
[cam] = self.sess.run(
[self.net.vid_net.cam], {
self.ims_ph: crop[None],
self.samples_ph: shifted.samples[None]
})
cam = cam[0, ..., 0]
if full is None:
full = np.zeros(cam.shape[:1] + ims.shape[1:3])
count = np.zeros_like(full)
cam_resized = scipy.ndimage.zoom(
cam,
np.array((full.shape[0], d, d), 'float32') /
np.array(cam.shape, 'float32'))
if 1:
print 'abs'
cam_resized = np.abs(cam_resized)
# print np.abs(cam_resized).max()
frame0 = int(max(-shift, 0) * self.pr.fps)
frame1 = cam_resized.shape[0] - int(
max(shift, 0) * self.pr.fps)
ok = np.ones(count.shape[0])
cam_resized[:frame0] = 0.
cam_resized[frame1:] = 0.
ok[:frame0] = 0
ok[frame1:] = 0
full[:, y:y + d, x:x + d] += cam_resized
count[:, y:y + d, x:x + d] += ok[:, None, None]
assert count.min() > 0
full /= np.maximum(count, 1e-5)
# ut.save('../results/full.pk', full)
# full = ut.load('../results/full.pk')
return full
def format_im(self, im):
return ig.scale(im, (crop_dim, crop_dim), 1) #.astype('float32')
def depth(x):
x = np.array(x).astype('float32')
x = ig.scale(x, (256, 256), 1)
return x
def im(x):
x = np.array(x)
x = ig.scale(x, (256, 256), 1)
return ig.compress(x)
def sc(x):
""" do a center crop (helps with gelsight) """
x = ig.scale(x, (256, 256))
return ut.crop_center(x, 224)
def center_crop(im, full_dim, crop_dim): im = ig.rgb_from_gray(im, remove_alpha = True) im = ig.scale(im, [full_dim]*2, 1) im = ut.crop_center(im, crop_dim) return im
