kienduynguyen / BoxeR

Code release for "BoxeR: Box-Attention for 2D and 3D Transformers"
MIT License
138 stars 22 forks source link

Dataset preparation #25

Open Xiang1004 opened 11 months ago

Xiang1004 commented 11 months ago

Hello, Thank you for share this work! I would like to inquire about which version of the Waymo dataset to download. image

Another question : Analysis and Visualization In the tools/analyze.py, the line (11) from detectron2.utils.logger import setup_logger Where could I find the 'detectron2' folder."

By the way In the create_gt_database.py, the line (11) should be changed from from detector.utils.det3d to from e2edet.utils.det3d.

kienduynguyen commented 10 months ago

Hi, I think I was using v1.2 for Waymo Dataset.

For detectron2, you should install it by yourself following their instruction. But the speed function does not require detectron2 so you can modify them.

Xiang1004 commented 10 months ago

Thank you for your response. I have successfully trained the model and obtained good results in BoxeR-2D.

I would like to inquire if there are any notebooks or tips available for BoxeR-3D visualization?

kienduynguyen commented 10 months ago

We use mayavi which supports visualization within jupyter notebook. However, like other 3D libraries, it needs a bit of effort to install properly :-).

For jupyter notebook:

from mayavi import mlab
mlab.options.offscreen = True

You can just call draw_scenes within jupyter notebook and it should work:

import mayavi.mlab as mlab
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch

from detector.utils.det3d.box_ops import center_to_corner_box2d, boxes_to_corners_3d

box_colormap = ["black", "peru", "red", "black", "purple"]

def visualize_pts(
    pts,
    fig=None,
    bgcolor=(0, 0, 0),
    fgcolor=(1.0, 1.0, 1.0),
    show_intensity=False,
    size=(600, 600),
    draw_origin=True,
):
    if isinstance(pts, torch.Tensor):
        pts = pts.cpu().numpy()

    if fig is None:
        fig = mlab.figure(
            figure=None, bgcolor=bgcolor, fgcolor=fgcolor, engine=None, size=size
        )

    if show_intensity:
        G = mlab.points3d(
            pts[:, 0],
            pts[:, 1],
            pts[:, 2],
            pts[:, 3],
            mode="point",
            colormap="gnuplot",
            scale_factor=1,
            figure=fig,
        )
    else:
        G = mlab.points3d(
            pts[:, 0],
            pts[:, 1],
            pts[:, 2],
            mode="point",
            colormap="gnuplot",
            scale_factor=1,
            figure=fig,
        )
    if draw_origin:
        mlab.points3d(0, 0, 0, color=(1, 1, 1), mode="cube", scale_factor=0.2)
        mlab.plot3d([0, 3], [0, 0], [0, 0], color=(0, 0, 1), tube_radius=0.1)
        mlab.plot3d([0, 0], [0, 3], [0, 0], color=(0, 1, 0), tube_radius=0.1)
        mlab.plot3d([0, 0], [0, 0], [0, 3], color=(1, 0, 0), tube_radius=0.1)

    return fig

def draw_sphere_pts(
    pts, color=(0, 1, 0), fig=None, bgcolor=(0, 0, 0), scale_factor=0.2
):
    if isinstance(pts, torch.Tensor):
        pts = pts.cpu().numpy()

    if fig is None:
        fig = mlab.figure(
            figure=None, bgcolor=bgcolor, fgcolor=None, engine=None, size=(600, 600)
        )

    if isinstance(color, np.ndarray) and color.shape[0] == 1:
        color = color[0]
        color = (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0)

    if isinstance(color, np.ndarray):
        pts_color = np.zeros((pts.__len__(), 4), dtype=np.uint8)
        pts_color[:, 0:3] = color
        pts_color[:, 3] = 255
        G = mlab.points3d(
            pts[:, 0],
            pts[:, 1],
            pts[:, 2],
            np.arange(0, pts_color.__len__()),
            mode="sphere",
            scale_factor=scale_factor,
            figure=fig,
        )
        G.glyph.color_mode = "color_by_scalar"
        G.glyph.scale_mode = "scale_by_vector"
        G.module_manager.scalar_lut_manager.lut.table = pts_color
    else:
        mlab.points3d(
            pts[:, 0],
            pts[:, 1],
            pts[:, 2],
            mode="sphere",
            color=color,
            colormap="gnuplot",
            scale_factor=scale_factor,
            figure=fig,
        )

    mlab.points3d(0, 0, 0, color=(1, 1, 1), mode="cube", scale_factor=0.2)
    mlab.plot3d(
        [0, 3],
        [0, 0],
        [0, 0],
        color=(0, 0, 1),
        line_width=3,
        tube_radius=None,
        figure=fig,
    )
    mlab.plot3d(
        [0, 0],
        [0, 3],
        [0, 0],
        color=(0, 1, 0),
        line_width=3,
        tube_radius=None,
        figure=fig,
    )
    mlab.plot3d(
        [0, 0],
        [0, 0],
        [0, 3],
        color=(1, 0, 0),
        line_width=3,
        tube_radius=None,
        figure=fig,
    )

    return fig

