Exemple #1
0
    def __call__(self, res, info):
        targets = {}

        # Calculate output feature map size for anchor generation.
        grid_size = res["lidar"]["voxels"]["shape"]  # [1408, 1600,  40]

        # get anchors: [x, y, z, w(x-axis), l(y-axis), h, ry].  [(70400, 7),]
        targets["anchors"] = [
            anchor_dict[self.target_class_names[i]]["anchors"].reshape([-1, 7])
            for i, anchor_dict in enumerate(self.anchor_dicts_by_task)
        ]

        # get gt labels of targeted classes; limit ry range in [-pi, pi].
        if res["mode"] == "train":

            gt_dict = res["lidar"]["annotations"]

            gt_mask = np.zeros(gt_dict["gt_classes"].shape, dtype=np.bool)
            for target_class_id in self.target_class_ids:
                gt_mask = np.logical_or(
                    gt_mask, gt_dict["gt_classes"] == target_class_id)

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

            gt_dict["gt_boxes"] = [gt_boxes]
            gt_dict["gt_classes"] = [gt_dict["gt_classes"][gt_mask]]
            gt_dict["gt_names"] = [gt_dict["gt_names"][gt_mask]]

            res["lidar"]["annotations"] = gt_dict

        # get anchor classification labels and localization regression labels
        if res["mode"] == "train":
            targets_dict = {}
            for idx, target_assigner in enumerate(self.target_assigners):
                targets_dict = target_assigner.assign_v2(
                    self.anchor_dicts_by_task[idx],
                    gt_dict["gt_boxes"][idx],  # (x, y, z, w, l, h, r)
                    anchors_mask=None,
                    gt_classes=gt_dict["gt_classes"][idx],
                    gt_names=gt_dict["gt_names"][idx],
                    enable_similar_type=self.enable_similar_type,
                )

            targets.update({
                "labels": [targets_dict["labels"]],
                "reg_targets": [targets_dict["bbox_targets"]],
                "reg_weights": [targets_dict["bbox_outside_weights"]],
                "positive_gt_id": [targets_dict["positive_gt_id"]],
            })

        res["lidar"]["targets"] = targets
        return res, info
