예제 #1
0
def prep_pointcloud(
        input_dict,
        root_path,
        voxel_generator,
        target_assigner,
        db_sampler=None,
        max_voxels=20000,
        remove_outside_points=False,
        training=True,
        create_targets=True,
        shuffle_points=False,
        remove_unknown=False,
        gt_rotation_noise=(-np.pi / 3, np.pi / 3),
        gt_loc_noise_std=(1.0, 1.0, 1.0),
        global_rotation_noise=(-np.pi / 4, np.pi / 4),
        global_scaling_noise=(0.95, 1.05),
        global_random_rot_range=(0.78, 2.35),
        global_translate_noise_std=(0, 0, 0),
        num_point_features=4,
        anchor_area_threshold=1,
        gt_points_drop=0.0,
        gt_drop_max_keep=10,
        remove_points_after_sample=True,
        anchor_cache=None,
        remove_environment=False,
        random_crop=False,
        reference_detections=None,
        out_size_factor=2,
        use_group_id=False,
        multi_gpu=False,
        min_points_in_gt=-1,
        random_flip_x=True,
        random_flip_y=True,
        sample_importance=1.0,
        out_dtype=np.float32,
        bcl_keep_voxels=6500,  #6000~8000 pillar
        seg_keep_points=8000,
        points_per_voxel=200,
        feature_map_size=[1, 200, 176],
        num_anchor_per_loc=2,
        segmentation=False,
        object_detection=True):
    """convert point cloud to voxels, create targets if ground truths
    exists.

    input_dict format: dataset.get_sensor_data format

    """

    class_names = target_assigner.classes
    points = input_dict["lidar"]["points"]

    if training or segmentation:
        anno_dict = input_dict["lidar"]["annotations"]
        gt_dict = {
            "gt_boxes":
            anno_dict["boxes"],
            "gt_names":
            anno_dict["names"],
            "gt_importance":
            np.ones([anno_dict["boxes"].shape[0]],
                    dtype=anno_dict["boxes"].dtype),
        }

        if "difficulty" not in anno_dict:
            difficulty = np.zeros([anno_dict["boxes"].shape[0]],
                                  dtype=np.int32)
            gt_dict["difficulty"] = difficulty
        else:
            gt_dict["difficulty"] = anno_dict["difficulty"]

        if use_group_id and "group_ids" in anno_dict:
            group_ids = anno_dict["group_ids"]
            gt_dict["group_ids"] = group_ids

    calib = None
    if "calib" in input_dict:
        calib = input_dict["calib"]

    if reference_detections is not None:
        assert calib is not None and "image" in input_dict
        C, R, T = box_np_ops.projection_matrix_to_CRT_kitti(P2)
        frustums = box_np_ops.get_frustum_v2(reference_detections, C)
        frustums -= T
        frustums = np.einsum('ij, akj->aki', np.linalg.inv(R), frustums)
        frustums = box_np_ops.camera_to_lidar(frustums, rect, Trv2c)
        surfaces = box_np_ops.corner_to_surfaces_3d_jit(frustums)
        masks = points_in_convex_polygon_3d_jit(points, surfaces)
        points = points[masks.any(-1)]
    if remove_outside_points:
        assert calib is not None
        image_shape = input_dict["image"]["image_shape"]
        points = box_np_ops.remove_outside_points(points, calib["rect"],
                                                  calib["Trv2c"], calib["P2"],
                                                  image_shape)
    if remove_environment is True and training:
        selected = kitti.keep_arrays_by_name(gt_names, target_assigner.classes)
        _dict_select(gt_dict, selected)
        masks = box_np_ops.points_in_rbbox(points, gt_dict["gt_boxes"])
        points = points[masks.any(-1)]

    if training:
        boxes_lidar = gt_dict["gt_boxes"]
        selected = kitti.drop_arrays_by_name(gt_dict["gt_names"], ["DontCare"])
        _dict_select(gt_dict, selected)
        if remove_unknown:
            remove_mask = gt_dict["difficulty"] == -1
            """
            gt_boxes_remove = gt_boxes[remove_mask]
            gt_boxes_remove[:, 3:6] += 0.25
            points = prep.remove_points_in_boxes(points, gt_boxes_remove)
            """
            keep_mask = np.logical_not(remove_mask)
            _dict_select(gt_dict, keep_mask)
        gt_dict.pop("difficulty")
        if min_points_in_gt > 0:
            # points_count_rbbox takes 10ms with 10 sweeps nuscenes data
            point_counts = box_np_ops.points_count_rbbox(
                points, gt_dict["gt_boxes"])
            mask = point_counts >= min_points_in_gt
            _dict_select(gt_dict, mask)
        gt_boxes_mask = np.array(
            [n in class_names for n in gt_dict["gt_names"]], dtype=np.bool_)
        if db_sampler is not None:
            group_ids = None
            if "group_ids" in gt_dict:
                group_ids = gt_dict["group_ids"]

            sampled_dict = db_sampler.sample_all(root_path,
                                                 gt_dict["gt_boxes"],
                                                 gt_dict["gt_names"],
                                                 num_point_features,
                                                 random_crop,
                                                 gt_group_ids=group_ids,
                                                 calib=calib)

            if sampled_dict is not None:
                sampled_gt_names = sampled_dict["gt_names"]
                sampled_gt_boxes = sampled_dict["gt_boxes"]
                sampled_points = sampled_dict["points"]
                sampled_gt_masks = sampled_dict["gt_masks"]
                gt_dict["gt_names"] = np.concatenate(
                    [gt_dict["gt_names"], sampled_gt_names], axis=0)
                gt_dict["gt_boxes"] = np.concatenate(
                    [gt_dict["gt_boxes"], sampled_gt_boxes])
                gt_boxes_mask = np.concatenate(
                    [gt_boxes_mask, sampled_gt_masks], axis=0)
                sampled_gt_importance = np.full([sampled_gt_boxes.shape[0]],
                                                sample_importance,
                                                dtype=sampled_gt_boxes.dtype)
                gt_dict["gt_importance"] = np.concatenate(
                    [gt_dict["gt_importance"], sampled_gt_importance])

                if group_ids is not None:
                    sampled_group_ids = sampled_dict["group_ids"]
                    gt_dict["group_ids"] = np.concatenate(
                        [gt_dict["group_ids"], sampled_group_ids])

                if remove_points_after_sample:
                    masks = box_np_ops.points_in_rbbox(points,
                                                       sampled_gt_boxes)
                    points = points[np.logical_not(masks.any(-1))]

                points = np.concatenate([sampled_points, points], axis=0)

        pc_range = voxel_generator.point_cloud_range
        group_ids = None
        if "group_ids" in gt_dict:
            group_ids = gt_dict["group_ids"]

        prep.noise_per_object_v3_(
            gt_dict["gt_boxes"],
            points,
            gt_boxes_mask,
            rotation_perturb=gt_rotation_noise,
            center_noise_std=gt_loc_noise_std,
            global_random_rot_range=global_random_rot_range,
            group_ids=group_ids,
            num_try=100)

        # should remove unrelated objects after noise per object
        # for k, v in gt_dict.items():
        #     print(k, v.shape)
        _dict_select(gt_dict, gt_boxes_mask)
        gt_classes = np.array(
            [class_names.index(n) + 1 for n in gt_dict["gt_names"]],
            dtype=np.int32)
        gt_dict["gt_classes"] = gt_classes
        gt_dict["gt_boxes"], points = prep.random_flip(gt_dict["gt_boxes"],
                                                       points, 0.5,
                                                       random_flip_x,
                                                       random_flip_y)
        gt_dict["gt_boxes"], points = prep.global_rotation_v2(
            gt_dict["gt_boxes"], points, *global_rotation_noise)
        gt_dict["gt_boxes"], points = prep.global_scaling_v2(
            gt_dict["gt_boxes"], points, *global_scaling_noise)
        prep.global_translate_(gt_dict["gt_boxes"], points,
                               global_translate_noise_std)

        bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]]
        mask = prep.filter_gt_box_outside_range_by_center(
            gt_dict["gt_boxes"], bv_range)
        _dict_select(gt_dict, mask)

        # limit rad to [-pi, pi]
        gt_dict["gt_boxes"][:, 6] = box_np_ops.limit_period(
            gt_dict["gt_boxes"][:, 6], offset=0.5, period=2 * np.pi)

    # add depth for point feature and remove intensity
    # points = points[...,:3]
    # points = AddDepthFeature(points, num_point_features)
    # num_point_features = points.shape[-1] #update point shape

    #remove points out of PC rannge
    pc_range = voxel_generator.point_cloud_range  # [0, -40, -3, 70.4, 40, 1] xmin,ymin.zmin. xmax. ymax, zmax
    points = box_np_ops.remove_out_pc_range_points(points, pc_range)

    if shuffle_points and not segmentation:
        np.random.shuffle(points)  # shuffle is a little slow.

    if not training and segmentation:
        #Keep Car Only
        gt_boxes_mask = np.array(
            [n in class_names for n in gt_dict["gt_names"]], dtype=np.bool_)
        _dict_select(gt_dict, gt_boxes_mask)

        points_in_box, points_out_box = box_np_ops.split_points_in_boxes(
            points, gt_dict["gt_boxes"])  #xyzr
        points_in_box, points_out_box = SamplePointsKeepALLPositive(
            points_in_box, points_out_box, seg_keep_points,
            num_point_features)  #fixed points
        data, label = PrepDataAndLabel(points_in_box, points_out_box)

        example = {
            'seg_points': data,  #data
            'seg_labels': label,  #label
            'gt_boxes': gt_dict["gt_boxes"],
            'image_idx': input_dict['metadata']['image_idx'],
        }

        ################# For feature map Focs
        # # # NOTE: For feature map Focs
        point_cloud_range = np.array(voxel_generator.point_cloud_range)
        anchor_strides = (point_cloud_range[3:] -
                          point_cloud_range[:3]) / feature_map_size[::-1]
        anchor_offsets = point_cloud_range[:3] + anchor_strides / 2
        centers = box_np_ops.create_anchors_3d_stride(
            feature_map_size,
            anchor_strides=anchor_strides,
            anchor_offsets=anchor_offsets,
            rotations=[0])
        centers = centers.squeeze()[..., :3].reshape(-1, 3)
        example.update({
            'coords_center':
            centers,  # if anchors free the 0 is the horizontal/vertical anchors
        })
        ##############

        ################ Fcos & points to voxel Test
        # NOTE: For voxel seg net
        # _, coords, coords_center, p2voxel_idx = box_np_ops.points_to_3dvoxel(data,
        #                                         feat_size=[100,80,10],
        #                                         max_voxels=bcl_keep_voxels,
        #                                         num_p_voxel=points_per_voxel)
        # example = {
        # 'seg_points': data, #data
        # 'coords': coords,
        # 'coords_center': coords_center,
        # 'p2voxel_idx': p2voxel_idx,
        # 'gt_boxes' : gt_dict["gt_boxes"],
        # 'image_idx' : input_dict['metadata']['image_idx'],
        # "gt_num" :  len(gt_dict["gt_boxes"]),
        # 'gt_boxes' : gt_dict["gt_boxes"],
        # 'seg_labels': label
        # }
        ################ Fcos & points to voxel

        if anchor_cache is not None:

            example.update({
                "gt_num":
                len(gt_dict["gt_boxes"]),  #how many objects in eval GT
                "anchors": anchor_cache["anchors"]
            })

        return example

    ################################Car point segmentation#####################
    if training and segmentation:
        # points_in_box = box_np_ops.points_in_rbbox(points, gt_dict["gt_boxes"]) #xyzr
        # enlarge bouding box
        # enlarge_size = 0.2
        # gt_dict["gt_boxes"][:, 3:6] = gt_dict["gt_boxes"][:, 3:6] + enlarge_size #xyzhwlr
        # masks = box_np_ops.points_in_rbbox(points, gt_dict["gt_boxes"])
        # points = points[np.logical_not(masks.any(-1))]
        # points = np.concatenate((points, points_in_box), axis=0)

        #above and below bouding box should have no points
        # gt_dict["gt_boxes"][:,3] += 2

        #random sample
        # points = SamplePoints(points, seg_keep_points, num_point_features) #Sample zero
        # points = PointRandomChoice(points, seg_keep_points) #Repeat sample
        points = PointRandomChoiceV2(
            points, seg_keep_points)  #Repeat sample according points distance
        points_in_box, points_out_box = box_np_ops.split_points_in_boxes(
            points, gt_dict["gt_boxes"])  #xyzr
        data, label = PrepDataAndLabel(points_in_box, points_out_box)

        #keep positive sample
        # points_in_box, points_out_box = box_np_ops.split_points_in_boxes(points, gt_dict["gt_boxes"]) #xyzr
        # points_in_box, points_out_box = SamplePointsKeepALLPositive(points_in_box, points_out_box, seg_keep_points, num_point_features) #fixed 18888 points
        # data, label = PrepDataAndLabel(points_in_box, points_out_box)
        """shuffle car seg points"""
        indices = np.arange(data.shape[0])
        np.random.shuffle(indices)
        data = data[indices]
        label = label[indices]

        example = {
            'seg_points': data,  #data
            'seg_labels': label,  #label
            'gt_boxes': gt_dict["gt_boxes"],
        }

        ################# For feature map Focs
        # # NOTE: For feature map Focs
        point_cloud_range = np.array(voxel_generator.point_cloud_range)
        anchor_strides = (point_cloud_range[3:] -
                          point_cloud_range[:3]) / feature_map_size[::-1]
        anchor_offsets = point_cloud_range[:3] + anchor_strides / 2
        centers = box_np_ops.create_anchors_3d_stride(
            feature_map_size,
            anchor_strides=anchor_strides,
            anchor_offsets=anchor_offsets,
            rotations=[0])
        centers = centers.squeeze()[..., :3].reshape(-1, 3)
        targets_dict = box_np_ops.fcos_box_encoder_v2(centers,
                                                      gt_dict["gt_boxes"])
        # bbox = box_np_ops.fcos_box_decoder_v2(np.expand_dims(centers, 0),
        #                                 np.expand_dims(targets_dict["bbox_targets"], 0))
        # labels = targets_dict["labels"]
        # with open(os.path.join('./debug_tool',"points.pkl") , 'wb') as f:
        #     pickle.dump(data,f)
        # with open(os.path.join('./debug_tool',"seg_points.pkl") , 'wb') as f:
        #     pickle.dump(centers[labels==1],f)
        # with open(os.path.join('./debug_tool',"pd_boxes.pkl") , 'wb') as f:
        #     pickle.dump(bbox.squeeze()[labels==1],f)
        # with open(os.path.join('./debug_tool',"gt_boxes.pkl") , 'wb') as f:
        #     pickle.dump(gt_dict["gt_boxes"],f)
        # exit()
        example.update({
            'labels': targets_dict[
                'labels'],  # if anchors free the 0 is the horizontal/vertical anchors
            # 'seg_labels': targets_dict['labels'], # if anchors free the 0 is the horizontal/vertical anchors
            'reg_targets':
            targets_dict['bbox_targets'],  # target assign get offsite
            'importance': targets_dict['importance'],
            # 'reg_weights': targets_dict['bbox_outside_weights'],
        })
        ##############

        ################ Fcos & points to voxel
        # NOTE: For voxel seg net
        # _, coords, coords_center, p2voxel_idx = box_np_ops.points_to_3dvoxel(data,
        #                                         feat_size=[200,176,10],
        #                                         max_voxels=bcl_keep_voxels,
        #                                         num_p_voxel=points_per_voxel)
        #
        # targets_dict = box_np_ops.fcos_box_encoder_v2(coords_center,
        #                                     gt_dict["gt_boxes"])
        # # Jim added
        # example.update({
        # 'coords': coords,
        # 'p2voxel_idx': p2voxel_idx,
        # 'cls_labels': targets_dict['labels'], # if anchors free the 0 is the horizontal/vertical anchors
        # 'reg_targets': targets_dict['bbox_targets'], # target assign get offsite
        # 'importance': targets_dict['importance'],
        # })
        ################ Fcos & points to voxel

        ################ Fcos & points to voxel
        if anchor_cache is not None:
            anchors = anchor_cache["anchors"]
            anchors_bv = anchor_cache["anchors_bv"]
            anchors_dict = anchor_cache["anchors_dict"]
            matched_thresholds = anchor_cache["matched_thresholds"]
            unmatched_thresholds = anchor_cache["unmatched_thresholds"]

            targets_dict = target_assigner.assign(
                anchors,
                anchors_dict,  #this is the key to control the number of anchors (input anchors) ['anchors, unmatch,match']
                gt_dict["gt_boxes"],
                anchors_mask=None,
                gt_classes=gt_dict["gt_classes"],
                gt_names=gt_dict["gt_names"],
                matched_thresholds=matched_thresholds,
                unmatched_thresholds=unmatched_thresholds,
                importance=gt_dict["gt_importance"])

            example.update({
                'labels': targets_dict[
                    'labels'],  # if anchors free the 0 is the horizontal/vertical anchors
                'reg_targets':
                targets_dict['bbox_targets'],  # target assign get offsite
                #'importance': targets_dict['importance'],
            })

            # boxes_lidar = gt_dict["gt_boxes"]
            # bev_map = simplevis.kitti_vis(points, boxes_lidar, gt_dict["gt_names"])
            # assigned_anchors = anchors[targets_dict['labels'] > 0]
            # ignored_anchors = anchors[targets_dict['labels'] == -1]
            # bev_map = simplevis.draw_box_in_bev(bev_map, [0, -40, -3, 70.4, 40, 1], ignored_anchors, [128, 128, 128], 2)
            # bev_map = simplevis.draw_box_in_bev(bev_map, [0, -40, -3, 70.4, 40, 1], assigned_anchors, [255, 0, 0])
            # cv2.imwrite('./visualization/anchors/anchors_{}.png'.format(input_dict['metadata']['image_idx']),bev_map)

        return example

    #################################voxel_generator############################
    '''
    voxel_size = voxel_generator.voxel_size # [0, -40, -3, 70.4, 40, 1]
    pc_range = voxel_generator.point_cloud_range
    grid_size = voxel_generator.grid_size # [352, 400]
    max_num_points_per_voxel = voxel_generator.max_num_points_per_voxel

    if not multi_gpu:
        res = voxel_generator.generate(
            points, max_voxels)
        voxels = res["voxels"]
        coordinates = res["coordinates"]
        num_points = res["num_points_per_voxel"]
        num_voxels = np.array([voxels.shape[0]], dtype=np.int64)
    else:
        res = voxel_generator.generate_multi_gpu(
            points, max_voxels)
        voxels = res["voxels"]
        coordinates = res["coordinates"]
        num_points = res["num_points_per_voxel"]
        num_voxels = np.array([res["voxel_num"]], dtype=np.int64)

    example = {
        'voxels': voxels,
        #'num_points': num_points,
        'coordinates': coordinates,
        "num_voxels": num_voxels,
    }


    ## WARNING:  For Simplex voxel Testing if bug comment this
    voxels= SimpleVoxel(voxels, num_points) #(V,100,C) -> (B, C, V, N) #For Second, if Pillar comment it
    max_num_points_per_voxel=1 #If SimpleVoxel max_num_points_per_voxel=1

    voxels, coordinates = VoxelRandomChoice(voxels, coordinates,
                                        bcl_keep_voxels, num_point_features,
                                        max_num_points_per_voexl=max_num_points_per_voxel)

    example['voxels']=voxels
    example['coordinates']=coordinates
    '''
    ############################################################################

    # if calib is not None:
    #     example["calib"] = calib

    if anchor_cache is not None:
        anchors = anchor_cache["anchors"]
        anchors_bv = anchor_cache["anchors_bv"]
        anchors_dict = anchor_cache["anchors_dict"]
        matched_thresholds = anchor_cache["matched_thresholds"]
        unmatched_thresholds = anchor_cache["unmatched_thresholds"]

    else:
        # generate anchors from ground truth
        """
        voxels= SimpleVoxel(voxels, num_points) #(V,100,C) -> (B, C, V, N)
        voxels, coordinates, num_points = VoxelRandomChoice(voxels, coordinates, num_points, bcl_keep_voxels)
        example['voxels']=voxels
        example['num_points']=num_points
        example['coordinates']=coordinates
        example['num_voxels']=bcl_keep_voxels

        if training:
            # for anchor free
            gt_boxes_coords = gt_dict["gt_boxes"][:,:3] #original gt xyz
            example['gt_boxes_coords']=gt_boxes_coords #GT save to example
            gt_boxes_coords = np.round(gt_dict["gt_boxes"][:,:3]).astype(int) #round xyz
            gt_boxes_coords = gt_boxes_coords[:,::-1] #zyx reverse
            ret = target_assigner.generate_anchors_from_gt(gt_boxes_coords) #for GT generate anchors
            anchors = ret["anchors"]
            anchors_dict = target_assigner.generate_anchors_dict_from_gt(gt_boxes_coords) #for GT generate anchors

        if not training:
            # for anchor free
            feature_map_size = grid_size[:2] // out_size_factor
            feature_map_size = [*feature_map_size, 1][::-1]
            ret = target_assigner.generate_anchors(feature_map_size)
            anchors_dict = target_assigner.generate_anchors_dict(feature_map_size)
            anchors = ret["anchors"]
        """

        # # generate anchors from anchor free (Voxel-wise)
        # ret = target_assigner.generate_anchors_from_voxels(coordinates) #for coordinates generate anchors
        # anchors_dict = target_assigner.generate_anchors_dict_from_voxels(coordinates) #this is the key to control the number of anchors (input anchors)
        # anchors = ret["anchors"]
        # matched_thresholds = ret["matched_thresholds"]
        # unmatched_thresholds = ret["unmatched_thresholds"]

        # generate anchors from  voxel + anchor free
        """
        gt_boxes_coords = gt_dict["gt_boxes"][:,:3] #original gt xyz
        #gt_boxes_coords = np.round(gt_dict["gt_boxes"][:,:3]).astype(int) #round xyz
        gt_boxes_coords = gt_boxes_coords[:,::-1] #zyx reverse

        #stack ret and ret_gt
        ret = target_assigner.generate_anchors_from_voxels(coordinates) #for coordinates generate anchors
        ret_gt = target_assigner.generate_anchors_from_gt(gt_boxes_coords) #for GT generate anchors
        for k in ret.keys():
            ret[k] = np.concatenate((ret[k], ret_gt[k]))
        anchors = ret["anchors"]

        #stack anchors_dict and anchors_dict_gt
        anchors_dict = target_assigner.generate_anchors_dict_from_voxels(coordinates) #this is the key to control the number of anchors (input anchors) ['anchors, unmatch,match']
        anchors_dict_gt = target_assigner.generate_anchors_dict_from_gt(gt_boxes_coords) #for GT generate anchors

        for order_k in anchors_dict.keys():
            for k in anchors_dict[order_k].keys():
                anchors_dict[order_k][k] = np.concatenate((anchors_dict[order_k][k], anchors_dict_gt[order_k][k]))
        """

        # generate anchors from groundtruth
        """
        if training:
            # generate anchors from car points
            points_in_box = points_in_box[:,:3] #xyz
            points_in_box = points_in_box[:,::-1] #zyx
            ret = target_assigner.generate_anchors_from_gt(points_in_box) #for GT generate anchors
            anchors = ret["anchors"]
            anchors_dict = target_assigner.generate_anchors_dict_from_gt(points_in_box) #for GT generate anchors


            anchors_bv = box_np_ops.rbbox2d_to_near_bbox(
                anchors[:, [0, 1, 3, 4, 6]])
            matched_thresholds = ret["matched_thresholds"]
            unmatched_thresholds = ret["unmatched_thresholds"]
        """

    # Fcos points sampling
    points = SamplePoints(points, bcl_keep_voxels, num_point_features)
    example = {
        'voxels': np.expand_dims(points, 0),
        #'num_points': num_points,
        'coordinates': points,
        # "num_voxels": None,
    }
    if not training:
        anno_dict = input_dict["lidar"]["annotations"]
        gt_dict = {
            "gt_boxes":
            anno_dict["boxes"],
            "gt_names":
            anno_dict["names"],
            "gt_importance":
            np.ones([anno_dict["boxes"].shape[0]],
                    dtype=anno_dict["boxes"].dtype),
        }
    targets_dict = box_np_ops.fcos_box_encoder_v2(points, gt_dict["gt_boxes"])
    # targets_dict = box_np_ops.fcos_box_encoder(points, gt_dict["gt_boxes"])
    example.update({
        'labels': targets_dict[
            'labels'],  # if anchors free the 0 is the horizontal/vertical anchors
        'seg_labels': targets_dict[
            'labels'],  # if anchors free the 0 is the horizontal/vertical anchors
        'reg_targets':
        targets_dict['bbox_targets'],  # target assign get offsite
        'importance': targets_dict['importance'],
        # 'reg_weights': targets_dict['bbox_outside_weights'],
    })
    # example["anchors"] = anchors
    # anchors_mask = None
    # if anchor_area_threshold >= 0:
    #     # slow with high resolution. recommend disable this forever.
    #     coors = coordinates
    #     dense_voxel_map = box_np_ops.sparse_sum_for_anchors_mask(
    #         coors, tuple(grid_size[::-1][1:]))
    #     dense_voxel_map = dense_voxel_map.cumsum(0)
    #     dense_voxel_map = dense_voxel_map.cumsum(1)
    #     anchors_area = box_np_ops.fused_get_anchors_area(
    #         dense_voxel_map, anchors_bv, voxel_size, pc_range, grid_size)
    #     anchors_mask = anchors_area > anchor_area_threshold
    #     # example['anchors_mask'] = anchors_mask.astype(np.uint8)
    #     example['anchors_mask'] = anchors_mask

    if not training:
        # Use it when debuging eval nms for good
        eval_classes = input_dict["lidar"]["annotations"]["names"]
        eval_gt_dict = {"gt_names": eval_classes}
        gt_boxes_mask = np.array([n in class_names for n in eval_classes],
                                 dtype=np.bool_)
        _dict_select(eval_gt_dict, gt_boxes_mask)
        example["gt_num"] = len(
            eval_gt_dict["gt_names"])  #how many objects in eval GT

        return example

    example["gt_names"] = gt_dict["gt_names"]
    # voxel_labels = box_np_ops.assign_label_to_voxel(gt_boxes, coordinates,
    #                                                 voxel_size, coors_range)
    """
    # bev anchors without screening
    boxes_lidar = gt_dict["gt_boxes"]
    bev_map = simplevis.kitti_vis(points, boxes_lidar, gt_dict["gt_names"])
    bev_map = simplevis.draw_box_in_bev(bev_map, [0, -40, -3, 70.4, 40, 1], anchors, [255, 0, 0]) #assigned_anchors blue
    cv2.imwrite('anchors/anchors_{}.png'.format(input_dict['metadata']['image_idx']),bev_map)
    # cv2.imshow('anchors', bev_map)
    # cv2.waitKey(0)
    """

    if create_targets:
        # No particular use
        where = None
        # Fcos target generator and encoder
        # targets_dict = target_assigner.assign(
        #     anchors,
        #     anchors_dict, #this is the key to control the number of anchors (input anchors) ['anchors, unmatch,match']
        #     gt_dict["gt_boxes"],
        #     anchors_mask,
        #     gt_classes=gt_dict["gt_classes"],
        #     gt_names=gt_dict["gt_names"],
        #     matched_thresholds=matched_thresholds,
        #     unmatched_thresholds=unmatched_thresholds,
        #     importance=gt_dict["gt_importance"])
        ################################Visualaiziton###########################
        """
        bev anchors with points
        boxes_lidar = gt_dict["gt_boxes"]
        bev_map = simplevis.kitti_vis(points, boxes_lidar, gt_dict["gt_names"])
        assigned_anchors = anchors[targets_dict['labels'] > 0]
        ignored_anchors = anchors[targets_dict['labels'] == -1]
        bev_map = simplevis.draw_box_in_bev(bev_map, [0, -40, -3, 70.4, 40, 1], ignored_anchors, [128, 128, 128], 2) #ignored_anchors gray    #[0, -30, -3, 64, 30, 1] for kitti
        bev_map = simplevis.draw_box_in_bev(bev_map, [0, -40, -3, 70.4, 40, 1], assigned_anchors, [255, 0, 0]) #assigned_anchors blue
        cv2.imwrite('anchors/anchors_{}.png'.format(input_dict['metadata']['image_idx']),bev_map)
        cv2.imshow('anchors', bev_map)
        cv2.waitKey(0)
        """
        """
        # bev boxes_lidar with voxels (put z in to the plane)
        boxes_lidar = gt_dict["gt_boxes"]
        pp_map = np.zeros(grid_size[:2], dtype=np.float32) # (1408, 1600)
        #print(voxels.shape)  #(16162, 5, 4) $ 4=bzyx
        voxels_max = np.max(voxels[:, :, 1], axis=1, keepdims=False)
        voxels_min = np.min(voxels[:, :, 1], axis=1, keepdims=False)
        voxels_height = voxels_max - voxels_min
        voxels_height = np.minimum(voxels_height, 4) #keep every voxels length less than 4

        # sns.distplot(voxels_height)
        # plt.show()
        pp_map[coordinates[:, 2], coordinates[:, 1]] = voxels_height / 4 #coordinates bzyx
        pp_map = (pp_map * 255).astype(np.uint8)
        pp_map = cv2.cvtColor(pp_map, cv2.COLOR_GRAY2RGB)
        pp_map = simplevis.draw_box_in_bev(pp_map, [0, -30, -3, 64, 30, 1], boxes_lidar, [128, 0, 128], 2) # for kitti 0, -30, -3, 64, 30, 1
        cv2.imwrite('bev_pp_map/pp_map{}.png'.format(input_dict['metadata']['image_idx']),pp_map)
        # cv2.imshow('heights', pp_map)
        # cv2.waitKey(0)
        """

        # example.update({
        #     'labels': targets_dict['labels'], # if anchors free the 0 is the horizontal/vertical anchors
        #     'reg_targets': targets_dict['bbox_targets'], # target assign get offsite
        #     'importance': targets_dict['importance'],
        #     # 'reg_weights': targets_dict['bbox_outside_weights'],
        # })

    return example
