utils/viz.py

bb = breakpoint
import torch
import numpy as np
from utils.utils import to_tensor, to_numpy
import open3d as o3d
import rerun as rr

OPENCV2OPENGL =  (1,-1,-1,1)

def pts_to_opengl(pts):
    return pts*OPENCV2OPENGL[:3]

def save_pointmaps(data,path='debug',view=False,color='novelty',frustrum_scale=20):
    # debug function to save points to a ply file
    import open3d as o3d
    pointmaps = data['pointmaps']
    B = pointmaps.shape[0]
    W, H = pointmaps.shape[-3:-1]
    n_cams = data['c2ws'].shape[1]
    geometries = []
    for b in range(B):
        geometry_b = []
        points = torch.cat([p.flatten(start_dim=0,end_dim=1) for p in pointmaps[b]],dim=0)
        if view:
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(to_numpy(points))
            if color == 'novelty':
                colors = torch.ones_like(points)
                pts_p_cam = W*H
                # make all novel points red
                colors[pts_p_cam:,1:]*=0.1

                # make all points from first camera blue
                colors[:pts_p_cam,0]*=0.1
                colors[:pts_p_cam,2]*=0.1
                colors*=255.0
            
            else:
                colors = torch.cat([p.flatten(start_dim=0,end_dim=1) for p in data['imgs'][b]],dim=0)
            pcd.colors = o3d.utility.Vector3dVector(to_numpy(colors)/255.0)
            geometry_b.append(pcd)
            origin = o3d.geometry.TriangleMesh.create_coordinate_frame(
                    size=10, origin=[0,0,0])
            geometry_b.append(origin)
            for i in range(n_cams):
                K = data['Ks'][b,i].cpu().numpy()
                K = o3d.camera.PinholeCameraIntrinsic(W,H,K)
                P = data['c2ws'][b,i].cpu().numpy()
                cam_frame = o3d.geometry.LineSet.create_camera_visualization(intrinsic=K,extrinsic=P,scale=frustrum_scale)
                geometry_b.append(cam_frame)
            o3d.visualization.draw_geometries(geometry_b)

            # add point at the origin
            o3d.io.write_point_cloud(f"{path}_{b}.ply", pcd)
            breakpoint()
        geometries.append(geometry_b)
    return geometries