Exemple #2
0
    def __call__(self, res, info):
        max_objs = self._max_objs
        class_names_by_task = [t.class_names for t in self.tasks]
        num_classes_by_task = [t.num_class for t in self.tasks]

        # Calculate output featuremap size
        grid_size = res["lidar"]["voxels"]["shape"] 
        pc_range = res["lidar"]["voxels"]["range"]
        voxel_size = res["lidar"]["voxels"]["size"]

        feature_map_size = grid_size[:2] // self.out_size_factor
        example = {}

        if res["mode"] == "train":
            gt_dict = res["lidar"]["annotations"]

            # reorganize the gt_dict by tasks
            task_masks = []
            flag = 0
            for class_name in class_names_by_task:
                task_masks.append(
                    [
                        np.where(
                            gt_dict["gt_classes"] == class_name.index(i) + 1 + flag
                        )
                        for i in class_name
                    ]
                )
                flag += len(class_name)

            task_boxes = []
            task_classes = []
            task_names = []
            flag2 = 0
            for idx, mask in enumerate(task_masks):
                task_box = []
                task_class = []
                task_name = []
                for m in mask:
                    task_box.append(gt_dict["gt_boxes"][m])
                    task_class.append(gt_dict["gt_classes"][m] - flag2)
                    task_name.append(gt_dict["gt_names"][m])
                task_boxes.append(np.concatenate(task_box, axis=0))
                task_classes.append(np.concatenate(task_class))
                task_names.append(np.concatenate(task_name))
                flag2 += len(mask)

            for task_box in task_boxes:
                # limit rad to [-pi, pi]
                task_box[:, -1] = box_np_ops.limit_period(
                    task_box[:, -1], offset=0.5, period=np.pi * 2
                )

            # print(gt_dict.keys())
            gt_dict["gt_classes"] = task_classes
            gt_dict["gt_names"] = task_names
            gt_dict["gt_boxes"] = task_boxes

            res["lidar"]["annotations"] = gt_dict

            draw_gaussian = draw_umich_gaussian

            hms, anno_boxs, inds, masks, cats = [], [], [], [], []

            for idx, task in enumerate(self.tasks):
                hm = np.zeros((len(class_names_by_task[idx]), feature_map_size[1], feature_map_size[0]),
                              dtype=np.float32)

                if res['type'] == 'NuScenesDataset':
                    # [reg, hei, dim, vx, vy, rots, rotc]
                    anno_box = np.zeros((max_objs, 10), dtype=np.float32)
                elif res['type'] == 'WaymoDataset':
                    anno_box = np.zeros((max_objs, 10), dtype=np.float32) 
                else:
                    raise NotImplementedError("Only Support nuScene for Now!")

                ind = np.zeros((max_objs), dtype=np.int64)
                mask = np.zeros((max_objs), dtype=np.uint8)
                cat = np.zeros((max_objs), dtype=np.int64)

                num_objs = min(gt_dict['gt_boxes'][idx].shape[0], max_objs)  

                for k in range(num_objs):
                    cls_id = gt_dict['gt_classes'][idx][k] - 1

                    w, l, h = gt_dict['gt_boxes'][idx][k][3], gt_dict['gt_boxes'][idx][k][4], \
                              gt_dict['gt_boxes'][idx][k][5]
                    w, l = w / voxel_size[0] / self.out_size_factor, l / voxel_size[1] / self.out_size_factor
                    if w > 0 and l > 0:
                        radius = gaussian_radius((l, w), min_overlap=self.gaussian_overlap)
                        radius = max(self._min_radius, int(radius))

                        # be really careful for the coordinate system of your box annotation. 
                        x, y, z = gt_dict['gt_boxes'][idx][k][0], gt_dict['gt_boxes'][idx][k][1], \
                                  gt_dict['gt_boxes'][idx][k][2]

                        coor_x, coor_y = (x - pc_range[0]) / voxel_size[0] / self.out_size_factor, \
                                         (y - pc_range[1]) / voxel_size[1] / self.out_size_factor

                        ct = np.array(
                            [coor_x, coor_y], dtype=np.float32)  
                        ct_int = ct.astype(np.int32)

                        # throw out not in range objects to avoid out of array area when creating the heatmap
                        if not (0 <= ct_int[0] < feature_map_size[0] and 0 <= ct_int[1] < feature_map_size[1]):
                            continue 

                        draw_gaussian(hm[cls_id], ct, radius)

                        new_idx = k
                        x, y = ct_int[0], ct_int[1]

                        cat[new_idx] = cls_id
                        ind[new_idx] = y * feature_map_size[0] + x
                        mask[new_idx] = 1

                        if res['type'] == 'NuScenesDataset': 
                            vx, vy = gt_dict['gt_boxes'][idx][k][6:8]
                            rot = gt_dict['gt_boxes'][idx][k][8]
                            anno_box[new_idx] = np.concatenate(
                                (ct - (x, y), z, np.log(gt_dict['gt_boxes'][idx][k][3:6]),
                                np.array(vx), np.array(vy), np.sin(rot), np.cos(rot)), axis=None)
                        elif res['type'] == 'WaymoDataset':
                            vx, vy = gt_dict['gt_boxes'][idx][k][6:8]
                            rot = gt_dict['gt_boxes'][idx][k][-1]
                            anno_box[new_idx] = np.concatenate(
                            (ct - (x, y), z, np.log(gt_dict['gt_boxes'][idx][k][3:6]),
                            np.array(vx), np.array(vy), np.sin(rot), np.cos(rot)), axis=None)
                        else:
                            raise NotImplementedError("Only Support Waymo and nuScene for Now")

                hms.append(hm)
                anno_boxs.append(anno_box)
                masks.append(mask)
                inds.append(ind)
                cats.append(cat)

            # used for two stage code 
            boxes = flatten(gt_dict['gt_boxes'])
            classes = merge_multi_group_label(gt_dict['gt_classes'], num_classes_by_task)

            if res["type"] == "NuScenesDataset":
                gt_boxes_and_cls = np.zeros((max_objs, 10), dtype=np.float32)
            elif res['type'] == "WaymoDataset":
                gt_boxes_and_cls = np.zeros((max_objs, 10), dtype=np.float32)
            else:
                raise NotImplementedError()

            boxes_and_cls = np.concatenate((boxes, 
                classes.reshape(-1, 1).astype(np.float32)), axis=1)
            num_obj = len(boxes_and_cls)
            assert num_obj <= max_objs
            # x, y, z, w, l, h, rotation_y, velocity_x, velocity_y, class_name
            boxes_and_cls = boxes_and_cls[:, [0, 1, 2, 3, 4, 5, 8, 6, 7, 9]]
            gt_boxes_and_cls[:num_obj] = boxes_and_cls

            example.update({'gt_boxes_and_cls': gt_boxes_and_cls})

            example.update({'hm': hms, 'anno_box': anno_boxs, 'ind': inds, 'mask': masks, 'cat': cats})
        else:
            pass

        res["lidar"]["targets"] = example

        return res, info