예제 #2
0
def prep_pointcloud(input_dict,
                    root_path,
                    voxel_generator,
                    target_assigner,
                    db_sampler=None,
                    max_voxels=20000,
                    remove_outside_points=False,
                    training=True,
                    create_targets=True,
                    shuffle_points=False,
                    remove_unknown=False,
                    gt_rotation_noise=(-np.pi / 3, np.pi / 3),
                    gt_loc_noise_std=(1.0, 1.0, 1.0),
                    global_rotation_noise=(-np.pi / 4, np.pi / 4),
                    global_scaling_noise=(0.95, 1.05),
                    global_random_rot_range=(0.78, 2.35),
                    global_translate_noise_std=(0, 0, 0),
                    num_point_features=4,
                    anchor_area_threshold=1,
                    gt_points_drop=0.0,
                    gt_drop_max_keep=10,
                    remove_points_after_sample=True,
                    anchor_cache=None,
                    remove_environment=False,
                    random_crop=False,
                    reference_detections=None,
                    out_size_factor=2,
                    use_group_id=False,
                    multi_gpu=False,
                    min_points_in_gt=-1,
                    random_flip_x=True,
                    random_flip_y=True,
                    sample_importance=1.0,
                    out_dtype=np.float32):
    """convert point cloud to voxels, create targets if ground truths 
    exists.

    input_dict format: dataset.get_sensor_data format

    """
    t = time.time()
    class_names = target_assigner.classes
    points = input_dict["lidar"]["points"]
    if training:
        anno_dict = input_dict["lidar"]["annotations"]
        gt_dict = {
            "gt_boxes": anno_dict["boxes"],
            "gt_names": anno_dict["names"],
            "gt_importance": np.ones([anno_dict["boxes"].shape[0]], dtype=anno_dict["boxes"].dtype),
        }
        if "difficulty" not in anno_dict:
            difficulty = np.zeros([anno_dict["boxes"].shape[0]],
                                  dtype=np.int32)
            gt_dict["difficulty"] = difficulty
        else:
            gt_dict["difficulty"] = anno_dict["difficulty"]
        if use_group_id and "group_ids" in anno_dict:
            group_ids = anno_dict["group_ids"]
            gt_dict["group_ids"] = group_ids
    calib = None
    if "calib" in input_dict:
        calib = input_dict["calib"]

    if reference_detections is not None:
        assert calib is not None and "image" in input_dict
        C, R, T = box_np_ops.projection_matrix_to_CRT_kitti(P2)
        frustums = box_np_ops.get_frustum_v2(reference_detections, C)
        frustums -= T
        frustums = np.einsum('ij, akj->aki', np.linalg.inv(R), frustums)
        frustums = box_np_ops.camera_to_lidar(frustums, rect, Trv2c)
        surfaces = box_np_ops.corner_to_surfaces_3d_jit(frustums)
        masks = points_in_convex_polygon_3d_jit(points, surfaces)
        points = points[masks.any(-1)]

    if remove_outside_points:
        assert calib is not None
        image_shape = input_dict["image"]["image_shape"]
        points = box_np_ops.remove_outside_points(
            points, calib["rect"], calib["Trv2c"], calib["P2"], image_shape)
    if remove_environment is True and training:
        selected = kitti.keep_arrays_by_name(gt_names, target_assigner.classes)
        _dict_select(gt_dict, selected)
        masks = box_np_ops.points_in_rbbox(points, gt_dict["gt_boxes"])
        points = points[masks.any(-1)]
    metrics = {}

    if training:
        """
        boxes_lidar = gt_dict["gt_boxes"]
        bev_map = simplevis.nuscene_vis(points, boxes_lidar)
        cv2.imshow('pre-noise', bev_map)
        """
        selected = kitti.drop_arrays_by_name(gt_dict["gt_names"], ["DontCare"])
        _dict_select(gt_dict, selected)
        if remove_unknown:
            remove_mask = gt_dict["difficulty"] == -1
            """
            gt_boxes_remove = gt_boxes[remove_mask]
            gt_boxes_remove[:, 3:6] += 0.25
            points = prep.remove_points_in_boxes(points, gt_boxes_remove)
            """
            keep_mask = np.logical_not(remove_mask)
            _dict_select(gt_dict, keep_mask)
        gt_dict.pop("difficulty")
        if min_points_in_gt > 0:
            # points_count_rbbox takes 10ms with 10 sweeps nuscenes data
            point_counts = box_np_ops.points_count_rbbox(points, gt_dict["gt_boxes"])
            mask = point_counts >= min_points_in_gt
            _dict_select(gt_dict, mask)
        gt_boxes_mask = np.array(
            [n in class_names for n in gt_dict["gt_names"]], dtype=np.bool_)
        if db_sampler is not None:
            group_ids = None
            if "group_ids" in gt_dict:
                group_ids = gt_dict["group_ids"]

            sampled_dict = db_sampler.sample_all(
                root_path,
                gt_dict["gt_boxes"],
                gt_dict["gt_names"],
                num_point_features,
                random_crop,
                gt_group_ids=group_ids,
                calib=calib)

            if sampled_dict is not None:
                sampled_gt_names = sampled_dict["gt_names"]
                sampled_gt_boxes = sampled_dict["gt_boxes"]
                sampled_points = sampled_dict["points"]
                sampled_gt_masks = sampled_dict["gt_masks"]
                gt_dict["gt_names"] = np.concatenate(
                    [gt_dict["gt_names"], sampled_gt_names], axis=0)
                gt_dict["gt_boxes"] = np.concatenate(
                    [gt_dict["gt_boxes"], sampled_gt_boxes])
                gt_boxes_mask = np.concatenate(
                    [gt_boxes_mask, sampled_gt_masks], axis=0)
                sampled_gt_importance = np.full([sampled_gt_boxes.shape[0]], sample_importance,
                                                dtype=sampled_gt_boxes.dtype)
                gt_dict["gt_importance"] = np.concatenate(
                    [gt_dict["gt_importance"], sampled_gt_importance])

                if group_ids is not None:
                    sampled_group_ids = sampled_dict["group_ids"]
                    gt_dict["group_ids"] = np.concatenate(
                        [gt_dict["group_ids"], sampled_group_ids])

                if remove_points_after_sample:
                    masks = box_np_ops.points_in_rbbox(points,
                                                       sampled_gt_boxes)
                    points = points[np.logical_not(masks.any(-1))]

                points = np.concatenate([sampled_points, points], axis=0)
        pc_range = voxel_generator.point_cloud_range
        group_ids = None
        if "group_ids" in gt_dict:
            group_ids = gt_dict["group_ids"]

        prep.noise_per_object_v3_(
            gt_dict["gt_boxes"],
            points,
            gt_boxes_mask,
            rotation_perturb=gt_rotation_noise,
            center_noise_std=gt_loc_noise_std,
            global_random_rot_range=global_random_rot_range,
            group_ids=group_ids,
            num_try=100)

        # should remove unrelated objects after noise per object
        # for k, v in gt_dict.items():
        #     print(k, v.shape)
        _dict_select(gt_dict, gt_boxes_mask)
        gt_classes = np.array(
            [class_names.index(n) + 1 for n in gt_dict["gt_names"]],
            dtype=np.int32)
        gt_dict["gt_classes"] = gt_classes
        gt_dict["gt_boxes"], points = prep.random_flip(gt_dict["gt_boxes"],
                                                       points, 0.5, random_flip_x, random_flip_y)
        gt_dict["gt_boxes"], points = prep.global_rotation_v2(
            gt_dict["gt_boxes"], points, *global_rotation_noise)
        gt_dict["gt_boxes"], points = prep.global_scaling_v2(
            gt_dict["gt_boxes"], points, *global_scaling_noise)
        prep.global_translate_(gt_dict["gt_boxes"], points, global_translate_noise_std)
        bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]]
        mask = prep.filter_gt_box_outside_range_by_center(gt_dict["gt_boxes"], bv_range)
        _dict_select(gt_dict, mask)

        # limit rad to [-pi, pi]
        gt_dict["gt_boxes"][:, 6] = box_np_ops.limit_period(
            gt_dict["gt_boxes"][:, 6], offset=0.5, period=2 * np.pi)

        # boxes_lidar = gt_dict["gt_boxes"]
        # bev_map = simplevis.nuscene_vis(points, boxes_lidar)
        # cv2.imshow('post-noise', bev_map)
        # cv2.waitKey(0)
    if shuffle_points:
        # shuffle is a little slow.
        np.random.shuffle(points)

    # [0, -40, -3, 70.4, 40, 1]
    voxel_size = voxel_generator.voxel_size
    pc_range = voxel_generator.point_cloud_range
    grid_size = voxel_generator.grid_size
    # [352, 400]
    t1 = time.time()
    if not multi_gpu:
        res = voxel_generator.generate(
            points, max_voxels)
        voxels = res["voxels"]
        coordinates = res["coordinates"]
        num_points = res["num_points_per_voxel"]
        num_voxels = np.array([voxels.shape[0]], dtype=np.int64)
    else:
        res = voxel_generator.generate_multi_gpu(
            points, max_voxels)
        voxels = res["voxels"]
        coordinates = res["coordinates"]
        num_points = res["num_points_per_voxel"]
        num_voxels = np.array([res["voxel_num"]], dtype=np.int64)
    metrics["voxel_gene_time"] = time.time() - t1
    example = {
        'voxels': voxels,
        'num_points': num_points,
        'coordinates': coordinates,
        "num_voxels": num_voxels,
        "metrics": metrics,
    }
    if calib is not None:
        example["calib"] = calib
    feature_map_size = grid_size[:2] // out_size_factor
    feature_map_size = [*feature_map_size, 1][::-1]
    if anchor_cache is not None:
        anchors = anchor_cache["anchors"]
        anchors_bv = anchor_cache["anchors_bv"]
        anchors_dict = anchor_cache["anchors_dict"]
        matched_thresholds = anchor_cache["matched_thresholds"]
        unmatched_thresholds = anchor_cache["unmatched_thresholds"]

    else:
        ret = target_assigner.generate_anchors(feature_map_size)
        anchors = ret["anchors"]
        anchors = anchors.reshape([-1, target_assigner.box_ndim])
        anchors_dict = target_assigner.generate_anchors_dict(feature_map_size)
        anchors_bv = box_np_ops.rbbox2d_to_near_bbox(
            anchors[:, [0, 1, 3, 4, 6]])
        matched_thresholds = ret["matched_thresholds"]
        unmatched_thresholds = ret["unmatched_thresholds"]
    example["anchors"] = anchors
    anchors_mask = None
    if anchor_area_threshold >= 0:
        # slow with high resolution. recommend disable this forever.
        coors = coordinates
        dense_voxel_map = box_np_ops.sparse_sum_for_anchors_mask(
            coors, tuple(grid_size[::-1][1:]))
        dense_voxel_map = dense_voxel_map.cumsum(0)
        dense_voxel_map = dense_voxel_map.cumsum(1)
        anchors_area = box_np_ops.fused_get_anchors_area(
            dense_voxel_map, anchors_bv, voxel_size, pc_range, grid_size)
        anchors_mask = anchors_area > anchor_area_threshold
        # example['anchors_mask'] = anchors_mask.astype(np.uint8)
        example['anchors_mask'] = anchors_mask
    # print("prep time", time.time() - t)
    metrics["prep_time"] = time.time() - t
    if not training:
        return example
    example["gt_names"] = gt_dict["gt_names"]
    # voxel_labels = box_np_ops.assign_label_to_voxel(gt_boxes, coordinates,
    #                                                 voxel_size, coors_range)
    if create_targets:
        t1 = time.time()
        targets_dict = target_assigner.assign(
            anchors,
            anchors_dict,
            gt_dict["gt_boxes"],
            anchors_mask,
            gt_classes=gt_dict["gt_classes"],
            gt_names=gt_dict["gt_names"],
            matched_thresholds=matched_thresholds,
            unmatched_thresholds=unmatched_thresholds,
            importance=gt_dict["gt_importance"])

        """
        boxes_lidar = gt_dict["gt_boxes"]
        bev_map = simplevis.nuscene_vis(points, boxes_lidar, gt_dict["gt_names"])
        assigned_anchors = anchors[targets_dict['labels'] > 0]
        ignored_anchors = anchors[targets_dict['labels'] == -1]
        bev_map = simplevis.draw_box_in_bev(bev_map, [-50, -50, 3, 50, 50, 1], ignored_anchors, [128, 128, 128], 2)
        bev_map = simplevis.draw_box_in_bev(bev_map, [-50, -50, 3, 50, 50, 1], assigned_anchors, [255, 0, 0])
        cv2.imshow('anchors', bev_map)
        cv2.waitKey(0)
        
        boxes_lidar = gt_dict["gt_boxes"]
        pp_map = np.zeros(grid_size[:2], dtype=np.float32)
        voxels_max = np.max(voxels[:, :, 2], axis=1, keepdims=False)
        voxels_min = np.min(voxels[:, :, 2], axis=1, keepdims=False)
        voxels_height = voxels_max - voxels_min
        voxels_height = np.minimum(voxels_height, 4)
        # sns.distplot(voxels_height)
        # plt.show()
        pp_map[coordinates[:, 1], coordinates[:, 2]] = voxels_height / 4
        pp_map = (pp_map * 255).astype(np.uint8)
        pp_map = cv2.cvtColor(pp_map, cv2.COLOR_GRAY2RGB)
        pp_map = simplevis.draw_box_in_bev(pp_map, [-50, -50, 3, 50, 50, 1], boxes_lidar, [128, 0, 128], 1)
        cv2.imshow('heights', pp_map)
        cv2.waitKey(0)
        """
        example.update({
            'labels': targets_dict['labels'],
            'reg_targets': targets_dict['bbox_targets'],
            # 'reg_weights': targets_dict['bbox_outside_weights'],
            'importance': targets_dict['importance'],
        })
    return example
