Visualize class information in demo.py

Hi, thanks for the amazing work!

I ran demo.py on my own point cloud and successfully generated bounding box predictions. Now I would like to generate a ply in which bounding boxes have different colors based on the predicted class. I’m guessing that can be achieved by modifying the models.dump_helper.dump_results function. Can anyone give me some pointers as to how to achieve this?

Thanks

[IliasMAOUDJ]

Hello again, So I tried something with 2 classes and it seems to work. What I did is: Change the write_oriented_bbox :

def write_oriented_bbox_rotmat(scene_bbox, out_filename, class_ids):
    ... 

    def convert_oriented_box_to_trimesh_fmt(box, class_id):
            ...
            box_trimesh_fmt = trimesh.creation.box(lengths, trns)
            if(class_id==0):
                box_trimesh_fmt.visual.face_colors= trimesh.visual.to_rgba(colors=[255,0,0])
            elif(class_id==1):
                box_trimesh_fmt.visual.face_colors= trimesh.visual.to_rgba(colors=[0,255,0])
            return box_trimesh_fmt

    scene = trimesh.scene.Scene()
    for box, class_id in zip(scene_bbox,class_ids):
        scene.add_geometry(convert_oriented_box_to_trimesh_fmt(box,class_id)) 
    mesh_list = trimesh.util.concatenate(scene.dump())
    # save to ply file    
    mesh_list.export(out_filename, file_type='ply')
    return

and change dump_helper accordingly. Basically, I just keep the predicted_heading_class in an array and put this array as a parameter for the write_oriented_bbox function depending on what we should keep (same indices as obbs, the first parameter).

[IliasMAOUDJ]

#Dump predicted bounding boxes

    if np.sum(objectness_prob>DUMP_CONF_THRESH)>0:
        num_proposal = pred_center.shape[1]
        obbs = []
        class_ids = []
        for j in range(num_proposal):
            obb = config.param2obb(pred_center[i,j,0:3], pred_heading_class[i,j], pred_heading_residual[i,j],
                            pred_size_class[i,j], pred_size_residual[i,j])
            obbs.append(obb)
            class_ids.append(pred_heading_class[i,j])
        if len(obbs)>0:
            obbs = np.vstack(tuple(obbs)) # (num_proposal, 7)
            class_ids = np.vstack(tuple(class_ids))
            pc_util.write_oriented_bbox_rotmat(obbs[objectness_prob>DUMP_CONF_THRESH,:],
                                                os.path.join(dump_dir, '%06d_pred_confident_bbox.ply'%(idx_beg+i)),
                                                class_ids[objectness_prob>DUMP_CONF_THRESH,:])

            pc_util.write_oriented_bbox_rotmat(obbs[np.logical_and(objectness_prob>DUMP_CONF_THRESH, pred_mask[i,:]==1),:],
                                             os.path.join(dump_dir, '%06d_pred_confident_nms_bbox.ply'%(idx_beg+i)),
                                             class_ids[np.logical_and(objectness_prob>DUMP_CONF_THRESH, pred_mask[i,:]==1),:])

            pc_util.write_oriented_bbox_rotmat(obbs[pred_mask[i,:]==1,:], os.path.join(dump_dir, '%06d_pred_nms_bbox.ply'%(idx_beg+i)),class_ids[pred_mask[i,:]==1,:])

            pc_util.write_oriented_bbox_rotmat(obbs, os.path.join(dump_dir, '%06d_pred_bbox.ply'%(idx_beg+i)),class_ids)

`

In the part of the code where all the bbox are computed, just add the “class” argument. It is possible that other parts of my code are changed compared to the original but I haven’t used this code for a year now, so I don’t remember everything.