tarashakhurana / 4d-occ-forecasting

CVPR 2023: Official code for `Point Cloud Forecasting as a Proxy for 4D Occupancy Forecasting'
https://www.cs.cmu.edu/~tkhurana/ff4d/index.html
MIT License
215 stars 24 forks source link

visualization for 4D occupany #6

Closed bbzh closed 1 year ago

bbzh commented 1 year ago

Hi, Thanks for the great work. Would you please release the code for visualizing the 4D occupancy as well?

tarashakhurana commented 1 year ago

Unfortunately, I won't be able to support the visualization code but here is a snippet to threshold the predicted occupancy and visualize it as a point cloud:

def color_by_z(pts, pc_range, cmap):            
    cmap = plt.get_cmap(cmap)                                                                                                                                                                    
    z = pts[:, 2]                                                                                                                                      
    z_min, z_max = pc_range[2], pc_range[5]
    n_z = (z - z_min) / (z_max - z_min)                                                         
    n_z = np.maximum(0, np.minimum(1, n_z))
    return cmap(n_z)[:, :3]

def get_occupancy_as_pcd(pog, thresh, voxel_size, pc_range, cmap):                                 
    def grid_to_pts(X, x_min, y_min, z_min, voxel_size):                                                                                                                                         
        pz, py, px = np.nonzero(X >= thresh)
        xx = px * voxel_size + x_min                                                            
        yy = py * voxel_size + y_min           
        zz = pz * voxel_size + z_min                                                            
        pts = np.stack((xx, yy, zz)).T
        return pts                                                                              

    x_min, y_min, z_min = pc_range[:3]                                                          
    pcds = []                                                                                   
    for t in range(len(pog)):                                                                   
        pred = (pog[t] >= thresh)                                                               
        pts = grid_to_pts(pred, x_min, y_min, z_min, voxel_size)         
        colors = color_by_z(pts, pc_range, cmap)                                                
        pcds.append((pts, colors))                                                       
    return pcds

thresh was set to 0.5. pog should already be getting returned from the model file if you put on the mode as testing.