예제 #3
0
def prep_pointcloud(input_dict,
                    root_path,
                    voxel_generator,
                    target_assigner,
                    use_quadrant=False,
                    db_sampler=None,
                    max_voxels=20000,
                    remove_outside_points=False,
                    training=True,
                    create_targets=True,
                    shuffle_points=False,
                    remove_unknown=False,
                    gt_rotation_noise=(-np.pi / 3, np.pi / 3),
                    gt_loc_noise_std=(1.0, 1.0, 1.0),
                    global_rotation_noise=(-np.pi / 4, np.pi / 4),
                    global_scaling_noise=(0.95, 1.05),
                    global_random_rot_range=(0.78, 2.35),
                    global_translate_noise_std=(0, 0, 0),
                    num_point_features=4,
                    remove_points_after_sample=True,
                    remove_environment=False,
                    random_crop=False,
                    reference_detections=None,
                    out_size_factor=2,
                    use_group_id=False,
                    multi_gpu=False,
                    min_points_in_gt=-1,
                    random_flip_x=True,
                    random_flip_y=True,
                    sample_importance=1.0,
                    dataset_name = 'KITTI'):
    """convert point cloud to voxels, create targets if ground truths 
    exists.

    input_dict format: dataset.get_sensor_data format

    """
    t = time.time()
    class_names = target_assigner.classes
    points = input_dict["lidar"]["points"]


    if training:
        anno_dict = input_dict["lidar"]["annotations"]
        gt_dict = {
            "gt_boxes": anno_dict["boxes"],
            "gt_names": anno_dict["names"],
            "gt_importance": np.ones([anno_dict["boxes"].shape[0]], dtype=anno_dict["boxes"].dtype),
        }
        if "difficulty" not in anno_dict:
            difficulty = np.zeros([anno_dict["boxes"].shape[0]],
                                  dtype=np.int32)
            gt_dict["difficulty"] = difficulty
        else:
            gt_dict["difficulty"] = anno_dict["difficulty"]
        if use_group_id and "group_ids" in anno_dict:
            group_ids = anno_dict["group_ids"]
            gt_dict["group_ids"] = group_ids
    calib = None
    if "calib" in input_dict:
        calib = input_dict["calib"]

    if reference_detections is not None:
        assert calib is not None and "image" in input_dict
        C, R, T = box_np_ops.projection_matrix_to_CRT_kitti(P2)
        frustums = box_np_ops.get_frustum_v2(reference_detections, C)
        frustums -= T
        frustums = np.einsum('ij, akj->aki', np.linalg.inv(R), frustums)
        frustums = box_np_ops.camera_to_lidar(frustums, rect, Trv2c)
        surfaces = box_np_ops.corner_to_surfaces_3d_jit(frustums)
        masks = points_in_convex_polygon_3d_jit(points, surfaces)
        points = points[masks.any(-1)]

    if remove_outside_points:
        assert calib is not None
        image_shape = input_dict["image"]["image_shape"]
        points = box_np_ops.remove_outside_points(
            points, calib["rect"], calib["Trv2c"], calib["P2"], image_shape)
    if remove_environment and training:
        selected = kitti.keep_arrays_by_name(gt_names, target_assigner.classes)
        _dict_select(gt_dict, selected)
        masks = box_np_ops.points_in_rbbox(points, gt_dict["gt_boxes"])
        points = points[masks.any(-1)]
    metrics = {}

    if training:
        # boxes_lidar = gt_dict["gt_boxes"]
        # bev_map = simplevis.nuscene_vis(points, boxes_lidar)
        # cv2.imshow('pre-noise', bev_map)
        selected = kitti.drop_arrays_by_name(gt_dict["gt_names"], ["DontCare"])
        _dict_select(gt_dict, selected)
        if remove_unknown:
            remove_mask = gt_dict["difficulty"] == -1
            """
            gt_boxes_remove = gt_boxes[remove_mask]
            gt_boxes_remove[:, 3:6] += 0.25
            points = prep.remove_points_in_boxes(points, gt_boxes_remove)
            """
            keep_mask = np.logical_not(remove_mask)
            _dict_select(gt_dict, keep_mask)
        gt_dict.pop("difficulty")
        if min_points_in_gt > 0:
            # points_count_rbbox takes 10ms with 10 sweeps nuscenes data
            point_counts = box_np_ops.points_count_rbbox(points, gt_dict["gt_boxes"])
            mask = point_counts >= min_points_in_gt
            _dict_select(gt_dict, mask)
        # select the gt_box in the specified classes
        gt_boxes_mask = np.array(
            [n in class_names for n in gt_dict["gt_names"]], dtype=np.bool_)
        # data augmentation
        if db_sampler is not None:
            group_ids = None
            if "group_ids" in gt_dict:
                group_ids = gt_dict["group_ids"]

            # data augmentation, using sample to add target
            sampled_dict = db_sampler.sample_all(
                root_path,
                gt_dict["gt_boxes"],
                gt_dict["gt_names"],
                num_point_features,
                random_crop,
                gt_group_ids=group_ids,
                calib=calib)

            if sampled_dict is not None:
                sampled_gt_names = sampled_dict["gt_names"]
                sampled_gt_boxes = sampled_dict["gt_boxes"]
                sampled_points = sampled_dict["points"]
                sampled_gt_masks = sampled_dict["gt_masks"]
                gt_dict["gt_names"] = np.concatenate(
                    [gt_dict["gt_names"], sampled_gt_names], axis=0)
                gt_dict["gt_boxes"] = np.concatenate(
                    [gt_dict["gt_boxes"], sampled_gt_boxes])
                gt_boxes_mask = np.concatenate(
                    [gt_boxes_mask, sampled_gt_masks], axis=0)
                sampled_gt_importance = np.full([sampled_gt_boxes.shape[0]], sample_importance, dtype=sampled_gt_boxes.dtype)
                gt_dict["gt_importance"] = np.concatenate(
                    [gt_dict["gt_importance"], sampled_gt_importance])

                if group_ids is not None:
                    sampled_group_ids = sampled_dict["group_ids"]
                    gt_dict["group_ids"] = np.concatenate(
                        [gt_dict["group_ids"], sampled_group_ids])
                # remove the raw points in the added box, to avoid overlap
                if remove_points_after_sample:
                    masks = box_np_ops.points_in_rbbox(points,
                                                       sampled_gt_boxes)
                    points = points[np.logical_not(masks.any(-1))]

                points = np.concatenate([sampled_points, points], axis=0)
        group_ids = None
        if "group_ids" in gt_dict:
            group_ids = gt_dict["group_ids"]

        prep.noise_per_object_v3_(gt_dict["gt_boxes"],
                                    points,
                                    gt_boxes_mask,
                                    rotation_perturb=gt_rotation_noise,
                                    center_noise_std=gt_loc_noise_std,
                                    global_random_rot_range=global_random_rot_range,
                                    group_ids=group_ids,
                                    num_try=100)

        # should remove unrelated objects after noise per object
        # for k, v in gt_dict.items():
        #     print(k, v.shape)
        _dict_select(gt_dict, gt_boxes_mask)
        gt_classes = np.array(
            [class_names.index(n) + 1 for n in gt_dict["gt_names"]],
            dtype=np.int32)
        gt_dict["gt_classes"] = gt_classes
        gt_dict["gt_boxes"], points = prep.random_flip(gt_dict["gt_boxes"],
                                                       points, 0.5, random_flip_x, random_flip_y)
        gt_dict["gt_boxes"], points = prep.global_rotation_v2(
            gt_dict["gt_boxes"], points, *global_rotation_noise)
        gt_dict["gt_boxes"], points = prep.global_scaling_v2(
            gt_dict["gt_boxes"], points, *global_scaling_noise)
        prep.global_translate_(gt_dict["gt_boxes"], points, global_translate_noise_std)
        bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]]
        mask = prep.filter_gt_box_outside_range_by_center(gt_dict["gt_boxes"], bv_range)
        _dict_select(gt_dict, mask)

        # limit rad to [-pi, pi]
        gt_dict["gt_boxes"][:, 6] = box_np_ops.limit_period(
            gt_dict["gt_boxes"][:, 6], offset=0.5, period=2 * np.pi)

        # boxes_lidar = gt_dict["gt_boxes"]
        # bev_map = simplevis.nuscene_vis(points, boxes_lidar)
        # cv2.imshow('post-noise', bev_map)
        # cv2.waitKey(0)
    if shuffle_points:
        # shuffle is a little slow.
        np.random.shuffle(points)

    # [0, -40, -3, 70.4, 40, 1]

    grid_size = voxel_generator.grid_size
    feature_map_size = grid_size[:2] // out_size_factor
    feature_map_size = [*feature_map_size, 1][::-1]
    # [352, 400]


    t1 = time.time()
    if not multi_gpu:
        res = voxel_generator.generate(points, max_voxels)
        voxels = res["voxels"]
        coordinates = res["coordinates"]
        num_points = res["num_points_per_voxel"]
        num_voxels = np.array([voxels.shape[0]], dtype=np.int64)
    else:
        res = voxel_generator.generate_multi_gpu(points, max_voxels)
        voxels = res["voxels"]
        coordinates = res["coordinates"]
        num_points = res["num_points_per_voxel"]
        num_voxels = np.array([res["voxel_num"]], dtype=np.int64)
    metrics["voxel_gene_time"] = time.time() - t1

    anchors_all = generate_anchors(target_assigner.classes_cfg, feature_map_size, use_quadrant)
    if use_quadrant:
        split_res = split_voxel_into_quadrants(res, grid_size)
        voxels = split_res["voxels"]
        coordinates = split_res["coordinates"]




    example = { 'voxels': voxels,
                'num_points': num_points,
                'coordinates': coordinates,
                "num_voxels": num_voxels,
                "metrics": metrics ,
                'anchors': anchors_all["anchors"]}
    if calib is not None:
        example["calib"] = calib


    metrics["prep_time"] = time.time() - t

    if not training:
        return example
    # voxel_labels = box_np_ops.assign_label_to_voxel(gt_boxes, coordinates,
    #                                                 voxel_size, coors_range)

    if create_targets:
        if target_assigner.name  == 'LabelAssigner':
            if use_quadrant:
                res = []
                gt_box = split_dict["gt_boxes"]
                gt_classes = split_dict["gt_classes"]
                gt_names = split_dict["gt_name"]
                gt_importance = split_dict["gt_importance"]
                for i in range(4):
                    d = target_assigner.assign(
                        gt_boxes=gt_box[i],
                        feature_map_size=feature_map_size,
                        gt_classes=gt_classes[i],
                        gt_names=gt_names[i],
                        importance=gt_importance[i],
                        training=training,
                        dataset_name=dataset_name
                    )
                    res.append(d)
            else:
                targets_dict = target_assigner.assign(
                    gt_dict["gt_boxes"],
                    feature_map_size,
                    gt_classes=gt_dict["gt_classes"],
                    gt_names=gt_dict["gt_names"],
                    importance=gt_dict["gt_importance"],
                    training=training,
                    dataset_name= dataset_name
                )
            example.update({
                'gt_dict': targets_dict['gt_dict'],
                'hm': targets_dict['targets']['hm'],
                'anno_box': targets_dict['targets']['anno_box'],
                'ind': targets_dict['targets']['ind'],
                'mask': targets_dict['targets']['mask'],
                'cat': targets_dict['targets']['cat']
            })
        else:
            if use_quadrant:
                split_dict = split_gt_into_quadrants(gt_dict)
                example['use_quadrant']= use_quadrant
                targets_dict = {}
                gt_box = split_dict["gt_boxes"]
                gt_classes = split_dict["gt_classes"]
                gt_names = split_dict["gt_names"]
                gt_importance = split_dict["gt_importance"]
                labels=[]
                bbox_targets=[]
                importance = []
                for i in range(4):
                    d = target_assigner.assign(
                        anchors_all,
                        gt_boxes=gt_box[i],
                        gt_classes=gt_classes[i],
                        gt_names=gt_names[i],
                        importance=gt_importance[i]
                    )

                    labels.append(d['labels'])
                    bbox_targets.append(d['bbox_targets'])
                    importance.append(d['importance'])
                targets_dict['labels']=np.stack(labels)
                targets_dict['bbox_targets'] = np.stack(bbox_targets)
                targets_dict['importance'] = np.stack(importance)

            else:
                targets_dict = target_assigner.assign(
                    anchors_all,
                    gt_dict["gt_boxes"],
                    gt_classes=gt_dict["gt_classes"],
                    gt_names=gt_dict["gt_names"],
                    importance=gt_dict["gt_importance"])

            example.update({'targets': targets_dict})


    return example
