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 demo():
net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2)
net = NetWrapper(net).cuda()
net = DataParallel(net)
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir = os.path.join(cfg.MODEL_DIR, "cat_demo")
load_model(net.module.net, optimizer, model_dir, -1)
data, points_3d, bb8_3d = read_data()
#print("BB8_3D: ",bb8_3d)
image, mask, vertex, vertex_weights, pose, corner_target = [
d.unsqueeze(0).cuda() for d in data
]
seg_pred, vertex_pred, loss_seg, loss_vertex, precision, recall = net(
image, mask, vertex, vertex_weights)
seg_mask = torch.argmax(seg_pred, 1)
print("seg_mask", seg_mask, type(seg_mask), seg_mask.shape, seg_mask[0])
visualize_mask(seg_mask)
visualize_mask(mask)
#visualize_vertex(vertex, vertex_weights)
#visualize_hypothesis(image, seg_pred, vertex_pred, corner_target)
visualize_voting_ellipse(image, seg_pred, vertex_pred, corner_target)
eval_net = DataParallel(EvalWrapper().cuda())
corner_pred = eval_net(seg_pred, vertex_pred).cpu().detach().numpy()[0]
print("Corner Predictions: ", corner_pred)
camera_matrix = np.array([[572.4114, 0., 325.2611],
[0., 573.57043, 242.04899], [0., 0., 1.]])
pose_pred = pnp(points_3d, corner_pred, camera_matrix)
projector = Projector()
bb8_2d_pred = projector.project(bb8_3d, pose_pred, 'linemod')
print("Pose prediction :\n", pose_pred)
print("GT pose: \n", pose[0].detach().cpu().numpy())
bb8_2d_gt = projector.project(bb8_3d, pose[0].detach().cpu().numpy(),
'linemod')
print(bb8_2d_gt)
image = imagenet_to_uint8(image.detach().cpu().numpy())[0]
visualize_bounding_box(image[None, ...], bb8_2d_pred[None, None, ...],
bb8_2d_gt[None, None, ...])