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 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)
def forward(self, seg_pred, vertex_pred, use_argmax=True, use_uncertainty=False): 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 if use_uncertainty: return ransac_voting_layer_v5(mask, vertex_pred, 128, inlier_thresh=0.99, max_num=100) else: return ransac_voting_layer_v3(mask, vertex_pred, 128, inlier_thresh=0.99, max_num=100)