def forward(self, seg_pred, vertex_pred, use_argmax=True):
vertex_pred=vertex_pred.permute(0,2,3,1)
b,h,w,vn_2=vertex_pred.shape
vertex_pred=vertex_pred.view(b,h,w,vn_2//2,2)
if use_argmax:
mask=torch.argmax(seg_pred,1)
else:
mask=seg_pred
mean=ransac_voting_layer_v3(mask, vertex_pred, 512, inlier_thresh=0.99)
mean, var=estimate_voting_distribution_with_mean(mask,vertex_pred,mean)
return mean, var
def vote_keypoints(self, pts2d_map, mask):
mask = mask[:, 0] # remove dummy dimension
mask = (mask > 0.5).long(
) # convert to binary and int64 to comply with pvnet interface
pts2d_map = pts2d_map.permute((0, 2, 3, 1))
bs, h, w, num_keypts_2 = pts2d_map.shape
pts2d_map = pts2d_map.view((bs, h, w, num_keypts_2 // 2, 2))
mean = ransac_voting_layer_v3(mask, pts2d_map, 512, inlier_thresh=0.99)
mean, var = estimate_voting_distribution_with_mean(
mask, pts2d_map, mean)
return mean, var
def visualize_voting_ellipse(image, seg_pred, vertex_pred, corner_target):
from lib.ransac_voting_gpu_layer.ransac_voting_gpu import estimate_voting_distribution_with_mean
vertex_pred = vertex_pred.permute(0, 2, 3, 1)
b, h, w, vn_2 = vertex_pred.shape
vertex_pred = vertex_pred.view(b, h, w, vn_2 // 2, 2)
mask = torch.argmax(seg_pred, 1)
mean = ransac_voting_layer_v3(mask, vertex_pred, 512, inlier_thresh=0.99)
mean, var = estimate_voting_distribution_with_mean(mask, vertex_pred, mean)
image = imagenet_to_uint8(image.detach().cpu().numpy())
mean = mean.detach().cpu().numpy()
var = var.detach().cpu().numpy()
corner_target = corner_target.detach().cpu().numpy()
from lib.utils.draw_utils import visualize_voting_ellipse
visualize_voting_ellipse(image, mean, var, corner_target)