Exemple #3
0
    def convert_detection_to_kitti_annos(self, detection, partial=False):
        class_names = self._class_names
        det_image_idxes = [k for k in detection.keys()]
        gt_image_idxes = [str(info["image"]["image_idx"]) for info in self._kitti_infos]
        image_idxes = [gt_image_idxes, det_image_idxes]
        # print(f"det_image_idxes: {det_image_idxes[:10]}")
        # print(f"gt_image_idxes: {gt_image_idxes[:10]}")
        annos = []
        # for i in range(len(detection)):
        for det_idx in image_idxes[int(partial==True)]:
            det = detection[det_idx]
            info = self._kitti_infos[gt_image_idxes.index(det_idx)]
            # info = self._kitti_infos[i]
            calib = info["calib"]
            rect = calib["R0_rect"]
            Trv2c = calib["Tr_velo_to_cam"]
            P2 = calib["P2"]
            final_box_preds = det["box3d_lidar"].detach().cpu().numpy()
            label_preds = det["label_preds"].detach().cpu().numpy()
            scores = det["scores"].detach().cpu().numpy()

            anno = get_start_result_anno()
            num_example = 0

            if final_box_preds.shape[0] != 0:
                final_box_preds[:, -1] = box_np_ops.limit_period(final_box_preds[:, -1], offset=0.5, period=np.pi * 2,)
                final_box_preds[:, 2] -= final_box_preds[:, 5] / 2   # center_z -> bottom_z

                # aim: x, y, z, w, l, h, r -> -y, -z, x, h, w, l, r
                # (x, y, z, w, l, h r) in lidar -> (x', y', z', l, h, w, r) in camera
                box3d_camera = box_np_ops.box_lidar_to_camera(final_box_preds, rect, Trv2c)
                camera_box_origin = [0.5, 1.0, 0.5]
                box_corners = box_np_ops.center_to_corner_box3d(box3d_camera[:, :3], box3d_camera[:, 3:6], box3d_camera[:, 6], camera_box_origin, axis=1,)
                box_corners_in_image = box_np_ops.project_to_image(box_corners, P2)
                # box_corners_in_image: [N, 8, 2]
                minxy = np.min(box_corners_in_image, axis=1)
                maxxy = np.max(box_corners_in_image, axis=1)
                bbox = np.concatenate([minxy, maxxy], axis=1)

                for j in range(box3d_camera.shape[0]):
                    image_shape = info["image"]["image_shape"]
                    if bbox[j, 0] > image_shape[1] or bbox[j, 1] > image_shape[0]:
                        continue
                    if bbox[j, 2] < 0 or bbox[j, 3] < 0:
                        continue
                    bbox[j, 2:] = np.minimum(bbox[j, 2:], image_shape[::-1])
                    bbox[j, :2] = np.maximum(bbox[j, :2], [0, 0])
                    anno["bbox"].append(bbox[j])

                    anno["alpha"].append(-np.arctan2(-final_box_preds[j, 1], final_box_preds[j, 0]) + box3d_camera[j, 6])
                    # anno["dimensions"].append(box3d_camera[j, [4, 5, 3]])
                    anno["dimensions"].append(box3d_camera[j, 3:6])
                    anno["location"].append(box3d_camera[j, :3])
                    anno["rotation_y"].append(box3d_camera[j, 6])
                    anno["name"].append(class_names[int(label_preds[j])])
                    anno["truncated"].append(0.0)
                    anno["occluded"].append(0)
                    anno["score"].append(scores[j])

                    num_example += 1

            if num_example != 0:
                anno = {n: np.stack(v) for n, v in anno.items()}
                annos.append(anno)
            else:
                annos.append(empty_result_anno())
            num_example = annos[-1]["name"].shape[0]
            annos[-1]["metadata"] = det["metadata"]
        return annos