def just_load_viz(pred_dict,gt_dict,batch,name='just_load_viz',addr='localhost:9000',fused_meshes=None,n_points=None):
    rr.init(name)
    rr.connect(addr)
    rr.set_time_seconds("stable_time", 0)

    context_views = batch['input_cams']['pointmaps']
    context_rgbs = batch['input_cams']['imgs']
    gt_pred_views = gt_dict['pointmaps']
    pred_views = pred_dict['pointmaps']
    
    # FIX this weird shape
    pred_masks = batch['new_cams']['valid_masks']
    context_masks = batch['input_cams']['valid_masks']

    B = batch['new_cams']['pointmaps'].shape[0]
    W,H = context_views.shape[-3:-1]
    n_pred_cams = pred_views.shape[1]

    for b in range(B):
        rr.set_time_seconds("stable_time", b)
        # Set world transform to identity (normal origin)
        rr.log("world", rr.Transform3D(translation=[0, 0, 0], mat3x3=np.eye(3)))
        ## show context views
        context_rgb = to_numpy(context_rgbs[b])
        
        for i in range(n_pred_cams):
            if 'conf_pointmaps' in pred_dict:
                conf_pts = pred_dict['conf_pointmaps'][b,i]
                
                #print(f"view {i} mean conf: {mean_conf}, std conf: {std_conf}")
                conf_pts = (conf_pts - conf_pts.min())/(conf_pts.max() - conf_pts.min())
                conf_pts = to_numpy(conf_pts)
                rr.log(f"view_{i}/pred_conf", rr.Image(conf_pts))
            if pred_masks[b,i].sum() == 0:
                continue
            if gt_pred_views is not None:
                gt_pred_pts = gt_pred_views[b,i][pred_masks[b,i]]
                gt_pred_pts = to_numpy(gt_pred_pts)
            else:
                gt_pred_pts = None 

            # red is color for gt points
            if gt_pred_pts is not None:
                color = np.array([1,0,0])
                colors = np.ones_like(gt_pred_pts)
                colors[:,0] = color[0]
                colors[:,1] = color[1]
                colors[:,2] = color[2]
                rr.log(
                    f"world/new_views_gt/view_{i}", rr.Points3D(gt_pred_pts,colors=colors)
                )
            # green is color for pred points
            pred_pts = pred_views[b,i][pred_masks[b,i]]
            pred_pts = to_numpy(pred_pts)

            depth = pred_views[b,i][:,:,2]
            depth -= depth[pred_masks[b,i]].min()
            depth[~pred_masks[b,i]] = 0
            depth /= depth.max()
            depth = to_numpy(depth)
            rr.log(f"world/new_views_pred/view_{i}/image", rr.Image(depth))

            if 'classifier' in pred_dict:
                classifier = (pred_dict['classifier'][b,i] > 0.0).float() # this is assuming the classifier is a sigmoid output
                classifier = to_numpy(classifier)
                rr.log(f"view_{i}/pred_mask", rr.Image(classifier))

            color = np.array([0,1,0])
            colors = np.ones_like(pred_pts)
            colors[:,0] = color[0]
            colors[:,1] = color[1]
            colors[:,2] = color[2]
            if n_points is None:
                rr.log(
                    f"world/new_views_pred/view_{i}/pred_points", rr.Points3D(pred_pts,colors=colors)
                )
            else:
                # randomly sample n_points from pred_pts
                n_points = min(n_points, pred_pts.shape[0])
                inds = np.random.choice(pred_pts.shape[0], n_points, replace=False)
                rr.log(
                    f"world/new_views_pred/view_{i}/pred_points", rr.Points3D(pred_pts[inds],colors=colors[inds])
                )

            K = batch['new_cams']['Ks'][b,i].cpu().numpy()
            P = batch['new_cams']['c2ws'][b,i].cpu().numpy()
            P = np.linalg.inv(P)
            rr.log(f"world/new_views_pred/view_{i}", rr.Transform3D(translation=P[:3,3], mat3x3=P[:3,:3], from_parent=True))
            
            rr.log(f"world/new_views_gt/view_{i}", rr.Transform3D(translation=P[:3,3], mat3x3=P[:3,:3], from_parent=True))

            if 'classifier' in pred_dict:
                classifier = gt_dict['valid_masks'][b,i].float()
                classifier = to_numpy(classifier)
                rr.log(f"view_{i}/gt_mask", rr.Image(classifier))

            rr.log(
                f"world/new_views_pred/view_{i}/image",
                rr.Pinhole(
                    resolution=[H, W],
                    focal_length=[K[0,0], K[1,1]],
                    principal_point=[K[0,2], K[1,2]],
                ),
            )

            rr.log(f"world/new_views_pred/view_{i}/image", rr.Image(to_numpy(pred_masks[b,i].float())))
        n_input_cams = context_masks.shape[1]

        for i in range(n_input_cams):
            context_pts = context_views[b][i][context_masks[b][i]]
            context_pts = to_numpy(context_pts)
            context_pts_rgb = context_rgbs[b][i][context_masks[b][i]]
            context_pts_rgb = to_numpy(context_pts_rgb)
            
            # depth imgs 
            #context_depths = batch['input_cams']['depths'][b][i]
            #context_depths  = (context_depths / context_depths.max() * 255.0).clamp(0,255)
            #context_depths = to_numpy(context_depths).astype(np.uint8)
            rr.log(
                f"world/context_views/view_{i}_points", rr.Points3D(context_pts,colors=(context_pts_rgb/255.0))
            )
            
            K = batch['input_cams']['Ks'][b,i].cpu().numpy()
            P = batch['input_cams']['c2ws'][b,i].cpu().numpy()
            P = np.linalg.inv(P)
            rr.log(f"world/context_views/view_{i}", rr.Transform3D(translation=P[:3,3], mat3x3=P[:3,:3], from_parent=True))

            rr.log(
                f"world/context_views/view_{i}/image",
                rr.Pinhole(
                    resolution=[H, W],
                    focal_length=[K[0,0], K[1,1]],
                    principal_point=[K[0,2], K[1,2]],
                        ),
                   )
            context_rgb_i = context_rgb[i]
            rr.log(
                f"world/context_views/view_{i}/image", rr.Image(context_rgb_i)
            )

            rr.log(
                f"world/context_camera_{i}/mask", rr.Image(to_numpy(context_masks[b,i].float()))
            )
        if fused_meshes is not None:
            rr.log(f"world/fused_mesh", rr.Mesh3D(vertex_positions=fused_meshes[b]['verts'], vertex_normals=fused_meshes[b]['norms'], vertex_colors=fused_meshes[b]['colors'], triangle_indices=fused_meshes[b]['faces']))