From 5cbc975abccf152eef07d660c473d7717b9cc6f6 Mon Sep 17 00:00:00 2001 From: Abyssaledge <1060056270@qq.com> Date: Fri, 1 Sep 2023 18:58:19 +0800 Subject: [PATCH] update visualization scripts --- tools/vis/show_bin.py | 90 +++++++++++++++++++ tools/vis/utils.py | 125 ++++++++++++++++++++++++++ tools/vis/visualizer.py | 191 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 406 insertions(+) create mode 100644 tools/vis/show_bin.py create mode 100644 tools/vis/utils.py create mode 100644 tools/vis/visualizer.py diff --git a/tools/vis/show_bin.py b/tools/vis/show_bin.py new file mode 100644 index 0000000..c589a7d --- /dev/null +++ b/tools/vis/show_bin.py @@ -0,0 +1,90 @@ +import numpy as np +import os +from os import path as osp +import argparse +from ipdb import set_trace +from tqdm import tqdm + +# from pipeline_vis import frame_visualization +from visualizer import Visualizer2D +from utils import get_obj_dict_from_bin_file, get_pc_from_time_stamp + +parser = argparse.ArgumentParser() +# running configurations +parser.add_argument('--bin-path', type=str, default='') +parser.add_argument('--gt-bin-path', type=str, default='./data/waymo/waymo_format/gt.bin') +parser.add_argument('--save-folder', type=str, default='') +parser.add_argument('--suffix', type=str, default='') +parser.add_argument('--split', type=str, default='training') +parser.add_argument('--interval', type=int, default=198) +parser.add_argument('--no-gt', action='store_true') +# process +args = parser.parse_args() + +def frame_visualization(pc, dets, gts, name='', save_path='./exp.png', figsize=(40, 40)): + visualizer = Visualizer2D(name=name, figsize=figsize) + visualizer.handler_pc(pc, s=0.1) + if gts is not None: + for _, bbox in enumerate(gts): + visualizer.handler_box(bbox, message='', color='black') + for det in dets: + # visualizer.handler_box(det, message='%.2f-' % det.s + str(int(det.type)), color='red', linestyle='dashed', text_color='green', fontsize='small') + visualizer.handler_box(det, message='%.2f' % det.s, color='red', linestyle='dashed', text_color='blue', fontsize='small', center_message=True) + visualizer.save(save_path) + visualizer.close() + +if __name__ == '__main__': + bin_path = osp.abspath(args.bin_path) + if args.save_folder == '': + save_folder = osp.join(osp.dirname(bin_path), 'vis_folder') + elif '/' in args.save_folder: + save_folder = args.save_folder + else: + save_folder = osp.join(osp.dirname(bin_path), args.save_folder) + + assert 'vis' in save_folder + + + os.makedirs(save_folder, exist_ok=True) + + pred_dict = get_obj_dict_from_bin_file(bin_path, debug=False) + if args.no_gt: + gt_dict = None + else: + gt_dict = get_obj_dict_from_bin_file(args.gt_bin_path, debug=False) + + if gt_dict is not None: + ts_list = sorted(list(gt_dict.keys())) + else: + ts_list = sorted(list(pred_dict.keys())) + + # with open + + for i, ts in tqdm(enumerate(ts_list)): + if i % args.interval != 0: + continue + try: + dets = pred_dict[ts] + except KeyError: + continue + + if gt_dict is None: + gts = None + elif ts not in gt_dict: + gts = [] + else: + gts = gt_dict[ts] + # set_trace() + pc = get_pc_from_time_stamp(ts, './data/waymo/kitti_format/idx2timestamp.pkl', split=args.split) + # set_trace() + + if len(args.suffix) > 0: + suffix = '_' + args.suffix + else: + suffix = '' + + frame_visualization( + pc, dets, gts, + save_path=osp.join(save_folder, str(ts) + suffix + '.png'), + figsize=(18, 18) + ) \ No newline at end of file diff --git a/tools/vis/utils.py b/tools/vis/utils.py new file mode 100644 index 0000000..d0c34cc --- /dev/null +++ b/tools/vis/utils.py @@ -0,0 +1,125 @@ +import os +import numpy as np +import pickle as pkl +from tqdm import tqdm +from collections import defaultdict + +from waymo_open_dataset import label_pb2 +from waymo_open_dataset.utils import box_utils +from waymo_open_dataset.protos import metrics_pb2 + +from ipdb import set_trace +from visualizer import BBox + +import torch + +def read_bin(file_path): + with open(file_path, 'rb') as f: + objects = metrics_pb2.Objects() + objects.ParseFromString(f.read()) + return objects + +def object2array(obj): + """transform box dict in waymo_open_format to array + Args: + box_dict ([dict]): waymo_open_dataset formatted bbox + """ + box = obj.object.box + result = np.array([ + box.center_x, + box.center_y, + box.center_z, + box.width, + box.length, + box.height, + box.heading, + obj.score, + float(obj.object.type), + ]) + return result + +def object2BBox(obj): + box = obj.object.box + result = BBox( + box.center_x, + box.center_y, + box.center_z, + box.height, + box.width, + box.length, + box.heading, + obj.score, + float(obj.object.type), + ) + return result + +def object2mmdetformat(obj): + ''' + According to https://github.com/waymo-research/waymo-open-dataset/blob/master/waymo_open_dataset/label.proto#L33 + and the definition of LiDARInstance3DBoxes + ''' + box = obj.object.box + heading = box.heading - 1.5 * np.pi + + while heading < -np.pi: + heading += 2 * np.pi + while heading > np.pi: + heading -= 2 * np.pi + + result = np.array( + [ + box.center_x, + box.center_y, + box.center_z, + box.width, + box.length, + box.height, + heading, + obj.score, + float(obj.object.type), + ] + ) + return result + +def get_obj_dict_from_bin_file(file_path, debug=False, to_mmdet=False, concat=False): + print(f'Reading {file_path} ...') + pred_data = read_bin(file_path) + objects = pred_data.objects + if debug: + objects = objects[:10] + obj_dict = defaultdict(list) + print('Collecting Bboxes ...') + for o in tqdm(objects): + seg_name = o.context_name + time_stamp = o.frame_timestamp_micros + if to_mmdet: + bbox_for_vis = object2mmdetformat(o) + else: + bbox_for_vis = object2BBox(o) + obj_dict[time_stamp].append(bbox_for_vis) + + if concat: + for k in obj_dict: + obj_list = obj_dict[k] + obj_dict[k] = np.stack(obj_list, dim=0) + + return obj_dict + +idx2ts = None + +def get_pc_from_time_stamp(timestamp, path, split='training'): + global idx2ts + if idx2ts is None: + with open(path, 'rb') as fr: + idx2ts = pkl.load(fr) + print('Read idx2ts') + ts2idx = {} + for idx, ts in idx2ts.items(): + ts2idx[ts] = idx + + curr_idx = ts2idx[timestamp] + pc_root = f'./data/waymo/kitti_format/{split}/velodyne' + pc_path = os.path.join(pc_root, curr_idx + '.bin') + pc = np.fromfile(pc_path, dtype=np.float32).reshape(-1, 6) + pc = pc[:, :3] + return pc diff --git a/tools/vis/visualizer.py b/tools/vis/visualizer.py new file mode 100644 index 0000000..c171e51 --- /dev/null +++ b/tools/vis/visualizer.py @@ -0,0 +1,191 @@ +import matplotlib.pyplot as plt, numpy as np +from copy import deepcopy + + +class BBox: + def __init__(self, x=None, y=None, z=None, h=None, w=None, l=None, o=None, s=None, type=None): + self.x = x # center x + self.y = y # center y + self.z = z # center z + self.h = h # height + self.w = w # width + self.l = l # length + self.o = o # orientation + self.s = s # detection score + self.type = type + + def __str__(self): + return 'x: {}, y: {}, z: {}, heading: {}, length: {}, width: {}, height: {}, score: {}'.format( + self.x, self.y, self.z, self.o, self.l, self.w, self.h, self.s) + + @classmethod + def bbox2dict(cls, bbox): + return { + 'center_x': bbox.x, 'center_y': bbox.y, 'center_z': bbox.z, + 'height': bbox.h, 'width': bbox.w, 'length': bbox.l, 'heading': bbox.o} + + @classmethod + def bbox2array(cls, bbox): + if bbox.s is None or bbox.s == -1: + return np.array([bbox.x, bbox.y, bbox.z, bbox.o, bbox.l, bbox.w, bbox.h]) + else: + return np.array([bbox.x, bbox.y, bbox.z, bbox.o, bbox.l, bbox.w, bbox.h, bbox.s]) + + @classmethod + def array2bbox(cls, data): + bbox = BBox() + bbox.x, bbox.y, bbox.z, bbox.o, bbox.l, bbox.w, bbox.h = data[:7] + if len(data) == 8: + bbox.s = data[-1] + return bbox + + @classmethod + def dict2bbox(cls, data): + bbox = BBox() + bbox.x = data['center_x'] + bbox.y = data['center_y'] + bbox.z = data['center_z'] + bbox.h = data['height'] + bbox.w = data['width'] + bbox.l = data['length'] + bbox.o = data['heading'] + if 'score' in data.keys(): + bbox.s = data['score'] + return bbox + + @classmethod + def copy_bbox(cls, bboxa, bboxb): + bboxa.x = bboxb.x + bboxa.y = bboxb.y + bboxa.z = bboxb.z + bboxa.l = bboxb.l + bboxa.w = bboxb.w + bboxa.h = bboxb.h + bboxa.o = bboxb.o + bboxa.s = bboxb.s + return + + @classmethod + def box2corners2d(cls, bbox): + """ the coordinates for bottom corners + """ + bottom_center = np.array([bbox.x, bbox.y, bbox.z - bbox.h / 2]) + cos, sin = np.cos(bbox.o), np.sin(bbox.o) + pc0 = np.array([bbox.x + cos * bbox.l / 2 + sin * bbox.w / 2, + bbox.y + sin * bbox.l / 2 - cos * bbox.w / 2, + bbox.z - bbox.h / 2]) + pc1 = np.array([bbox.x + cos * bbox.l / 2 - sin * bbox.w / 2, + bbox.y + sin * bbox.l / 2 + cos * bbox.w / 2, + bbox.z - bbox.h / 2]) + pc2 = 2 * bottom_center - pc0 + pc3 = 2 * bottom_center - pc1 + + return [pc0.tolist(), pc1.tolist(), pc2.tolist(), pc3.tolist()] + + @classmethod + def box2corners3d(cls, bbox): + """ the coordinates for bottom corners + """ + center = np.array([bbox.x, bbox.y, bbox.z]) + bottom_corners = np.array(BBox.box2corners2d(bbox)) + up_corners = 2 * center - bottom_corners + corners = np.concatenate([up_corners, bottom_corners], axis=0) + return corners.tolist() + + @classmethod + def motion2bbox(cls, bbox, motion): + result = deepcopy(bbox) + result.x += motion[0] + result.y += motion[1] + result.z += motion[2] + result.o += motion[3] + return result + + @classmethod + def set_bbox_size(cls, bbox, size_array): + result = deepcopy(bbox) + result.l, result.w, result.h = size_array + return result + + @classmethod + def set_bbox_with_states(cls, prev_bbox, state_array): + prev_array = BBox.bbox2array(prev_bbox) + prev_array[:4] += state_array[:4] + prev_array[4:] = state_array[4:] + bbox = BBox.array2bbox(prev_array) + return bbox + + @classmethod + def box_pts2world(cls, ego_matrix, pcs): + new_pcs = np.concatenate((pcs, + np.ones(pcs.shape[0])[:, np.newaxis]), + axis=1) + new_pcs = ego_matrix @ new_pcs.T + new_pcs = new_pcs.T[:, :3] + return new_pcs + + @classmethod + def edge2yaw(cls, center, edge): + vec = edge - center + yaw = np.arccos(vec[0] / np.linalg.norm(vec)) + if vec[1] < 0: + yaw = -yaw + return yaw + + @classmethod + def bbox2world(cls, ego_matrix, box): + # center and corners + corners = np.array(BBox.box2corners2d(box)) + center = BBox.bbox2array(box)[:3][np.newaxis, :] + center = BBox.box_pts2world(ego_matrix, center)[0] + corners = BBox.box_pts2world(ego_matrix, corners) + # heading + edge_mid_point = (corners[0] + corners[1]) / 2 + yaw = BBox.edge2yaw(center[:2], edge_mid_point[:2]) + + result = deepcopy(box) + result.x, result.y, result.z = center + result.o = yaw + return result + + +class Visualizer2D: + def __init__(self, name='', figsize=(8, 8)): + self.figure = plt.figure(name, figsize=figsize) + plt.axis('equal') + self.COLOR_MAP = { + 'gray': np.array([140, 140, 136]) / 256, + 'light_blue': np.array([4, 157, 217]) / 256, + 'blue': np.array([0, 0, 255]) / 256, + 'wine_red': np.array([191, 4, 54]) / 256, + 'red': np.array([255, 0, 0]) / 256, + 'black': np.array([0, 0, 0]) / 256, + 'purple': np.array([224, 133, 250]) / 256, + 'dark_green': np.array([32, 64, 40]) / 256, + 'green': np.array([77, 115, 67]) / 256 + } + + def show(self): + plt.show() + + def close(self): + plt.close() + + def save(self, path): + plt.savefig(path) + + def handler_pc(self, pc, color='gray', s=0.25): + vis_pc = np.asarray(pc) + plt.scatter(vis_pc[:, 0], vis_pc[:, 1], s=s, marker='o', color=self.COLOR_MAP[color]) + + def handler_box(self, box: BBox, message: str='', color='red', linestyle='solid', text_color=None, fontsize='xx-small', center_message=False): + corners = np.array(BBox.box2corners2d(box))[:, :2] + corners = np.concatenate([corners, corners[0:1, :2]]) + plt.plot(corners[:, 0], corners[:, 1], color=self.COLOR_MAP[color], linestyle=linestyle) + corner_index = np.random.randint(0, 4, 1) + if text_color is None: + text_color = color + if center_message: + plt.text(box.x - 1, box.y, message, color=self.COLOR_MAP[text_color], fontsize=fontsize) + else: + plt.text(corners[corner_index, 0] - 1, corners[corner_index, 1] - 1, message, color=self.COLOR_MAP[text_color], fontsize=fontsize) \ No newline at end of file