def visualize(self, **kwargs):
o = self.o
im_scale = 1
obj_scale = 1
n = 0 #self.n
self.n = (self.n + 1) % o.N
H, W = o.H * im_scale, o.W * im_scale
h, w = o.h * obj_scale, o.w * obj_scale
if o.v == 2:
save_dir = path.join(o.pic_dir, str(n))
show_dict = {'input': kwargs['X'], 'input_recon': kwargs['X_r']}
if o.bg == 1:
if o.metric == 1:
show_dict['org'] = kwargs['X_org']
att_hor = 1
if att_hor == 1:
att = self.tracker_array.ntm.att.permute(0, 2, 1, 3).reshape(
o.T, self.tracker_array.ntm.ntm_cell.ha, -1)
mem = self.tracker_array.ntm.mem.permute(0, 2, 1, 3).reshape(
o.T, self.tracker_array.ntm.ntm_cell.ha, -1)
else:
att = self.tracker_array.ntm.att.view(
o.T, -1, self.tracker_array.ntm.ntm_cell.wa)
mem = self.tracker_array.ntm.mem.view(
o.T, -1, self.tracker_array.ntm.ntm_cell.wa)
mem_max = 1.8 if o.task == 'mnist' else 3.8 #mem.max()
mem_min = 0 #mem.min()
# print(mem_min, mem_max)
mem = (mem - mem_min) / (mem_max - mem_min + 1e-20)
frames = []
for t in range(0, o.T):
tao = o.batch_id * o.T + t
# Images
for img_kw, img_arg in show_dict.items():
img = img_arg.data[n, t].permute(1, 2, 0).clamp(0, 1)
if o.v == 1:
utils.imshow(img, H, W, img_kw)
else:
if img_kw == 'input' or img_kw == 'org':
utils.mkdir(path.join(save_dir, img_kw))
utils.imwrite(
img, path.join(save_dir, img_kw, "%05d" % (tao)))
# Enforce to show object bounding boxes on the image
if o.metric == 1 and "no_mem" not in o.exp_config:
y_e = kwargs['y_e'].data[n:n + 1].clone().round()
else:
y_e = kwargs['y_e'].data[n:n + 1].clone()
y_e_vis = y_e #kwargs['y_e'].data[n:n+1].clone().fill_(1)
y_l = kwargs['y_l'].data[n:n + 1].clone()
y_p = kwargs['y_p'].data[n:n + 1].clone()
Y_s = kwargs['Y_s'].data[n:n + 1].clone() # 1 * T * O * 1 * h * w
Y_a = kwargs['Y_a'].data[n:n + 1].clone() # 1 * T * O * D * h * w
# Y_s.data[:, :, :, :, 0, :].fill_(1)
# Y_s.data[:, :, :, :, -1, :].fill_(1)
# Y_s.data[:, :, :, :, :, 0].fill_(1)
# Y_s.data[:, :, :, :, :, -1].fill_(1)
# Y_a.data[:, :, :, :, 0, :].fill_(1)
# Y_a.data[:, :, :, :, -1, :].fill_(1)
# Y_a.data[:, :, :, :, :, 0].fill_(1)
# Y_a.data[:, :, :, :, :, -1].fill_(1)
if o.bg == 0:
X_r_vis, _a, X_s_split_vis = self.renderer_vis(
y_e_vis, y_l, y_p, Y_s, Y_a) # 1 * T * D * H * W
else:
Y_b = kwargs['Y_b'].data[n:n + 1].clone()
X_r_vis, _a, X_s_split_vis = self.renderer_vis(
y_e_vis, y_l, y_p, Y_s, Y_a, Y_b=Y_b) # 1 * T * D * H * W
img = X_r_vis.data[0, t, 0:o.D].permute(1, 2, 0).clamp(0, 1)
if o.v == 1:
utils.imshow(img, H, W, 'X_r_vis')
else:
utils.mkdir(path.join(save_dir, 'X_r_vis'))
utils.imwrite(img,
path.join(save_dir, 'X_r_vis', "%05d" % (tao)))
# Objects
y_e, Y_s, Y_a = y_e.data[0, t], Y_s.data[0, t], Y_a.data[
0, t] # O * D * h * w
if o.task == 'mnist':
# Y_o = (y_e.view(-1, 1, 1, 1) * Y_a).permute(2, 0, 3, 1).reshape(o.h, o.O*o.w, o.D)
Y_o_v = (y_e.view(-1, 1, 1, 1) * Y_a).permute(
0, 2, 3, 1).reshape(o.O * o.h, o.w, o.D)
else:
# Y_o = (y_e.view(-1, 1, 1, 1) * Y_s * Y_a).permute(2, 0, 3, 1).reshape(o.h, o.O*o.w, o.D)
Y_o_v = (y_e.view(-1, 1, 1, 1) * Y_a * Y_a).permute(
0, 2, 3, 1).reshape(o.O * o.h, o.w, o.D)
if o.v == 2:
# utils.mkdir(path.join(save_dir, 'Y_o'))
# utils.imwrite(Y_o, path.join(save_dir, 'Y_o', "%05d" % (tao)))
utils.mkdir(path.join(save_dir, 'Y_o_v'))
utils.imwrite(Y_o_v,
path.join(save_dir, 'Y_o_v', "%05d" % (tao)))
# Object masks
X_s_s = torch.stack(X_s_split_vis,
dim=2).data[t,
0].reshape(1, 1, o.O * o.H, o.W)
# Resize and clear noise
X_s_s_pool = (1 - nn.functional.max_pool2d(1 - X_s_s, 2)).reshape(
o.O, o.H // 2, o.W // 2)
frame = {}
frame['timestamp'] = int(t)
frame['num'] = int(tao)
frame['class'] = 'frame'
annotations = []
for j in range(0, o.O):
bin_img = X_s_s_pool[j].cpu().numpy().astype(np.int)
if bin_img.sum() == 0:
continue
else:
annotation = {
'mask': utils.rle_encode(bin_img),
'id': int(o.O * o.batch_id + j),
}
annotations.append(annotation)
frame['annotations'] = annotations
frames.append(frame)
X_s_s_pool = X_s_s_pool.reshape(o.O * o.H // 2, o.W // 2,
1).repeat(1, 1, 3)
if o.v == 2:
utils.mkdir(path.join(save_dir, 'X_s_s'))
utils.imwrite(X_s_s_pool,
path.join(save_dir, 'X_s_s', "%05d" % (tao)))
# Attention and memory
if o.task != 'duke':
cmap = 'hot'
att_c = utils.heatmap(att[t], cmap)
mem_c = utils.heatmap(mem[t], cmap)
if o.v == 1:
sa = 10
utils.imshow(att_c,
att_c.size(0) * sa,
att_c.size(1) * sa, 'att')
utils.imshow(mem_c,
mem_c.size(0) * sa,
mem_c.size(1) * sa, 'mem')
else:
utils.mkdir(path.join(save_dir, 'att'))
utils.mkdir(path.join(save_dir, 'mem'))
utils.imwrite(att_c,
path.join(save_dir, 'att', "%05d" % (tao)))
utils.imwrite(mem_c,
path.join(save_dir, 'mem', "%05d" % (tao)))
return frames
def visualize(self, **kwargs):
o = self.o
im_scale = 1
obj_scale = 1
n = 0 #self.n
self.n = (self.n + 1) % o.N
H, W = o.H * im_scale, o.W * im_scale
h, w = o.h * obj_scale, o.w * obj_scale
if o.v == 2:
save_dir = os.path.join(o.pic_dir, str(n))
show_dict = {'input': kwargs['X'], 'input_recon': kwargs['X_r']}
if o.bg == 1:
if o.metric == 1:
show_dict['org'] = kwargs['X_org']
att_hor = 1
if att_hor == 1:
att = self.tracker_array.ntm.att.permute(
0, 2, 1,
3).contiguous().view(o.T, self.tracker_array.ntm.ntm_cell.ha,
-1)
mem = self.tracker_array.ntm.mem.permute(
0, 2, 1,
3).contiguous().view(o.T, self.tracker_array.ntm.ntm_cell.ha,
-1)
else:
att = self.tracker_array.ntm.att.view(
o.T, -1, self.tracker_array.ntm.ntm_cell.wa)
mem = self.tracker_array.ntm.mem.view(
o.T, -1, self.tracker_array.ntm.ntm_cell.wa)
mem_max = 4 #mem.max()
mem_min = 0 #mem.min()
# print(mem_min, mem_max)
mem = (mem - mem_min) / (mem_max - mem_min + 1e-20)
for t in range(0, o.T):
tao = o.batch_id * o.T + t
# Images
for img_kw, img_arg in show_dict.items():
img = img_arg.data[n, t].permute(1, 2, 0).clamp(0, 1)
if o.v == 1:
utils.imshow(img, H, W, img_kw)
else:
if img_kw == 'input' or img_kw == 'org':
utils.mkdir(os.path.join(save_dir, img_kw))
utils.imwrite(
img, os.path.join(save_dir, img_kw,
"%05d" % (tao)))
# Enforce to show object bounding boxes on the image
if o.metric == 1 and "no_mem" not in o.exp_config:
y_e = Variable(kwargs['y_e'].data[n:n + 1].clone().round())
else:
y_e = Variable(kwargs['y_e'].data[n:n + 1].clone())
y_e_vis = y_e #Variable(kwargs['y_e'].data[n:n+1].clone().fill_(1))
y_l = Variable(kwargs['y_l'].data[n:n + 1].clone())
y_p = Variable(kwargs['y_p'].data[n:n + 1].clone())
Y_s = Variable(
kwargs['Y_s'].data[n:n + 1].clone()) # 1 * T * O * 1 * h * w
Y_a = Variable(
kwargs['Y_a'].data[n:n + 1].clone()) # 1 * T * O * D * h * w
Y_s.data[:, :, :, :, 0, :].fill_(1)
Y_s.data[:, :, :, :, -1, :].fill_(1)
Y_s.data[:, :, :, :, :, 0].fill_(1)
Y_s.data[:, :, :, :, :, -1].fill_(1)
Y_a.data[:, :, :, :, 0, :].fill_(1)
Y_a.data[:, :, :, :, -1, :].fill_(1)
Y_a.data[:, :, :, :, :, 0].fill_(1)
Y_a.data[:, :, :, :, :, -1].fill_(1)
if o.bg == 0:
X_r_vis, _a = self.renderer_vis(y_e_vis, y_l, y_p, Y_s,
Y_a) # 1 * T * D * H * W
else:
Y_b = Variable(kwargs['Y_b'].data[n:n + 1].clone())
X_r_vis, _a = self.renderer_vis(y_e_vis,
y_l,
y_p,
Y_s,
Y_a,
Y_b=Y_b) # 1 * T * D * H * W
img = X_r_vis.data[0, t, 0:o.D].permute(1, 2, 0).clamp(0, 1)
if o.v == 1:
utils.imshow(img, H, W, 'X_r_vis')
else:
utils.mkdir(os.path.join(save_dir, 'X_r_vis'))
utils.imwrite(
img, os.path.join(save_dir, 'X_r_vis', "%05d" % (tao)))
# Objects
y_e, Y_s, Y_a = y_e.data[0, t], Y_s.data[0, t], Y_a.data[
0, t] # O * D * h * w
if o.task == 'mnist':
#print(y_e)
y_e_save = torch.ones_like(y_e)
Y_o = (y_e_save.view(-1, 1, 1, 1) * Y_a).permute(
2, 0, 3, 1).contiguous().view(o.h, o.O * o.w, o.D)
Y_o_v = (y_e_save.view(-1, 1, 1, 1) * Y_a).permute(
0, 2, 3, 1).contiguous().view(o.O * o.h, o.w, o.D)
else:
Y_o = (y_e.view(-1, 1, 1, 1) * Y_s * Y_a).permute(
2, 0, 3, 1).contiguous().view(o.h, o.O * o.w, o.D)
Y_o_v = (y_e.view(-1, 1, 1, 1) * Y_a * Y_a).permute(
0, 2, 3, 1).contiguous().view(o.O * o.h, o.w, o.D)
if o.v == 1:
utils.imshow(Y_o, h, w * o.O, 'Y_o', 1)
else:
utils.mkdir(os.path.join(save_dir, 'Y_o'))
utils.imwrite(Y_o, os.path.join(save_dir, 'Y_o',
"%05d" % (tao)))
# utils.imshow(Y_o_v, h, w * o.O, 'Y_o_v')
# utils.mkdir(path.join(save_dir, 'Y_o_v'))
# utils.imwrite(Y_o_v, path.join(save_dir, 'Y_o_v', "%05d" % (tao)))
# Attention and memory
if o.task != 'duke':
cmap = 'hot'
att_c = utils.heatmap(att[t], cmap)
mem_c = utils.heatmap(mem[t], cmap)
if o.v == 1:
sa = 10
utils.imshow(att_c,
att_c.size(0) * sa,
att_c.size(1) * sa, 'att')
utils.imshow(mem_c,
mem_c.size(0) * sa,
mem_c.size(1) * sa, 'mem')
else:
utils.mkdir(os.path.join(save_dir, 'att'))
utils.mkdir(os.path.join(save_dir, 'mem'))
utils.imwrite(
att_c, os.path.join(save_dir, 'att', "%05d" % (tao)))
utils.imwrite(
mem_c, os.path.join(save_dir, 'mem', "%05d" % (tao)))
if o.v == 2:
if 'X_base_img' in kwargs.keys():
utils.mkdir(os.path.join(save_dir, 'base'))
data, path, actions, phase = kwargs['X_base_img']
torch.save((data[n], path[n], actions[n]),
os.path.join(save_dir, 'base',
str(o.batch_id) + '.pt'))