Exemple #4
0
    def __call__(self, res, info):
        max_objs = self._max_objs * self.dense_reg  # dense reg要扩充那个维度?

        # 2d list dim1:task dim2:classes
        class_names_by_task = [t.class_names for t in self.tasks]

        dxy = [(0, 0)]

        # Calculate output featuremap size
        grid_size = res["lidar"]["voxels"]["shape"]  # 448 x 512
        pc_range = res["lidar"]["voxels"]["range"]  # pcl范围实际大小
        voxel_size = res["lidar"]["voxels"]["size"]  # voxel实际大小

        feature_map_size = grid_size[:2] // self.out_size_factor
        example = {}

        if res["mode"] == "train":
            gt_dict = res["lidar"]["annotations"]

            # reorganize the gt_dict by tasks
            task_masks = []
            flag = 0

            # 在class group里面的class_name是一个name list
            # 每一个group内每一个class都要一个mask
            # task_masks:[[m1,m2,m3,m4], [m1,m2,m3,m4]..]
            for class_name in class_names_by_task:
                # print("classes: ", gt_dict["gt_classes"], "name", class_name)
                task_masks.append([
                    np.where(gt_dict["gt_classes"] == class_name.index(i) + 1 +
                             flag) for i in class_name
                ])
                flag += len(class_name)

            task_boxes = []
            task_classes = []
            task_names = []
            flag2 = 0
            for idx, mask in enumerate(task_masks):
                task_box = []
                task_class = []
                task_name = []
                for m in mask:
                    task_box.append(gt_dict["gt_boxes"][m])
                    task_class.append(gt_dict["gt_classes"][m] - flag2)
                    task_name.append(gt_dict["gt_names"][m])
                task_boxes.append(np.concatenate(task_box, axis=0))
                task_classes.append(np.concatenate(task_class))
                task_names.append(np.concatenate(task_name))
                flag2 += len(mask)

            for task_box in task_boxes:
                # limit rad to [-pi, pi]
                task_box[:, -1] = box_np_ops.limit_period(task_box[:, -1],
                                                          offset=0.5,
                                                          period=np.pi * 2)

            # print(gt_dict.keys())
            gt_dict["gt_classes"] = task_classes
            gt_dict["gt_names"] = task_names
            gt_dict["gt_boxes"] = task_boxes

            res["lidar"]["annotations"] = gt_dict

            draw_gaussian = draw_umich_gaussian

            hms, anno_boxs, inds, masks, cats = [], [], [], [], []

            for idx, task in enumerate(self.tasks):
                hm = np.zeros((len(class_names_by_task[idx]),
                               feature_map_size[1], feature_map_size[0]),
                              dtype=np.float32)

                if res['type'] == 'NuScenesDataset':
                    # [reg, hei, dim, vx, vy, rots, rotc]
                    anno_box = np.zeros((max_objs, 10), dtype=np.float32)
                else:
                    raise NotImplementedError("Only Support nuScene for Now!")

                ind = np.zeros((max_objs), dtype=np.int64)
                mask = np.zeros((max_objs), dtype=np.uint8)
                cat = np.zeros((max_objs), dtype=np.int64)
                direction = np.zeros((max_objs), dtype=np.int64)

                # 统计gtbox个数
                num_objs = min(gt_dict['gt_boxes'][idx].shape[0], max_objs)

                for k in range(num_objs):
                    cls_id = gt_dict['gt_classes'][idx][k] - 1

                    w, l, h = gt_dict['gt_boxes'][idx][k][3], gt_dict['gt_boxes'][idx][k][4], \
                              gt_dict['gt_boxes'][idx][k][5]
                    w, l = w / voxel_size[
                        0] / self.out_size_factor, l / voxel_size[
                            1] / self.out_size_factor
                    if w > 0 and l > 0:
                        radius = gaussian_radius(
                            (l, w), min_overlap=self.gaussian_overlap)
                        radius = max(self._min_radius, int(radius))

                        # be really careful for the coordinate system of your box annotation.
                        x, y, z = gt_dict['gt_boxes'][idx][k][0], gt_dict['gt_boxes'][idx][k][1], \
                                  gt_dict['gt_boxes'][idx][k][2]

                        coor_x, coor_y = (x - pc_range[0]) / voxel_size[0] / self.out_size_factor, \
                                         (y - pc_range[1]) / voxel_size[1] / self.out_size_factor

                        ct = np.array([coor_x, coor_y], dtype=np.float32)
                        ct_int = ct.astype(np.int32)

                        # throw out not in range objects to avoid out of array area when creating the heatmap
                        if not (0 <= ct_int[0] < feature_map_size[0]
                                and 0 <= ct_int[1] < feature_map_size[1]):
                            continue

                        draw_gaussian(hm[cls_id], ct, radius)

                        new_idx = k
                        x, y = ct_int[0], ct_int[1]

                        if not (y * feature_map_size[0] + x <
                                feature_map_size[0] * feature_map_size[1]):
                            # a double check, should never happen
                            print(x, y, y * feature_map_size[0] + x)
                            assert False

                        cat[new_idx] = cls_id
                        ind[new_idx] = y * feature_map_size[0] + x
                        mask[new_idx] = 1

                        if res['type'] == 'NuScenesDataset':
                            vx, vy = gt_dict['gt_boxes'][idx][k][6:8]
                            rot = gt_dict['gt_boxes'][idx][k][8]
                            if not self.no_log:
                                anno_box[new_idx] = np.concatenate(
                                    (ct - (x, y), z,
                                     np.log(gt_dict['gt_boxes'][idx][k][3:6]),
                                     np.array(vx), np.array(vy), np.sin(rot),
                                     np.cos(rot)),
                                    axis=None)
                            else:
                                anno_box[new_idx] = np.concatenate(
                                    (ct - (x, y), z,
                                     gt_dict['gt_boxes'][idx][k][3:6],
                                     np.array(vx), np.array(vy), np.sin(rot),
                                     np.cos(rot)),
                                    axis=None)

                        else:
                            raise NotImplementedError(
                                "Only Support KITTI and nuScene for Now!")

                hms.append(hm)
                anno_boxs.append(anno_box)
                masks.append(mask)
                inds.append(ind)
                cats.append(cat)

            example.update({
                'hm': hms,
                'anno_box': anno_boxs,
                'ind': inds,
                'mask': masks,
                'cat': cats
            })
        else:
            pass

        res["lidar"]["targets"] = example

        return res, info
