How to visualize the output bbox

[rzhevcherkasy]

Thanks for your amazing work!
I would like to visualize the predicted bounding box output, yet I dont know how to do it.
I would really appreciate it if anyone can offer some guidance!

[xuxiaoxxxx]

Hi, do you slove this problem?

[ch3cook-fdu]

You can use the codes from https://github.com/ch3cook-fdu/3d-pc-box-viz/blob/main/ply_helper.py to save the bounding boxes as well as the point clouds to .ply files and visualize in meshlab, or use https://github.com/ch3cook-fdu/3d-pc-box-viz/blob/main/o3d_helper.py to render in an open window.

[xuxiaoxxxx]

ok，thanks

[FedericoMafrici]

You can use the codes from https://github.com/ch3cook-fdu/3d-pc-box-viz/blob/main/ply_helper.py to save the bounding boxes as well as the point clouds to .ply files and visualize in meshlab, or use https://github.com/ch3cook-fdu/3d-pc-box-viz/blob/main/o3d_helper.py to render in an open window.

hi, may i have further explanaition on how to do that? i mean what function should i call and at what point of the process? it would be really helpfull for me since i'm trying to visualize the mesh and the bounding box, thank you!

[Pablesky]

Hey, I created a function to create visualizations using Cloud Compare.

I will explain the functionality using the example of the SUNRGD dataset.

First, add this new line to the code in datasets/sunrgbd.py, at the end of the file:

ret_dict = {}
### This is the new line
ret_dict["original_clouds"] = np.load(scan_path + "_pc.npz")["pc"].astype(np.float32)
###
ret_dict["point_clouds"] = point_cloud.astype(np.float32)
ret_dict["gt_box_corners"] = box_corners.astype(np.float32)
ret_dict["gt_box_centers"] = box_centers.astype(np.float32)
ret_dict["gt_box_centers_normalized"] = box_centers_normalized.astype(np.float32)

With this, we will save the original point cloud. Notice that in the case of the SUNRGBD dataset, all the point clouds have the same size. If your dataset has different point cloud sizes, you might encounter some problems.

Then I created this function (rotation is not supported, so only the bounding box is going to be represented):

import torch
import datetime
import logging
import math
import time
import sys
import os
import shutil
import numpy as np

from torch.distributed.distributed_c10d import reduce
from utils.ap_calculator import APCalculator
from utils.ap_calculator import parse_predictions
from utils.ap_calculator import flip_axis_to_depth
import utils.pc_util as pc_util
from utils.pc_util import shift_scale_points
from utils.misc import SmoothedValue
from utils.dist import (
    all_gather_dict,
    all_reduce_average,
    is_primary,
    reduce_dict,
    barrier,
)