예제 #4
0
def prep_data_aug(points,
                  gt_dict,
                  calib,
                  root_path,
                  voxel_generator,
                  target_assigner,
                  db_sampler=None,
                  remove_unknown=False,
                  gt_rotation_noise=(-np.pi / 3, np.pi / 3),
                  gt_loc_noise_std=(1.0, 1.0, 1.0),
                  global_rotation_noise=(-np.pi / 4, np.pi / 4),
                  global_scaling_noise=(0.95, 1.05),
                  global_random_rot_range=(0.78, 2.35),
                  global_translate_noise_std=(0, 0, 0),
                  num_point_features=4,
                  remove_points_after_sample=True,
                  remove_environment=False,
                  random_crop=False,
                  min_points_in_gt=-1,
                  random_flip_x=True,
                  random_flip_y=True,
                  sample_importance=1.0):
    """
    Performs data augmentation (which is usually done during training time).

    Takes as input points and gt_dict, and returns updated points and gt_dict after applying
    the various data augmentation methods.

    Main args:
        points: np.ndarray(dtype=float32, shape=(N, 3+)) TODO: confirm this.
        gt_dict: dict with keys "gt_boxes", "gt_names", "gt_importance", "gt_difficulty",
        calib: calibration.
    Returns:
        points: np.ndarray(dtype=float32, shape=(N, 3+)) Modified points. TODO: confirm this.
        gt_dict: dict with keys "gt_boxes", "gt_names", "gt_importance", "gt_difficulty"
    """
    gt_dict = gt_dict.copy()
    class_names = target_assigner.classes

    if remove_environment is True:
        selected = kitti.keep_arrays_by_name(gt_dict["gt_names"], target_assigner.classes)
        _dict_select(gt_dict, selected)
        masks = box_np_ops.points_in_rbbox(points, gt_dict["gt_boxes"])
        points = points[masks.any(-1)]
    metrics = {}

    """
    boxes_lidar = gt_dict["gt_boxes"]
    bev_map = simplevis.nuscene_vis(points, boxes_lidar)
    cv2.imshow('pre-noise', bev_map)
    """
    selected = kitti.drop_arrays_by_name(gt_dict["gt_names"], ["DontCare"])
    _dict_select(gt_dict, selected)
    if remove_unknown:
        remove_mask = gt_dict["difficulty"] == -1
        """
        gt_boxes_remove = gt_boxes[remove_mask]
        gt_boxes_remove[:, 3:6] += 0.25
        points = prep.remove_points_in_boxes(points, gt_boxes_remove)
        """
        keep_mask = np.logical_not(remove_mask)
        _dict_select(gt_dict, keep_mask)
    gt_dict.pop("difficulty")
    if min_points_in_gt > 0:
        # points_count_rbbox takes 10ms with 10 sweeps nuscenes data
        point_counts = box_np_ops.points_count_rbbox(points, gt_dict["gt_boxes"])
        mask = point_counts >= min_points_in_gt
        _dict_select(gt_dict, mask)
    gt_boxes_mask = np.array(
        [n in class_names for n in gt_dict["gt_names"]], dtype=np.bool_)
    if db_sampler is not None:
        group_ids = None
        if "group_ids" in gt_dict:
            group_ids = gt_dict["group_ids"]

        sampled_dict = db_sampler.sample_all(
            root_path,
            gt_dict["gt_boxes"],
            gt_dict["gt_names"],
            num_point_features,
            random_crop,
            gt_group_ids=group_ids,
            calib=calib)

        if sampled_dict is not None:
            sampled_gt_names = sampled_dict["gt_names"]
            sampled_gt_boxes = sampled_dict["gt_boxes"]
            sampled_points = sampled_dict["points"]
            sampled_gt_masks = sampled_dict["gt_masks"]
            gt_dict["gt_names"] = np.concatenate(
                [gt_dict["gt_names"], sampled_gt_names], axis=0)
            gt_dict["gt_boxes"] = np.concatenate(
                [gt_dict["gt_boxes"], sampled_gt_boxes])
            gt_boxes_mask = np.concatenate(
                [gt_boxes_mask, sampled_gt_masks], axis=0)
            sampled_gt_importance = np.full([sampled_gt_boxes.shape[0]], sample_importance, dtype=sampled_gt_boxes.dtype)
            gt_dict["gt_importance"] = np.concatenate(
                [gt_dict["gt_importance"], sampled_gt_importance])

            if group_ids is not None:
                sampled_group_ids = sampled_dict["group_ids"]
                gt_dict["group_ids"] = np.concatenate(
                    [gt_dict["group_ids"], sampled_group_ids])

            if remove_points_after_sample:
                masks = box_np_ops.points_in_rbbox(points,
                                                    sampled_gt_boxes)
                points = points[np.logical_not(masks.any(-1))]

            points = np.concatenate([sampled_points, points], axis=0)
    group_ids = None
    if "group_ids" in gt_dict:
        group_ids = gt_dict["group_ids"]

    prep.noise_per_object_v3_(
        gt_dict["gt_boxes"],
        points,
        gt_boxes_mask,
        rotation_perturb=gt_rotation_noise,
        center_noise_std=gt_loc_noise_std,
        global_random_rot_range=global_random_rot_range,
        group_ids=group_ids,
        num_try=100)

    # should remove unrelated objects after noise per object
    # for k, v in gt_dict.items():
    #     print(k, v.shape)
    _dict_select(gt_dict, gt_boxes_mask)
    gt_classes = np.array(
        [class_names.index(n) + 1 for n in gt_dict["gt_names"]],
        dtype=np.int32)
    gt_dict["gt_classes"] = gt_classes
    gt_dict["gt_boxes"], points = prep.random_flip(gt_dict["gt_boxes"],
                                                    points, 0.5, random_flip_x, random_flip_y)
    gt_dict["gt_boxes"], points = prep.global_rotation_v2(
        gt_dict["gt_boxes"], points, *global_rotation_noise)
    gt_dict["gt_boxes"], points = prep.global_scaling_v2(
        gt_dict["gt_boxes"], points, *global_scaling_noise)
    prep.global_translate_(gt_dict["gt_boxes"], points, global_translate_noise_std)
    bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]]
    mask = prep.filter_gt_box_outside_range_by_center(gt_dict["gt_boxes"], bv_range)
    _dict_select(gt_dict, mask)

    # limit rad to [-pi, pi]
    gt_dict["gt_boxes"][:, 6] = box_np_ops.limit_period(
        gt_dict["gt_boxes"][:, 6], offset=0.5, period=2 * np.pi)

    # boxes_lidar = gt_dict["gt_boxes"]
    # bev_map = simplevis.nuscene_vis(points, boxes_lidar)
    # cv2.imshow('post-noise', bev_map)
    # cv2.waitKey(0)

    return points, gt_dict