Exemple #5
0
    def __call__(self, res, info):

        class_names_by_task = [t.classes for t in self.target_assigners]

        # Calculate output featuremap size
        grid_size = res["lidar"]["voxels"]["shape"]
        feature_map_size = grid_size[:2] // self.out_size_factor
        feature_map_size = [*feature_map_size, 1][::-1]

        anchors_by_task = [
            t.generate_anchors(feature_map_size) for t in self.target_assigners
        ]
        anchor_dicts_by_task = [
            t.generate_anchors_dict(feature_map_size)
            for t in self.target_assigners
        ]
        reshaped_anchors_by_task = [
            t["anchors"].reshape([-1, t["anchors"].shape[-1]])
            for t in anchors_by_task
        ]
        matched_by_task = [t["matched_thresholds"] for t in anchors_by_task]
        unmatched_by_task = [
            t["unmatched_thresholds"] for t in anchors_by_task
        ]

        bv_anchors_by_task = [
            box_np_ops.rbbox2d_to_near_bbox(anchors[:, [0, 1, 3, 4, -1]])
            for anchors in reshaped_anchors_by_task
        ]

        anchor_caches_by_task = dict(
            anchors=reshaped_anchors_by_task,
            anchors_bv=bv_anchors_by_task,
            matched_thresholds=matched_by_task,
            unmatched_thresholds=unmatched_by_task,
            anchors_dict=anchor_dicts_by_task,
        )

        if res["mode"] == "train":
            gt_dict = res["lidar"]["annotations"]

            task_masks = []
            flag = 0
            for class_name in class_names_by_task:
                task_masks.append([
                    np.where(gt_dict["gt_classes"] == class_name.index(i) + 1 +
                             flag) for i in class_name
                ])
                flag += len(class_name)

            task_boxes = []
            task_classes = []
            task_names = []
            flag2 = 0
            for idx, mask in enumerate(task_masks):
                task_box = []
                task_class = []
                task_name = []
                for m in mask:
                    task_box.append(gt_dict["gt_boxes"][m])
                    task_class.append(gt_dict["gt_classes"][m] - flag2)
                    task_name.append(gt_dict["gt_names"][m])
                task_boxes.append(np.concatenate(task_box, axis=0))
                task_classes.append(np.concatenate(task_class))
                task_names.append(np.concatenate(task_name))
                flag2 += len(mask)

            for task_box in task_boxes:
                # limit rad to [-pi, pi]
                task_box[:, -1] = box_np_ops.limit_period(task_box[:, -1],
                                                          offset=0.5,
                                                          period=np.pi * 2)

            # print(gt_dict.keys())
            gt_dict["gt_classes"] = task_classes
            gt_dict["gt_names"] = task_names
            gt_dict["gt_boxes"] = task_boxes

            res["lidar"]["annotations"] = gt_dict

        anchorss = anchor_caches_by_task["anchors"]
        anchors_bvs = anchor_caches_by_task["anchors_bv"]
        anchors_dicts = anchor_caches_by_task["anchors_dict"]

        example = {}
        example["anchors"] = anchorss

        if self.anchor_area_threshold >= 0:
            example["anchors_mask"] = []
            for idx, anchors_bv in enumerate(anchors_bvs):
                anchors_mask = None
                # 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"].append(anchors_mask)

        if res["mode"] == "train":
            targets_dicts = []
            for idx, target_assigner in enumerate(self.target_assigners):
                if "anchors_mask" in example:
                    anchors_mask = example["anchors_mask"][idx]
                else:
                    anchors_mask = None
                targets_dict = target_assigner.assign_v2(
                    anchors_dicts[idx],
                    gt_dict["gt_boxes"][idx],
                    anchors_mask,
                    gt_classes=gt_dict["gt_classes"][idx],
                    gt_names=gt_dict["gt_names"][idx],
                )
                targets_dicts.append(targets_dict)

            example.update({
                "labels":
                [targets_dict["labels"] for targets_dict in targets_dicts],
                "reg_targets": [
                    targets_dict["bbox_targets"]
                    for targets_dict in targets_dicts
                ],
                "reg_weights": [
                    targets_dict["bbox_outside_weights"]
                    for targets_dict in targets_dicts
                ],
            })

        res["lidar"]["targets"] = example

        return res, info
