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, ...])