def get_model_part_iou_rasterization(self, output, target):
if output.shape == (3, 3):
output = output[None]
if target.shape == (3, 3):
target = target[None]
assert output.shape == target.shape, 'output shape does not match target shape'
batch_size = output.shape[0]
fake_frame = torch.ones(
[batch_size, 1, self.frame_height, self.frame_width],
device=output.device)
fake_template = torch.ones(
[batch_size, 1, self.template_height, self.template_width],
device=output.device)
output_mask = warp.warp_image(fake_frame,
output.inverse(),
out_shape=fake_template.shape[-2:])
target_mask = warp.warp_image(fake_frame,
target.inverse(),
out_shape=fake_template.shape[-2:])
output_mask[output_mask > 0] = 1
target_mask[target_mask > 0] = 1
intersection_mask = output_mask * target_mask
output = output_mask.sum(dim=[1, 2, 3])
target = target_mask.sum(dim=[1, 2, 3])
intersection = intersection_mask.sum(dim=[1, 2, 3])
union = output + target - intersection
iou = intersection / (union + self.epsilon)
return utils.to_numpy(iou)
def __getitem__(self, index):
img = self.get_image_by_index(index)
homography = self.get_homography_by_index(index)
warped_template = warp.warp_image(self.template,
homography,
out_shape=img.shape[-2:])[0]
return img, warped_template, homography
def first_order_main_optimization_loop(self, frame, template, optim_tools,
get_corners_fun, corner_to_mat_fun):
loss_hist = []
corners_optim_list = []
optimizer = optim_tools['optimizer']
B = frame.shape[0]
for i in tqdm(range(0, self.opt.optim_iters)):
corners_optim = get_corners_fun()
corners_optim_list.append(corners_optim)
inferred_transformation_mat = corner_to_mat_fun(corners_optim)
warped_tmp = warp.warp_image(template,
inferred_transformation_mat,
out_shape=frame.shape[-2:])
inferred_dist = self.optim_net((frame, warped_tmp))
optim_loss = self.get_loss(inferred_dist,
self.target_dist.repeat(B, 1))
loss_hist.append(optim_loss.clone().detach().cpu().numpy())
if torch.isnan(optim_loss.data):
assert 0, 'loss is nan during optimization'
else:
optimizer.zero_grad()
optim_loss.backward()
optimizer.step()
if optim_loss.data < 0.000000:
break
loss_hist = np.array(loss_hist)
return loss_hist, corners_optim_list
def get_warped_tmp_by_id(self, data_id):
homo_mat = self.get_homography_by_id(data_id)
warped_tmp = warp.warp_image(self.template_torch,
utils.to_torch(homo_mat),
out_shape=(self.frame_h, self.frame_w))
warped_tmp = utils.torch_img_to_np_img(warped_tmp[0])
return warped_tmp
def get_model_whole_iou_rasterization(self, output, target):
'''calculate the model whole iou
'''
if output.shape == (3, 3):
output = output[None]
if target.shape == (3, 3):
target = target[None]
assert output.shape == target.shape, 'output shape does not match target shape'
ZOOM_OUT_SCALE = 4
batch_size = output.shape[0]
fake_template = torch.ones([
batch_size, 1, self.template_height * ZOOM_OUT_SCALE,
self.template_width * ZOOM_OUT_SCALE
],
device=output.device)
scaling_mat = torch.eye(3,
device=output.device).repeat(batch_size, 1, 1)
scaling_mat[:, 0, 0] = scaling_mat[:, 1, 1] = ZOOM_OUT_SCALE
target_mask = warp.warp_image(fake_template,
scaling_mat,
out_shape=fake_template.shape[-2:])
mapping_mat = torch.matmul(output, scaling_mat)
mapping_mat = torch.matmul(target.inverse(), mapping_mat)
output_mask = warp.warp_image(fake_template,
mapping_mat,
out_shape=fake_template.shape[-2:])
output_mask = (output_mask >= 0.5).float()
target_mask = (target_mask >= 0.5).float()
intersection_mask = output_mask * target_mask
output = output_mask.sum(dim=[1, 2, 3])
target = target_mask.sum(dim=[1, 2, 3])
intersection = intersection_mask.sum(dim=[1, 2, 3])
union = output + target - intersection
iou = intersection / (union + self.epsilon)
return utils.to_numpy(iou)
orig_homography, optim_homography = e2e.optim(goal_image[None], template_image)
# reload image and template for visualization
# overload goal image
goal_image_draw = imageio.imread(opt.goal_image_path, pilmode='RGB')
goal_image_draw = goal_image_draw / 255.0
outshape = goal_image_draw.shape[0:2]
# overload template image
template_image_draw = imageio.imread(opt.template_path, pilmode='RGB')
template_image_draw = template_image_draw / 255.0
template_image_draw = image_utils.rgb_template_to_coord_conv_template(template_image_draw)
template_image_draw = utils.np_img_to_torch_img(template_image_draw)
# warp template image with initial guess
warped_tmp_orig = warp.warp_image(template_image_draw, orig_homography, out_shape=outshape)[0]
warped_tmp_orig = utils.torch_img_to_np_img(warped_tmp_orig)
imageio.imwrite('warped_tmp_orig.jpg', warped_tmp_orig)
# plt.imshow(warped_tmp_orig)
# plt.show()
# warp template image with optimized guess
warped_tmp_optim = warp.warp_image(template_image_draw, optim_homography, out_shape=outshape)[0]
warped_tmp_optim = utils.torch_img_to_np_img(warped_tmp_optim)
imageio.imwrite('warped_tmp_optim.jpg', warped_tmp_optim)
# plt.imshow(warped_tmp_optim)
# plt.show()
# show initial guess overlay