Exemple #6
0
    def __call__(self, res, info):
        targets, targets_raw = {}, {}
        targets["anchors"] = [
            anchor_dict[self.target_class_names[i]]["anchors"].reshape([-1, 7])
            for i, anchor_dict in enumerate(self.anchor_dicts_by_task)
        ]
        targets_raw["anchors"] = targets["anchors"].copy()

        # get gt labels of targeted classes; limit ry range in [-pi, pi].
        if res["mode"] == "train" and res['labeled']:
            gt_dict = res["lidar"]["annotations"]
            gt_mask = np.zeros(gt_dict["gt_classes"].shape, dtype=np.bool)
            for target_class_id in self.target_class_ids:
                gt_mask = np.logical_or(
                    gt_mask, gt_dict["gt_classes"] == target_class_id)

            gt_boxes = gt_dict["gt_boxes"][gt_mask]
            gt_boxes[:,
                     -1] = box_np_ops.limit_period(gt_boxes[:, -1],
                                                   offset=0.5,
                                                   period=np.pi *
                                                   2)  # limit ry to [-pi, pi]
            gt_dict["gt_boxes"] = [gt_boxes]
            gt_dict["gt_classes"] = [gt_dict["gt_classes"][gt_mask]]
            gt_dict["gt_names"] = [gt_dict["gt_names"][gt_mask]]
            res["lidar"]["annotations"] = gt_dict

            targets_dict = {}
            for idx, target_assigner in enumerate(self.target_assigners):
                targets_dict = target_assigner.assign_v2(
                    self.anchor_dicts_by_task[idx],
                    gt_dict["gt_boxes"][idx],  # (x, y, z, w, l, h, r)
                    anchors_mask=None,
                    gt_classes=gt_dict["gt_classes"][idx],
                    gt_names=gt_dict["gt_names"][idx],
                    enable_similar_type=self.enable_similar_type,
                )

            targets.update({
                "labels": [targets_dict["labels"]],
                "reg_targets": [targets_dict["bbox_targets"]],
                "reg_weights": [targets_dict["bbox_outside_weights"]],
                "positive_gt_id": [targets_dict["positive_gt_id"]],
            })

            ################# for raw points labels [unused] #######################
            gt_dict_raw = res["lidar"]["annotations_raw"]
            gt_mask = np.zeros(gt_dict_raw["gt_classes"].shape, dtype=np.bool)
            for target_class_id in self.target_class_ids:
                gt_mask = np.logical_or(
                    gt_mask, gt_dict_raw["gt_classes"] == target_class_id)

            gt_boxes = gt_dict_raw["gt_boxes"][gt_mask]
            gt_boxes[:, -1] = box_np_ops.limit_period(gt_boxes[:, -1],
                                                      offset=0.5,
                                                      period=np.pi * 2)
            gt_dict_raw["gt_boxes"] = [gt_boxes]
            gt_dict_raw["gt_classes"] = [gt_dict_raw["gt_classes"][gt_mask]]
            gt_dict_raw["gt_names"] = [gt_dict_raw["gt_names"][gt_mask]]
            res["lidar"]["annotations_raw"] = gt_dict_raw

            targets_dict = {}
            for idx, target_assigner in enumerate(self.target_assigners):
                targets_dict = target_assigner.assign_v2(
                    self.anchor_dicts_by_task[idx],
                    gt_dict_raw["gt_boxes"][idx],
                    anchors_mask=None,
                    gt_classes=gt_dict_raw["gt_classes"][idx],
                    gt_names=gt_dict_raw["gt_names"][idx],
                    enable_similar_type=self.enable_similar_type,
                )

            targets_raw.update({
                "labels": [targets_dict["labels"]],
                "reg_targets": [targets_dict["bbox_targets"]],
                "reg_weights": [targets_dict["bbox_outside_weights"]],
                "positive_gt_id": [targets_dict["positive_gt_id"]],
            })
            ######################## end #########################

        res["lidar"]["targets"] = targets
        res["lidar"]["targets_raw"] = targets_raw
        return res, info
