def viz_imgrid(tensors_list, path="tmp.png", fig_height=3, show_axis=False):
row_nb = len(tensors_list)
img_nb = len(tensors_list[0])
fig, axes = plt.subplots(
row_nb,
img_nb,
figsize=(int(img_nb * fig_height), int(row_nb * fig_height)),
)
for row_idx, row_imgs in enumerate(tensors_list):
imgs = [conv2img(tens) for tens in row_imgs]
for img_idx, img in enumerate(imgs):
ax = vizmp.get_axis(axes, row_idx, img_idx, row_nb, img_nb)
ax.imshow(img)
if show_axis:
ax.axis("off")
if path:
fig.savefig(path)
return fig
def add_imgrow(
axes,
tensors,
row_idx=0,
row_nb=1,
overlays=None,
points=None,
over_alpha=0.5,
show_axis=False,
point_s=1,
point_c="r",
overval=1,
):
img_nb = len(tensors)
imgs = [conv2img(tens) for tens in tensors]
for img_idx, img in enumerate(imgs):
ax = vizmp.get_axis(
axes,
row_idx=row_idx,
col_idx=img_idx,
row_nb=row_nb,
col_nb=img_nb,
)
if points is not None:
pts = npt.numpify(points[img_idx])
if pts is not None:
ax.scatter(pts[:, 0], pts[:, 1], s=point_s, c=point_c)
if overlays is not None:
overlay_img = conv2img(overlays[img_idx])
if overlay_img.ndim == 2:
overlay_mask = (overlay_img != overval)[:, :]
else:
overlay_mask = (overlay_img.sum(2) != overval * 3)[
:, :, np.newaxis
]
over_img = (
overlay_mask * img + (1 - overlay_mask) * img * over_alpha
)
ax.imshow(over_img)
else:
ax.imshow(img)
if not show_axis:
ax.axis("off")
def add_pointsrow(
axes,
tensors,
overlay_list=None,
overlay_colors=["c", "k", "b"],
row_idx=0,
row_nb=1,
point_s=1,
point_c="r",
axis_equal=True,
show_axis=True,
alpha=1,
over_alpha=1,
):
point_nb = len(tensors)
points = [conv2img(tens) for tens in tensors]
for point_idx, point in enumerate(points):
ax = vizmp.get_axis(
axes,
row_idx=row_idx,
col_idx=point_idx,
row_nb=row_nb,
col_nb=point_nb,
)
pts = npt.numpify(points[point_idx])
if point_c == "rainbow":
pt_nb = pts.shape[0]
point_c = cm.rainbow(np.linspace(0, 1, pt_nb))
ax.scatter(pts[:, 0], pts[:, 1], s=point_s, c=point_c, alpha=alpha)
if overlay_list is not None:
for overlay, over_color in zip(overlay_list, overlay_colors):
over_pts = npt.numpify(overlay[point_idx])
ax.scatter(
over_pts[:, 0],
over_pts[:, 1],
s=point_s,
c=over_color,
alpha=over_alpha,
)
if axis_equal:
ax.axis("equal")
if not show_axis:
ax.axis("off")
def fitobj2mask(
masks,
bboxes,
obj_paths,
z_off=0.5,
radius=0.1,
faces_per_pixel=1,
lr=0.01,
loss_type="l2",
iters=100,
viz_step=1,
save_folder="tmp/",
viz_rows=12,
crop_box=True,
crop_size=(200, 200),
rot_nb=1,
):
# Initialize logging info
opts = {
"z_off": z_off,
"loss_type": loss_type,
"iters": iters,
"radius": radius,
"lr": lr,
"obj_paths": obj_paths,
"faces_per_pix": faces_per_pixel,
}
results = {"opts": opts}
save_folder = Path(save_folder)
print(f"Saving to {save_folder}")
metrics = defaultdict(list)
batch_size = len(obj_paths)
# Load normalized object
batch_faces = []
batch_verts = []
for obj_path in obj_paths:
verts_loc, faces_idx, _ = py3dload_obj(obj_path)
faces = faces_idx.verts_idx
batch_faces.append(faces.cuda())
verts = normalize.normalize_verts(verts_loc, radius).cuda()
batch_verts.append(verts)
batch_verts = torch.stack(batch_verts)
batch_faces = torch.stack(batch_faces)
# Dummy intrinsic camera
height, width = masks[0].shape
focal = min(masks[0].shape)
camintr = (
torch.Tensor(
[[focal, 0, width // 2], [0, focal, height // 2], [0, 0, 1]]
)
.cuda()
.unsqueeze(0)
.repeat(batch_size, 1, 1)
)
if crop_box:
adaptive_loss = AdaptiveLossFunction(
num_dims=crop_size[0] * crop_size[1],
float_dtype=np.float32,
device="cuda:0",
)
else:
adaptive_loss = AdaptiveLossFunction(
num_dims=height * width, float_dtype=np.float32, device="cuda:0"
)
# Prepare rigid parameters
if rot_nb > 1:
rot_mats = [special_ortho_group.rvs(3) for _ in range(rot_nb)]
rot_vecs = torch.Tensor(
[np.linalg.svd(rot_mat)[0][:2].reshape(-1) for rot_mat in rot_mats]
)
rot_vec = rot_vecs.repeat(batch_size, 1).cuda()
# Ordering b1 rot1, b1 rot2, ..., b2 rot1, ...
else:
rot_vec = torch.Tensor(
[[1, 0, 0, 0, 1, 0] for _ in range(batch_size)]
).cuda()
bboxes_tight = torch.stack(bboxes)
# trans = ops3d.trans_init_from_boxes(bboxes, camintr, (z_off, z_off)).cuda()
trans = ops3d.trans_init_from_boxes_autodepth(
bboxes_tight, camintr, batch_verts, z_guess=z_off
).cuda()
# Repeat to match rots
trans = repeatdim(trans, rot_nb, 1)
bboxes = boxutils.preprocess_boxes(bboxes_tight, padding=10, squarify=True)
if crop_box:
camintr_crop = camutils.get_K_crop_resize(camintr, bboxes, crop_size)
camintr_crop = repeatdim(camintr_crop, rot_nb, 1)
trans.requires_grad = True
rot_vec.requires_grad = True
optim_params = [rot_vec, trans]
if "adapt" in loss_type:
optim_params = optim_params + list(adaptive_loss.parameters())
optimizer = torch.optim.Adam([rot_vec, trans], lr=lr)
ref_masks = torch.stack(masks).cuda()
if crop_box:
ref_masks = cropping.crops(ref_masks.float(), bboxes, crop_size)[:, 0]
# Prepare reference mask
if "dtf" in loss_type:
target_masks = torch.stack(
[torch.Tensor(dtf.distance_transform(mask)) for mask in ref_masks]
).cuda()
else:
target_masks = ref_masks
ref_masks = repeatdim(ref_masks, rot_nb, 1)
target_masks = repeatdim(target_masks, rot_nb, 1)
batch_verts = repeatdim(batch_verts, rot_nb, 1)
batch_faces = repeatdim(batch_faces, rot_nb, 1)
col_nb = 5
fig_res = 1.5
# Aggregate images
clip_data = []
for iter_idx in tqdm(range(iters)):
rot_mat = rotations.compute_rotation_matrix_from_ortho6d(rot_vec)
optim_verts = batch_verts.bmm(rot_mat) + trans.unsqueeze(1)
if crop_box:
rendres = batch_render(
optim_verts,
batch_faces,
K=camintr_crop,
image_sizes=[(crop_size[1], crop_size[0])],
mode="silh",
faces_per_pixel=faces_per_pixel,
)
else:
rendres = batch_render(
optim_verts,
batch_faces,
K=camintr,
image_sizes=[(width, height)],
mode="silh",
faces_per_pixel=faces_per_pixel,
)
optim_masks = rendres[:, :, :, -1]
mask_diff = ref_masks - optim_masks
mask_l2 = (mask_diff ** 2).mean()
mask_l1 = mask_diff.abs().mean()
mask_iou = lyiou.batch_mask_iou(
(optim_masks > 0), (ref_masks > 0)
).mean()
metrics["l1"].append(mask_l1.item())
metrics["l2"].append(mask_l2.item())
metrics["mask"].append(mask_iou.item())
optim_mask_diff = target_masks - optim_masks
if "l2" in loss_type:
loss = (optim_mask_diff ** 2).mean()
elif "l1" in loss_type:
loss = optim_mask_diff.abs().mean()
elif "adapt" in loss_type:
loss = adaptive_loss.lossfun(
optim_mask_diff.view(rot_nb * batch_size, -1)
).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
if iter_idx % viz_step == 0:
row_idxs = np.linspace(
0, batch_size * rot_nb - 1, viz_rows
).astype(np.int)
row_nb = viz_rows
fig, axes = plt.subplots(
row_nb,
col_nb,
figsize=(int(col_nb * fig_res), int(row_nb * fig_res)),
)
for row_idx in range(row_nb):
show_idx = row_idxs[row_idx]
ax = vizmp.get_axis(
axes, row_idx, 0, row_nb=row_nb, col_nb=col_nb
)
ax.imshow(npt.numpify(optim_masks[show_idx]))
ax.set_title("optim mask")
ax = vizmp.get_axis(
axes, row_idx, 1, row_nb=row_nb, col_nb=col_nb
)
ax.imshow(npt.numpify(ref_masks[show_idx]))
ax.set_title("ref mask")
ax = vizmp.get_axis(
axes, row_idx, 2, row_nb=row_nb, col_nb=col_nb
)
ax.imshow(
npt.numpify(ref_masks[show_idx] - optim_masks[show_idx]),
vmin=-1,
vmax=1,
)
ax.set_title("ref masks diff")
ax = vizmp.get_axis(
axes, row_idx, 3, row_nb=row_nb, col_nb=col_nb
)
ax.imshow(npt.numpify(target_masks[show_idx]), vmin=-1, vmax=1)
ax.set_title("target mask")
ax = vizmp.get_axis(
axes, row_idx, 4, row_nb=row_nb, col_nb=col_nb
)
ax.imshow(
npt.numpify(
target_masks[show_idx] - optim_masks[show_idx]
),
vmin=-1,
vmax=1,
)
ax.set_title("masks diff")
viz_folder = save_folder / "viz"
viz_folder.mkdir(parents=True, exist_ok=True)
data = vizmp.fig2np(fig)
clip_data.append(data)
fig.savefig(viz_folder / f"{iter_idx:04d}.png")
clip = mpy.ImageSequenceClip(clip_data, fps=4)
clip.write_videofile(str(viz_folder / "out.mp4"))
clip.write_videofile(str(viz_folder / "out.webm"))
results["metrics"] = metrics
return results
row_nb = args.batch_size
col_nb = 3
diffs = (rendres - img_th[:, :, :, :])[:, :, :, :3]
if args.loss_type == "l1":
loss = diffs.abs().mean()
if args.loss_type == "l2":
loss = (diffs ** 2).sum(-1).mean()
# loss = (rendres - img_th).abs().mean()
optimizer.zero_grad()
loss.backward()
print(loss)
optimizer.step()
if iter_idx % args.viz_step == 0:
fig, axes = plt.subplots(row_nb, col_nb)
for row_idx in range(row_nb):
ax = vizmp.get_axis(
axes, row_idx=row_idx, col_idx=0, row_nb=row_nb, col_nb=col_nb
)
ax.imshow(egoviz.imagify(rendres[row_idx], normalize_colors=False))
ax = vizmp.get_axis(
axes, row_idx=row_idx, col_idx=1, row_nb=row_nb, col_nb=col_nb
)
ax.imshow(
egoviz.imagify(img_th[row_idx][:, :], normalize_colors=False)
)
ax = vizmp.get_axis(
axes, row_idx=row_idx, col_idx=2, row_nb=row_nb, col_nb=col_nb
)
ax.imshow(npt.numpify(diffs[row_idx]))
fig.savefig(f"tmp_{iter_idx:04d}.png", bbox_inches="tight")
def ego_viz(
data,
supervision,
scene_outputs,
loss_metas=None,
fig_res=2,
step_idx=0,
save_folder="tmp",
sample_nb=4,
):
# segm_rend = npt.numpify(scene_outputs["segm_rend"])
viz_rends = [npt.numpify(rend) for rend in scene_outputs["scene_viz_rend"]]
ref_hand_rends = supervision["ref_hand_rends"]
col_nb = 3 + len(viz_rends)
fig, axes = plt.subplots(
sample_nb,
col_nb,
figsize=(int(3 / 2 * col_nb * fig_res), int(sample_nb * fig_res)),
)
scene_size = len(supervision["imgs"])
sample_idxs = np.linspace(0, scene_size - 1, sample_nb).astype(np.int)
mask_diffs = npt.numpify(loss_metas["mask_diffs"])
for row_idx, sample_idx in enumerate(sample_idxs):
img = supervision["imgs"][sample_idx][:, :, ::-1]
# obj_mask = supervision["masks"][sample_idx]
sample_data = data[sample_idx]
# Column 1: image and supervision
ax = vizmp.get_axis(axes,
row_idx=row_idx,
col_idx=0,
row_nb=sample_nb,
col_nb=col_nb)
ax.imshow(img)
ax.axis("off")
add_hands(ax, sample_data)
# Column 2: Rendered prediction on top of GT hands
ax = vizmp.get_axis(axes,
row_idx=row_idx,
col_idx=1,
row_nb=sample_nb,
col_nb=col_nb)
ax.imshow(ref_hand_rends[sample_idx])
ax.imshow(make_alpha(viz_rends[0][sample_idx][:, :, :3]), alpha=0.8)
ax.axis("off")
# Column 3: Rendered Mask
ax = vizmp.get_axis(axes,
row_idx=row_idx,
col_idx=2,
row_nb=sample_nb,
col_nb=col_nb)
mask_diff = mask_diffs[sample_idx]
mask_diff_img = imagify(mask_diff)
ax.imshow(mask_diff_img)
# Column 4+: Rendered prediction
ax.axis("off")
for view_idx, viz_rend in enumerate(viz_rends):
ax = vizmp.get_axis(
axes,
row_idx=row_idx,
col_idx=3 + view_idx,
row_nb=sample_nb,
col_nb=col_nb,
)
if view_idx == 0:
ax.imshow(img)
alpha = 0.8
# Special treatment of first view, which is camera view
else:
alpha = 1
ax.imshow(viz_rend[sample_idx, :, :, :3], alpha=alpha)
ax.axis("off")
os.makedirs(save_folder, exist_ok=True)
save_path = os.path.join(save_folder, f"tmp_{step_idx:04d}.png")
fig.suptitle(f"optim iter : {step_idx}")
fig.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0.1, wspace=0)
fig.gca().xaxis.set_major_locator(ticker.NullLocator())
fig.gca().yaxis.set_major_locator(ticker.NullLocator())
fig.savefig(save_path, bbox_inches="tight")
print(f"Saved to {save_path}")
return save_path
def ego_viz_old(
pred_verts,
pred_proj_verts,
gt_proj_verts,
vert_flags,
imgs=None,
fig_res=2,
cam=None,
faces=None,
step_idx=0,
save_folder="tmp",
):
# Render predicted human
render_verts = pred_verts.cuda()
batch_size = len(render_verts)
faces_th = (faces.unsqueeze(0).repeat(batch_size, 1,
1).to(render_verts.device))
camintr = (cam.get_camintr().to(render_verts.device).unsqueeze(0).repeat(
batch_size, 1, 1))
rot = (cam.get_camrot().unsqueeze(0).repeat(batch_size, 1,
1).to(render_verts.device))
img_size = (imgs[0].shape[1], imgs[0].shape[0])
with torch.no_grad():
rends = batch_render(
render_verts,
faces_th,
K=camintr,
rot=rot,
image_sizes=[img_size for _ in range(batch_size)],
)
show_pred_verts = pred_verts.cpu().detach().numpy()
row_nb = len(pred_verts)
col_nb = 4
fig, axes = plt.subplots(row_nb,
col_nb,
figsize=(int(col_nb * fig_res),
int(row_nb * fig_res)))
for row_idx in range(row_nb):
ax = vizmp.get_axis(axes,
row_idx=row_idx,
col_idx=0,
row_nb=row_nb,
col_nb=col_nb)
super_proj_verts = (gt_proj_verts[row_idx][(vert_flags[row_idx] >
0)].cpu().detach().numpy())
super_pred_proj_verts = (pred_proj_verts[row_idx][(
vert_flags[row_idx] > 0)].cpu().detach().numpy())
if imgs is not None:
ax.imshow(imgs[row_idx])
point_nb = super_pred_proj_verts.shape[0]
colors = cm.rainbow(np.linspace(0, 1, point_nb))
ax.scatter(
super_proj_verts[:, 0],
super_proj_verts[:, 1],
s=0.5,
alpha=0.2,
c="k",
)
ax.scatter(
super_pred_proj_verts[:, 0],
super_pred_proj_verts[:, 1],
s=0.5,
alpha=0.2,
c=colors,
)
ax.axis("equal")
row_pred_verts = show_pred_verts[row_idx]
ax = vizmp.get_axis(axes,
row_idx=row_idx,
col_idx=1,
row_nb=row_nb,
col_nb=col_nb)
ax.scatter(row_pred_verts[:, 0], row_pred_verts[:, 2], s=1)
ax.axis("equal")
ax = vizmp.get_axis(axes,
row_idx=row_idx,
col_idx=2,
row_nb=row_nb,
col_nb=col_nb)
ax.scatter(row_pred_verts[:, 1], row_pred_verts[:, 2], s=1)
ax.axis("equal")
ax = vizmp.get_axis(axes,
row_idx=row_idx,
col_idx=3,
row_nb=row_nb,
col_nb=col_nb)
ax.imshow(imgs[row_idx])
ax.imshow(rends[row_idx].sum(-1).cpu().numpy(), alpha=0.5)
os.makedirs(save_folder, exist_ok=True)
save_path = os.path.join(save_folder, f"tmp_{step_idx:04d}.png")
fig.savefig(save_path)
print(f"Saved to {save_path}")