@torch.no_grad()
def visualize(
    args,
    model,
    dataset_config,
    dataset_loader,
):

    name_folder = args.checkpoint_dir
    name_folder = os.path.join(name_folder, args.visualize)

    if os.path.exists(name_folder):
        shutil.rmtree(name_folder)

    os.makedirs(name_folder)

    config_dict = {
        'remove_empty_box': True,
        'use_3d_nms': True,
        'nms_iou': 0.25,
        'use_old_type_nms': False,
        'cls_nms': True,
        'per_class_proposal': True,
        'use_cls_confidence_only': False,
        'conf_thresh': 0.25,
        'no_nms': False,
        'dataset_config': dataset_config,
    }

    net_device = next(model.parameters()).device

    model.eval()
    barrier()

    for batch_idx, batch_data_label in enumerate(dataset_loader):
        curr_time = time.time()
        for key in batch_data_label:
            batch_data_label[key] = batch_data_label[key].to(net_device)

        inputs = {
            "point_clouds": batch_data_label["point_clouds"],
            "point_cloud_dims_min": batch_data_label["point_cloud_dims_min"],
            "point_cloud_dims_max": batch_data_label["point_cloud_dims_max"],
        }
        outputs = model(inputs)

        # Memory intensive as it gathers point cloud GT tensor across all ranks
        outputs["outputs"] = all_gather_dict(outputs["outputs"])
        batch_data_label = all_gather_dict(batch_data_label)

        if args.visualize is not None:
            predicted_boxes = outputs['outputs']['box_corners']
            sem_cls_probs = outputs['outputs']['sem_cls_prob']
            objectness_probs = outputs['outputs']['objectness_prob']
            point_cloud = batch_data_label["original_clouds"]
            point_cloud_modelo = batch_data_label["point_clouds"]
            point_cloud_dims_min = batch_data_label["point_cloud_dims_min"]
            point_cloud_dims_max = batch_data_label["point_cloud_dims_max"]
            gt_bounding_boxes = batch_data_label["gt_box_corners"]

            batch_pred_map_cls = parse_predictions(
                predicted_boxes, sem_cls_probs, objectness_probs, point_cloud, config_dict
            )               
            
            # Create a directory for each element in the batch
            for i in range(point_cloud.size(0)):
                element_dir = os.path.join(name_folder, f'element_{batch_idx}_{i}')
                os.makedirs(element_dir, exist_ok=True)

                GT = os.path.join(element_dir, 'GT')
                if os.path.exists(GT)):
                    shutil.rmtree(GT)
                os.makedirs(GT)     
                
                # Save bounding boxes
                for j, pred in enumerate(batch_pred_map_cls[i]):
                    _, box_params, _ = pred
                    bbox_file_path = os.path.join(element_dir, f'bbox_{j}.txt')
                    with open(bbox_file_path, 'w') as bbox_file:
                        for corner in box_params:
                            corner_str = " ".join(map(str, corner))
                            bbox_file.write(f"{corner_str}\n")

                # Save ground truth bounding boxes
                for j, gt in enumerate(gt_bounding_boxes[i]):
                    box_params = gt
                    is_all_zeros = torch.all(box_params == 0)

                    if not is_all_zeros:
                        box_params = box_params.cpu().numpy()
                        bbox_file_path = os.path.join(GT, f'gt_bbox_{j}.txt')
                        with open(bbox_file_path, 'w') as bbox_file:
                            for corner in box_params:
                                corner_str = " ".join(map(str, corner))
                                bbox_file.write(f"{corner_str}\n")

                # Save point cloud data
                pc_file_path = os.path.join(element_dir, 'point_cloud.txt')
                with open(pc_file_path, 'w') as pc_file:
                    pc_data = point_cloud[i].cpu().numpy()

                    pc_data = flip_axis_to_depth(pc_data)
                    pc_data[:, 2] *= -1
                    pc_data[:, 1] *= -1

                    pc_data[:, 3:] = pc_data[:, 3:] * 255
                        
                    for point in pc_data:
                        pc_file.write(" ".join(map(str, point)) + "\n")

                pc_file_path = os.path.join(element_dir, 'point_cloud_modelo.txt')
                with open(pc_file_path, 'w') as pc_file:
                    pc_data = point_cloud_modelo[i].cpu()

                    pc_data = flip_axis_to_depth(pc_data)
                    pc_data[:, 2] *= -1
                    pc_data[:, 1] *= -1
                        
                    for point in pc_data:
                        pc_file.write(" ".join(map(str, point)) + "\n")

Basically, this function will save the points in txt format to be represented using the Cloud Compare tool. With this, you will have a folder inside your checkpoint_dir with all the point clouds saved.

Now on the main.py I add this line, on the evaluate call function. Remember to add the import on the file to your new function ;)

if args.test_only:
        sd = torch.load(args.test_ckpt, map_location=torch.device("cpu"))
        model_no_ddp.load_state_dict(sd["model"])

        if args.visualize is not None:
            visualize(
                args = args,
                model = model,
                dataset_config = dataset_config,
                dataset_loader = dataloaders[dataset_splits[0]],
            )
        
        criterion = None  # faster evaluation
        test_model(args, model, model_no_ddp, criterion, dataset_config, dataloaders)

Hope this helps you. I know it is not the best code, but it is functional.

Regards,
Pablesky.