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)