def draw_grid(x1, y1, x2, y2, fig, tube_radius=None, color=(0.5, 0.5, 0.5)):
    mlab.plot3d(
        [x1, x1],
        [y1, y2],
        [0, 0],
        color=color,
        tube_radius=tube_radius,
        line_width=1,
        figure=fig,
    )
    mlab.plot3d(
        [x2, x2],
        [y1, y2],
        [0, 0],
        color=color,
        tube_radius=tube_radius,
        line_width=1,
        figure=fig,
    )
    mlab.plot3d(
        [x1, x2],
        [y1, y1],
        [0, 0],
        color=color,
        tube_radius=tube_radius,
        line_width=1,
        figure=fig,
    )
    mlab.plot3d(
        [x1, x2],
        [y2, y2],
        [0, 0],
        color=color,
        tube_radius=tube_radius,
        line_width=1,
        figure=fig,
    )
    return fig

def draw_multi_grid_range(fig, grid_size=20, bv_range=(-60, -60, 60, 60)):
    for x in range(bv_range[0], bv_range[2], grid_size):
        for y in range(bv_range[1], bv_range[3], grid_size):
            fig = draw_grid(x, y, x + grid_size, y + grid_size, fig)

    return fig

def draw_scenes(
    points, gt_boxes=None, ref_boxes=None, ref_scores=None, ref_labels=None
):
    if isinstance(points, torch.Tensor):
        points = points.cpu().numpy()

    if ref_boxes is not None and not isinstance(ref_boxes, np.ndarray):
        ref_boxes = ref_boxes.cpu().numpy()
    if gt_boxes is not None and not isinstance(gt_boxes, np.ndarray):
        gt_boxes = gt_boxes.cpu().numpy()
    if ref_scores is not None and not isinstance(ref_scores, np.ndarray):
        ref_scores = ref_scores.cpu().numpy()
    if ref_labels is not None and not isinstance(ref_labels, np.ndarray):
        ref_labels = ref_labels.cpu().numpy()

    fig = visualize_pts(points)
    fig = draw_multi_grid_range(fig, bv_range=(0, -40, 80, 40))
    if gt_boxes is not None:
        corners3d = boxes_to_corners_3d(gt_boxes)
        fig = draw_corners3d(corners3d, fig=fig, color=(0, 0, 1), max_num=100)

    if ref_boxes is not None and len(ref_boxes) > 0:
        ref_corners3d = boxes_to_corners_3d(ref_boxes)
        if ref_labels is None:
            fig = draw_corners3d(
                ref_corners3d, fig=fig, color=(0, 1, 0), cls=ref_scores, max_num=100
            )
        else:
            for k in range(ref_labels.min(), ref_labels.max() + 1):
                cur_color = tuple(box_colormap[k % len(box_colormap)])
                mask = ref_labels == k
                fig = draw_corners3d(
                    ref_corners3d[mask],
                    fig=fig,
                    color=cur_color,
                    cls=ref_scores[mask],
                    max_num=100,
                )
    mlab.view(azimuth=-179, elevation=54.0, distance=104.0, roll=90.0)
    return fig

def draw_corners3d(
    corners3d,
    fig,
    color=(1, 1, 1),
    line_width=2,
    cls=None,
    tag="",
    max_num=500,
    tube_radius=None,
):
    """
    :param corners3d: (N, 8, 3)
    :param fig:
    :param color:
    :param line_width:
    :param cls:
    :param tag:
    :param max_num:
    :return:
    """
    import mayavi.mlab as mlab

    num = min(max_num, len(corners3d))
    for n in range(num):
        b = corners3d[n]  # (8, 3)

        if cls is not None:
            if isinstance(cls, np.ndarray):
                mlab.text3d(
                    b[6, 0],
                    b[6, 1],
                    b[6, 2],
                    "%.2f" % cls[n],
                    scale=(0.3, 0.3, 0.3),
                    color=color,
                    figure=fig,
                )
            else:
                mlab.text3d(
                    b[6, 0],
                    b[6, 1],
                    b[6, 2],
                    "%s" % cls[n],
                    scale=(0.3, 0.3, 0.3),
                    color=color,
                    figure=fig,
                )

        for k in range(0, 4):
            i, j = k, (k + 1) % 4
            mlab.plot3d(
                [b[i, 0], b[j, 0]],
                [b[i, 1], b[j, 1]],
                [b[i, 2], b[j, 2]],
                color=color,
                tube_radius=tube_radius,
                line_width=line_width,
                figure=fig,
            )

            i, j = k + 4, (k + 1) % 4 + 4
            mlab.plot3d(
                [b[i, 0], b[j, 0]],
                [b[i, 1], b[j, 1]],
                [b[i, 2], b[j, 2]],
                color=color,
                tube_radius=tube_radius,
                line_width=line_width,
                figure=fig,
            )

            i, j = k, k + 4
            mlab.plot3d(
                [b[i, 0], b[j, 0]],
                [b[i, 1], b[j, 1]],
                [b[i, 2], b[j, 2]],
                color=color,
                tube_radius=tube_radius,
                line_width=line_width,
                figure=fig,
            )

        i, j = 0, 5
        mlab.plot3d(
            [b[i, 0], b[j, 0]],
            [b[i, 1], b[j, 1]],
            [b[i, 2], b[j, 2]],
            color=color,
            tube_radius=tube_radius,
            line_width=line_width,
            figure=fig,
        )
        i, j = 1, 4
        mlab.plot3d(
            [b[i, 0], b[j, 0]],
            [b[i, 1], b[j, 1]],
            [b[i, 2], b[j, 2]],
            color=color,
            tube_radius=tube_radius,
            line_width=line_width,
            figure=fig,
        )

    return fig