예제 #5
0
def prep_pointcloud(input_dict,
                    root_path,
                    voxel_generator,
                    target_assigner,
                    db_sampler=None,
                    max_voxels=20000,
                    max_sweeps=10,
                    remove_outside_points=False,
                    training=True,
                    create_targets=True,
                    shuffle_points=False,
                    remove_unknown=False,
                    gt_rotation_noise=(-np.pi / 3, np.pi / 3),
                    gt_loc_noise_std=(1.0, 1.0, 1.0),
                    global_rotation_noise=(-np.pi / 4, np.pi / 4),
                    global_scaling_noise=(0.95, 1.05),
                    global_random_rot_range=(0.78, 2.35),
                    global_translate_noise_std=(0, 0, 0),
                    num_point_features=4,
                    anchor_area_threshold=1,
                    gt_points_drop=0.0,
                    gt_drop_max_keep=10,
                    remove_points_after_sample=True,
                    anchor_cache=None,
                    remove_environment=False,
                    random_crop=False,
                    reference_detections=None,
                    out_size_factor=2,
                    use_group_id=False,
                    multi_gpu=False,
                    min_points_in_gt=-1,
                    random_flip_x=True,
                    random_flip_y=True,
                    sample_importance=1.0,
                    out_dtype=np.float32):
    """convert point cloud to voxels, create targets if ground truths
    exists.

    input_dict format: dataset.get_sensor_data format

    """
    t = time.time()
    class_names = target_assigner.classes
    points = input_dict["lidar"]["points"]
    indices = input_dict["lidar"]["indices"]
    origins = input_dict["lidar"]["origins"]
    if training:
        anno_dict = input_dict["lidar"]["annotations"]
        gt_dict = {
            "gt_boxes":
            anno_dict["boxes"],
            "gt_names":
            anno_dict["names"],
            "gt_importance":
            np.ones([anno_dict["boxes"].shape[0]],
                    dtype=anno_dict["boxes"].dtype),
        }
        if "difficulty" not in anno_dict:
            difficulty = np.zeros([anno_dict["boxes"].shape[0]],
                                  dtype=np.int32)
            gt_dict["difficulty"] = difficulty
        else:
            gt_dict["difficulty"] = anno_dict["difficulty"]
        if use_group_id and "group_ids" in anno_dict:
            group_ids = anno_dict["group_ids"]
            gt_dict["group_ids"] = group_ids
    calib = None
    if "calib" in input_dict:
        calib = input_dict["calib"]

    # # Disable these two since we do not do this for NuScenes
    # if reference_detections is not None:
    #     assert calib is not None and "image" in input_dict
    #     C, R, T = box_np_ops.projection_matrix_to_CRT_kitti(P2)
    #     frustums = box_np_ops.get_frustum_v2(reference_detections, C)
    #     frustums -= T
    #     frustums = np.einsum('ij, akj->aki', np.linalg.inv(R), frustums)
    #     frustums = box_np_ops.camera_to_lidar(frustums, rect, Trv2c)
    #     surfaces = box_np_ops.corner_to_surfaces_3d_jit(frustums)
    #     masks = points_in_convex_polygon_3d_jit(points, surfaces)
    #     points = points[masks.any(-1)]
    # if remove_outside_points:
    #     assert calib is not None
    #     image_shape = input_dict["image"]["image_shape"]
    #     points = box_np_ops.remove_outside_points(
    #         points, calib["rect"], calib["Trv2c"], calib["P2"], image_shape)

    # # Very interesting attempt
    # # I have tried the same and found it doesn't really work
    # if remove_environment is True and training:
    #     selected = kitti.keep_arrays_by_name(gt_names, target_assigner.classes)
    #     _dict_select(gt_dict, selected)
    #     masks = box_np_ops.points_in_rbbox(points, gt_dict["gt_boxes"])
    #     points = points[masks.any(-1)]

    metrics = {}

    point_indices_to_remove = None
    if training:
        """
        boxes_lidar = gt_dict["gt_boxes"]
        bev_map = simplevis.nuscene_vis(points, boxes_lidar)
        cv2.imshow('pre-noise', bev_map)
        """
        selected = kitti.drop_arrays_by_name(gt_dict["gt_names"], ["Denture"])
        _dict_select(gt_dict, selected)
        if remove_unknown:
            remove_mask = gt_dict["difficulty"] == -1
            """
            gt_boxes_remove = gt_boxes[remove_mask]
            gt_boxes_remove[:, 3:6] += 0.25
            points = prep.remove_points_in_boxes(points, gt_boxes_remove)
            """
            keep_mask = np.logical_not(remove_mask)
            _dict_select(gt_dict, keep_mask)
        gt_dict.pop("difficulty")

        # This part is interesting - we will need to do the same
        if min_points_in_gt > 0:
            # points_count_rbbox takes 10ms with 10 sweeps nuscenes data
            point_counts = box_np_ops.points_count_rbbox(
                points, gt_dict["gt_boxes"])
            mask = point_counts >= min_points_in_gt
            _dict_select(gt_dict, mask)

        gt_boxes_mask = np.array(
            [n in class_names for n in gt_dict["gt_names"]], dtype=np.bool_)

        if db_sampler is not None:
            group_ids = None
            if "group_ids" in gt_dict:
                group_ids = gt_dict["group_ids"]

            sampled_dict = db_sampler.sample_all(root_path,
                                                 gt_dict["gt_boxes"],
                                                 gt_dict["gt_names"],
                                                 num_point_features,
                                                 random_crop,
                                                 gt_group_ids=group_ids,
                                                 calib=calib)

            if sampled_dict is not None:
                sampled_gt_names = sampled_dict["gt_names"]
                sampled_gt_boxes = sampled_dict["gt_boxes"]
                sampled_points = sampled_dict["points"]
                sampled_gt_masks = sampled_dict["gt_masks"]
                gt_dict["gt_names"] = np.concatenate(
                    [gt_dict["gt_names"], sampled_gt_names], axis=0)
                gt_dict["gt_boxes"] = np.concatenate(
                    [gt_dict["gt_boxes"], sampled_gt_boxes])
                gt_boxes_mask = np.concatenate(
                    [gt_boxes_mask, sampled_gt_masks], axis=0)
                sampled_gt_importance = np.full([sampled_gt_boxes.shape[0]],
                                                sample_importance,
                                                dtype=sampled_gt_boxes.dtype)
                gt_dict["gt_importance"] = np.concatenate(
                    [gt_dict["gt_importance"], sampled_gt_importance])

                if group_ids is not None:
                    sampled_group_ids = sampled_dict["group_ids"]
                    gt_dict["group_ids"] = np.concatenate(
                        [gt_dict["group_ids"], sampled_group_ids])

                # # Commented out because we have a new way of removing points
                # if remove_points_after_sample:
                #     masks = box_np_ops.points_in_rbbox(points, sampled_gt_boxes)
                #     point_indices_to_remove = np.flatnonzero(masks.any(-1))
                #     # # Delay this process so we can use the full point cloud
                #     # # when we do the ray stopping algorithm
                #     # points = points[np.logical_not(masks.any(-1))]

                # # Paste objects behind so that we don't have to update indices
                # points = np.concatenate([sampled_points, points], axis=0)
                points = np.concatenate([points, sampled_points], axis=0)

        pc_range = voxel_generator.point_cloud_range
        group_ids = None
        if "group_ids" in gt_dict:
            group_ids = gt_dict["group_ids"]

        # # Disable this one for now (not used in PointPillars anyways)
        # prep.noise_per_object_v3_(
        #     gt_dict["gt_boxes"],
        #     points,
        #     gt_boxes_mask,
        #     rotation_perturb=gt_rotation_noise,
        #     center_noise_std=gt_loc_noise_std,
        #     global_random_rot_range=global_random_rot_range,
        #     group_ids=group_ids,
        #     num_try=100)

        # should remove unrelated objects after noise per object
        # for k, v in gt_dict.items():
        #     print(k, v.shape)
        _dict_select(gt_dict, gt_boxes_mask)
        gt_classes = np.array(
            [class_names.index(n) + 1 for n in gt_dict["gt_names"]],
            dtype=np.int32)
        gt_dict["gt_classes"] = gt_classes
        gt_dict["gt_boxes"], points, origins = prep.random_flip(
            gt_dict["gt_boxes"], points, origins, 0.5, random_flip_x,
            random_flip_y)
        gt_dict["gt_boxes"], points, origins = prep.global_rotation_v2(
            gt_dict["gt_boxes"], points, origins, *global_rotation_noise)
        gt_dict["gt_boxes"], points, origins = prep.global_scaling_v2(
            gt_dict["gt_boxes"], points, origins, *global_scaling_noise)
        prep.global_translate_(gt_dict["gt_boxes"], points, origins,
                               global_translate_noise_std)
        bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]]
        mask = prep.filter_gt_box_outside_range_by_center(
            gt_dict["gt_boxes"], bv_range)
        _dict_select(gt_dict, mask)

        # limit rad to [-pi, pi]
        gt_dict["gt_boxes"][:, 6] = box_np_ops.limit_period(
            gt_dict["gt_boxes"][:, 6], offset=0.5, period=2 * np.pi)

        # boxes_lidar = gt_dict["gt_boxes"]
        # bev_map = simplevis.nuscene_vis(points, boxes_lidar)
        # cv2.imshow('post-noise', bev_map)
        # cv2.waitKey(0)

    # # Disable this for now (not used in PointPillars anyways)
    # if shuffle_points:
    #     # shuffle is a little slow.
    #     np.random.shuffle(points)

    # [0, -40, -3, 70.4, 40, 1]
    voxel_size = voxel_generator.voxel_size
    pc_range = voxel_generator.point_cloud_range
    grid_size = voxel_generator.grid_size

    # organize points into lists based on timestamps
    time_stamps = points[
        indices[:-1],
        -1]  # counting on the fact we do not miss points from any intermediate time_stamps
    time_stamps = (time_stamps[:-1] + time_stamps[1:]) / 2
    time_stamps = [-1000.0] + time_stamps.tolist() + [1000.0]  # add boundaries
    time_stamps = np.array(time_stamps)

    # # LL_OCCUPIED, LL_FREE = 0.85, -0.4
    # lo_occupied = np.log(0.7 / (1 - 0.7))
    # lo_free = np.log(0.4 / (1 - 0.4))

    # is there are additional points (from database sampling)
    num_original = indices[-1]
    if len(points) > num_original:
        # split data into two half (indexed and un-indexed)
        original_points, sampled_points = points[:num_original], points[
            num_original:]
        # compute occupancy and masks
        # visibility, original_mask, sampled_mask = mapping.compute_visibility_and_masks(
        #     original_points, sampled_points, origins, time_stamps, pc_range, min(voxel_size)
        # )
        logodds, original_mask, sampled_mask = mapping.compute_logodds_and_masks(
            original_points,
            sampled_points,
            origins,
            time_stamps,
            pc_range,
            min(voxel_size)  # , lo_occupied, lo_free
        )
        # apply visible mask
        points = np.concatenate(
            (original_points[original_mask], sampled_points[sampled_mask]))
    else:
        # visibility = mapping.compute_visibility(
        #     points, origins, time_stamps, pc_range, min(voxel_size)
        # )
        logodds = mapping.compute_logodds(
            points,
            origins,
            time_stamps,
            pc_range,
            min(voxel_size)  #, lo_occupied, lo_free
        )

    # T = len(time_stamps)-1
    # visibility = visibility.reshape(T, -1)
    # if T < (1 + max_sweeps):
    #     visibility = np.pad(visibility, ((0, (1+max_sweeps)-T), (0,0)), 'edge')

    # with open(f'./utils/mapping/examples/{time.time()}.pkl', 'wb') as f:
    #     ##
    #     pickle.dump(original_points, f)
    #     pickle.dump(sampled_points, f)
    #     pickle.dump(origins, f)
    #     pickle.dump(time_stamps, f)
    #     pickle.dump(pc_range, f)
    #     pickle.dump(voxel_size, f)
    #     ##
    #     pickle.dump(occupancy, f)
    #     pickle.dump(original_mask, f)
    #     pickle.dump(sampled_mask, f)

    if training:
        if min_points_in_gt > 0:
            # points_count_rbbox takes 10ms with 10 sweeps nuscenes data
            point_counts = box_np_ops.points_count_rbbox(
                points, gt_dict["gt_boxes"])
            mask = point_counts >= min_points_in_gt
            _dict_select(gt_dict, mask)

    # [352, 400]
    t1 = time.time()
    if not multi_gpu:
        res = voxel_generator.generate(points, max_voxels)
        voxels = res["voxels"]
        coordinates = res["coordinates"]
        num_points = res["num_points_per_voxel"]
        num_voxels = np.array([voxels.shape[0]], dtype=np.int64)
    else:
        res = voxel_generator.generate_multi_gpu(points, max_voxels)
        voxels = res["voxels"]
        coordinates = res["coordinates"]
        num_points = res["num_points_per_voxel"]
        num_voxels = np.array([res["voxel_num"]], dtype=np.int64)
    metrics["voxel_gene_time"] = time.time() - t1
    example = {
        'voxels': voxels,
        # 'visibility': visibility,
        'logodds': logodds,
        'num_points': num_points,
        'coordinates': coordinates,
        "num_voxels": num_voxels,
        "metrics": metrics,
    }
    if calib is not None:
        example["calib"] = calib
    feature_map_size = grid_size[:2] // out_size_factor
    feature_map_size = [*feature_map_size, 1][::-1]
    # print(f'feature_map_size in prep_pointcloud(): {feature_map_size}')
    if anchor_cache is not None:
        # print('having anchor cache')
        anchors = anchor_cache["anchors"]
        anchors_bv = anchor_cache["anchors_bv"]
        anchors_dict = anchor_cache["anchors_dict"]
        matched_thresholds = anchor_cache["matched_thresholds"]
        unmatched_thresholds = anchor_cache["unmatched_thresholds"]

    else:
        # print('NOT having anchor cache')
        ret = target_assigner.generate_anchors(feature_map_size)
        anchors = ret["anchors"]
        anchors = anchors.reshape([-1, target_assigner.box_ndim])
        anchors_dict = target_assigner.generate_anchors_dict(feature_map_size)
        anchors_bv = box_np_ops.rbbox2d_to_near_bbox(anchors[:,
                                                             [0, 1, 3, 4, 6]])
        matched_thresholds = ret["matched_thresholds"]
        unmatched_thresholds = ret["unmatched_thresholds"]
    # print(f'anchors.shape: {anchors.shape}')

    example["anchors"] = anchors
    anchors_mask = None
    if anchor_area_threshold >= 0:
        # slow with high resolution. recommend disable this forever.
        coors = coordinates
        dense_voxel_map = box_np_ops.sparse_sum_for_anchors_mask(
            coors, tuple(grid_size[::-1][1:]))
        dense_voxel_map = dense_voxel_map.cumsum(0)
        dense_voxel_map = dense_voxel_map.cumsum(1)
        anchors_area = box_np_ops.fused_get_anchors_area(
            dense_voxel_map, anchors_bv, voxel_size, pc_range, grid_size)
        anchors_mask = anchors_area > anchor_area_threshold
        # example['anchors_mask'] = anchors_mask.astype(np.uint8)
        example['anchors_mask'] = anchors_mask
    # print("prep time", time.time() - t)
    metrics["prep_time"] = time.time() - t
    if not training:
        return example
    example["gt_names"] = gt_dict["gt_names"]
    # voxel_labels = box_np_ops.assign_label_to_voxel(gt_boxes, coordinates,
    #                                                 voxel_size, coors_range)
    if create_targets:
        t1 = time.time()
        targets_dict = target_assigner.assign(
            anchors,
            anchors_dict,
            gt_dict["gt_boxes"],
            anchors_mask,
            gt_classes=gt_dict["gt_classes"],
            gt_names=gt_dict["gt_names"],
            matched_thresholds=matched_thresholds,
            unmatched_thresholds=unmatched_thresholds,
            importance=gt_dict["gt_importance"])
        """
        boxes_lidar = gt_dict["gt_boxes"]
        bev_map = simplevis.nuscene_vis(points, boxes_lidar, gt_dict["gt_names"])
        assigned_anchors = anchors[targets_dict['labels'] > 0]
        ignored_anchors = anchors[targets_dict['labels'] == -1]
        bev_map = simplevis.draw_box_in_bev(bev_map, [-50, -50, 3, 50, 50, 1], ignored_anchors, [128, 128, 128], 2)
        bev_map = simplevis.draw_box_in_bev(bev_map, [-50, -50, 3, 50, 50, 1], assigned_anchors, [255, 0, 0])
        cv2.imshow('anchors', bev_map)
        cv2.waitKey(0)

        boxes_lidar = gt_dict["gt_boxes"]
        pp_map = np.zeros(grid_size[:2], dtype=np.float32)
        voxels_max = np.max(voxels[:, :, 2], axis=1, keepdims=False)
        voxels_min = np.min(voxels[:, :, 2], axis=1, keepdims=False)
        voxels_height = voxels_max - voxels_min
        voxels_height = np.minimum(voxels_height, 4)
        # sns.distplot(voxels_height)
        # plt.show()
        pp_map[coordinates[:, 1], coordinates[:, 2]] = voxels_height / 4
        pp_map = (pp_map * 255).astype(np.uint8)
        pp_map = cv2.cvtColor(pp_map, cv2.COLOR_GRAY2RGB)
        pp_map = simplevis.draw_box_in_bev(pp_map, [-50, -50, 3, 50, 50, 1], boxes_lidar, [128, 0, 128], 1)
        cv2.imshow('heights', pp_map)
        cv2.waitKey(0)
        """
        example.update({
            'labels': targets_dict['labels'],
            'reg_targets': targets_dict['bbox_targets'],
            # 'reg_weights': targets_dict['bbox_outside_weights'],
            'importance': targets_dict['importance'],
        })
    return example