Exemple #7
0
    def __call__(self, res, info):

        # 需要检测的物体类别
        class_names_by_task = [t.classes for t in self.target_assigners]

        # Calculate output featuremap size
        grid_size = res["lidar"]["voxels"]["shape"]
        feature_map_size = grid_size[:2] // self.out_size_factor
        feature_map_size = [*feature_map_size, 1][::-1]    # z轴的size是1

        # 生成不同类的不同anchor
        # anchors list
        anchors_by_task = [
            t.generate_anchors(feature_map_size) for t in self.target_assigners
        ]
        # anchors 类别有dict
        anchor_dicts_by_task = [
            t.generate_anchors_dict(feature_map_size)
            for t in self.target_assigners
        ]
        # reshape乘n*7的形式
        reshaped_anchors_by_task = [
            t["anchors"].reshape([-1, t["anchors"].shape[-1]])
            for t in anchors_by_task
        ]
        matched_by_task = [t["matched_thresholds"] for t in anchors_by_task]
        unmatched_by_task = [
            t["unmatched_thresholds"] for t in anchors_by_task
        ]

        # bev下的anchor
        bv_anchors_by_task = [
            box_np_ops.rbbox2d_to_near_bbox(anchors[:, [0, 1, 3, 4, -1]])
            for anchors in reshaped_anchors_by_task
        ]

        # 整合为字典
        anchor_caches_by_task = dict(
            anchors=reshaped_anchors_by_task,
            anchors_bv=bv_anchors_by_task,
            matched_thresholds=matched_by_task,
            unmatched_thresholds=unmatched_by_task,
            anchors_dict=anchor_dicts_by_task,
        )

        if res["mode"] == "train":
            # 得到当前帧的标注
            gt_dict = res["lidar"]["annotations"]

            task_masks = []
            flag = 0
            # 不同任务的不同类gt,得到mask,task_id:[]
            for class_name in class_names_by_task:
                task_masks.append([
                    np.where(gt_dict["gt_classes"] == class_name.index(i) + 1 +
                             flag) for i in class_name
                ])
                flag += len(class_name)

            # task存储
            task_boxes = []
            task_boxes_1 = []
            task_boxes_2 = []
            task_classes = []
            task_names = []
            flag2 = 0
            for idx, mask in enumerate(task_masks):
                task_box = []
                task_box_1 = []
                task_box_2 = []
                task_class = []
                task_name = []
                for m in mask:
                    # mask 应用
                    # TODO: 在这里应用mask时需要把前后两帧的label加入
                    task_box.append(gt_dict["gt_boxes"][m])
                    task_box_1.append(gt_dict["gt_boxes_1"][m])
                    task_box_2.append(gt_dict["gt_boxes_2"][m])
                    task_class.append(gt_dict["gt_classes"][m] - flag2)
                    task_name.append(gt_dict["gt_names"][m])
                task_boxes.append(np.concatenate(task_box, axis=0))
                task_boxes_1.append(np.concatenate(task_box_1, axis=0))
                task_boxes_2.append(np.concatenate(task_box_2, axis=0))
                task_classes.append(np.concatenate(task_class))
                task_names.append(np.concatenate(task_name))
                flag2 += len(mask)

            for task_box, task_box_1, task_box_2 in zip(
                    task_boxes, task_boxes_1, task_boxes_2):
                # TODO:将朝向角限制在-pi~pi
                # limit rad to [-pi, pi]
                task_box[:, -1] = box_np_ops.limit_period(task_box[:, -1],
                                                          offset=0.5,
                                                          period=np.pi * 2)
                task_box_1[:, -1] = box_np_ops.limit_period(task_box_1[:, -1],
                                                            offset=0.5,
                                                            period=np.pi * 2)
                task_box_2[:, -1] = box_np_ops.limit_period(task_box_2[:, -1],
                                                            offset=0.5,
                                                            period=np.pi * 2)

            # print(gt_dict.keys())
            # 对应的gt_boxes_1
            gt_dict["gt_classes"] = task_classes
            gt_dict["gt_names"] = task_names
            gt_dict["gt_boxes"] = task_boxes
            gt_dict["gt_boxes_1"] = task_boxes_1
            gt_dict["gt_boxes_2"] = task_boxes_2

            res["lidar"]["annotations"] = gt_dict

        anchorss = anchor_caches_by_task["anchors"]
        anchors_bvs = anchor_caches_by_task["anchors_bv"]
        anchors_dicts = anchor_caches_by_task["anchors_dict"]

        example = {}
        # anchor是reshape出来的D*H*W*7
        example["anchors"] = anchorss

        if self.anchor_area_threshold >= 0:
            # 暂时没有用到
            example["anchors_mask"] = []
            for idx, anchors_bv in enumerate(anchors_bvs):
                anchors_mask = None
                # 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"].append(anchors_mask)

        if res["mode"] == "train":
            targets_dicts = []
            # 主体部分,利用anchors和gt_boxes得到target
            # TODO: 修改的核心部分
            for idx, target_assigner in enumerate(self.target_assigners):
                if "anchors_mask" in example:
                    anchors_mask = example["anchors_mask"][idx]
                else:
                    anchors_mask = None
                targets_dict = target_assigner.assign_v3(
                    anchors_dicts[idx],
                    gt_dict["gt_boxes"][idx],
                    gt_dict["gt_boxes_1"][idx],
                    gt_dict["gt_boxes_2"][idx],
                    anchors_mask,
                    gt_classes=gt_dict["gt_classes"][idx],
                    gt_names=gt_dict["gt_names"][idx],
                )
                targets_dicts.append(targets_dict)

            example.update({
                "labels":
                [targets_dict["labels"] for targets_dict in targets_dicts],
                "reg_targets": [
                    targets_dict["bbox_targets"]
                    for targets_dict in targets_dicts
                ],
                "reg_targets_1": [
                    targets_dict["bbox_targets_1"]
                    for targets_dict in targets_dicts
                ],
                "reg_targets_2": [
                    targets_dict["bbox_targets_2"]
                    for targets_dict in targets_dicts
                ],
                "reg_weights": [
                    targets_dict["bbox_outside_weights"]
                    for targets_dict in targets_dicts
                ],
            })

        res["lidar"]["targets"] = example

        return res, info