Ejemplo n.º 1
0
def kitti_bbox2results(boxes_lidar, scores, labels, meta, class_names=None):
    calib = meta['calib']
    sample_id = meta['sample_idx']
    image_shape = meta['img_shape'][: 2]
    if scores is None \
        or len(scores) == 0 \
            or boxes_lidar is None \
                or len(boxes_lidar) == 0:

        return kitti.empty_result_anno()

    hehe = kitti.get_start_result_anno()
    hehe.update({'image_idx': []})

    boxes_lidar[:, -1] = limit_period(
        boxes_lidar[:, -1], offset=0.5, period=np.pi * 2,
    )
    boxes_cam = np.zeros_like(boxes_lidar)
    boxes_cam[:, :3] = project_velo_to_rect(boxes_lidar[:, :3], calib)
    boxes_cam[:, 3:] = boxes_lidar[:, [4, 5, 3, 6]]
    corners_cam = center_to_corner_box3d(boxes_cam, origin=[0.5, 1.0, 0.5], axis=1)
    corners_rgb = project_rect_to_image(corners_cam, calib)
    minxy = np.min(corners_rgb, axis=1)
    maxxy = np.max(corners_rgb, axis=1)
    box2d_rgb = np.concatenate([minxy, maxxy], axis=1)
    alphas = -np.arctan2(-boxes_lidar[:, 1], boxes_lidar[:, 0]) + boxes_lidar[:, 6]

    for i, (lb, score, box3d, box2d, alpha) in enumerate(zip(labels, scores, boxes_cam, box2d_rgb, alphas)):
        if box2d[0] > image_shape[1] or box2d[1] > image_shape[0]:
            continue
        if box2d[2] < 0 or box2d[3] < 0:
            continue
        box2d[2:] = np.minimum(box2d[2:], image_shape[::-1])
        box2d[:2] = np.maximum(box2d[:2], [0, 0])

        hehe["name"].append(class_names[lb])
        hehe["truncated"].append(0.0)
        hehe["occluded"].append(0)
        hehe["alpha"].append(alpha)
        hehe["bbox"].append(box2d)
        hehe["dimensions"].append(box3d[[3, 4, 5]])
        hehe["location"].append(box3d[:3])
        hehe["rotation_y"].append(box3d[6])
        hehe["score"].append(score)
        hehe['image_idx'].append(int(sample_id))

    try:
        hehe = {n: np.stack(v) for n, v in hehe.items()}
    except:
        return kitti.empty_result_anno()

    return hehe
Ejemplo n.º 2
0
def single_test(model, data_loader, saveto=None, class_names=['Car']):
    template = '{} ' + ' '.join(['{:.4f}' for _ in range(15)]) + '\n'
    if saveto is not None:
        mmcv.mkdir_or_exist(saveto)

    model.eval()
    annos = []
    prog_bar = mmcv.ProgressBar(len(data_loader.dataset))

    for i, data in enumerate(data_loader):
        with torch.no_grad():
            results = model(return_loss=False, **data)
        image_shape = (375,1242)
        for re in results:
            img_idx = re['image_idx']
            if re['bbox'] is not None:
                box2d = re['bbox']
                box3d = re['box3d_camera']
                labels = re['label_preds']
                scores = re['scores']
                alphas = re['alphas']
                anno = kitti.get_start_result_anno()
                num_example = 0
                for bbox2d, bbox3d, label, score, alpha in zip(box2d, box3d, labels, scores, alphas):
                    if bbox2d[0] > image_shape[1] or bbox2d[1] > image_shape[0]:
                        continue
                    if bbox2d[2] < 0 or bbox2d[3] < 0:
                        continue
                    bbox2d[2:] = np.minimum(bbox2d[2:], image_shape[::-1])
                    bbox2d[:2] = np.maximum(bbox2d[:2], [0, 0])
                    anno["name"].append(class_names[int(label)])
                    anno["truncated"].append(0.0)
                    anno["occluded"].append(0)
                    # anno["alpha"].append(-10)
                    anno["alpha"].append(alpha)
                    anno["bbox"].append(bbox2d)
                    # anno["dimensions"].append(np.array([-1,-1,-1]))
                    anno["dimensions"].append(bbox3d[[3, 4, 5]])
                    # anno["location"].append(np.array([-1000,-1000,-1000]))
                    anno["location"].append(bbox3d[:3])
                    # anno["rotation_y"].append(-10)
                    anno["rotation_y"].append(bbox3d[6])
                    anno["score"].append(score)
                    num_example += 1
                if num_example != 0:
                    if saveto is not None:
                        of_path = os.path.join(saveto, '%06d.txt' % img_idx)
                        with open(of_path, 'w+') as f:
                            for name, bbox, dim, loc, ry, score, alpha in zip(anno['name'], anno["bbox"], \
                            anno["dimensions"], anno["location"], anno["rotation_y"], anno["score"],anno["alpha"]):
                                line = template.format(name, 0, 0, alpha, *bbox, *dim[[1,2,0]], *loc, ry, score)
                                f.write(line)

                    anno = {n: np.stack(v) for n, v in anno.items()}
                    annos.append(anno)
                else:
                    if saveto is not None:
                        of_path = os.path.join(saveto, '%06d.txt' % img_idx)
                        f = open(of_path, 'w+')
                        f.close()
                    annos.append(kitti.empty_result_anno())
            else:
                if saveto is not None:
                    of_path = os.path.join(saveto, '%06d.txt' % img_idx)
                    f = open(of_path, 'w+')
                    f.close()
                annos.append(kitti.empty_result_anno())

            num_example = annos[-1]["name"].shape[0]
            annos[-1]["image_idx"] = np.array(
                [img_idx] * num_example, dtype=np.int64)

        batch_size = len(results)
        for _ in range(batch_size):
            prog_bar.update()

    return annos
Ejemplo n.º 3
0
    def after_train_epoch(self, runner):
        if not self.every_n_epochs(runner, self.interval):
            return
        if runner.rank == 0:
            runner.model.eval()
            prog_bar = mmcv.ProgressBar(len(self.dataset))
            class_names = get_classes('kitti')
            results = []
            for i, data in enumerate(self.dataloader):
                # compute output
                with torch.no_grad():
                    result = runner.model(return_loss=False, **data)

                image_shape = (600, 1987)

                for re in result:
                    img_idx = re['image_idx']
                    if re['bbox'] is not None:
                        box2d = re['bbox']
                        box3d = re['box3d_camera']
                        labels = re['label_preds']
                        scores = re['scores']
                        alphas = re['alphas']

                        anno = kitti.get_start_result_anno()
                        num_example = 0
                        for bbox2d, bbox3d, label, score, alpha in zip(
                                box2d, box3d, labels, scores, alphas):

                            if bbox2d[0] > image_shape[1] or bbox2d[
                                    1] > image_shape[0]:
                                continue
                            if bbox2d[2] < 0 or bbox2d[3] < 0:
                                continue

                            bbox2d[2:] = np.minimum(bbox2d[2:],
                                                    image_shape[::-1])
                            bbox2d[:2] = np.maximum(bbox2d[:2], [0, 0])

                            anno["name"].append(class_names[label])
                            anno["truncated"].append(0.0)
                            anno["occluded"].append(0)
                            #anno["alpha"].append(-10)
                            anno["alpha"].append(alpha)
                            anno["bbox"].append(bbox2d)

                            #anno["dimensions"].append(np.array([-1,-1,-1]))
                            anno["dimensions"].append(bbox3d[[3, 4, 5]])
                            #anno["location"].append(np.array([-1000,-1000,-1000]))
                            anno["location"].append(bbox3d[:3])
                            #anno["rotation_y"].append(-10)
                            anno["rotation_y"].append(bbox3d[6])

                            anno["score"].append(score)
                            num_example += 1
                        if num_example != 0:
                            anno = {n: np.stack(v) for n, v in anno.items()}
                            results.append(anno)
                        else:
                            results.append(kitti.empty_result_anno())
                    else:
                        results.append(kitti.empty_result_anno())

                    num_example = results[-1]["name"].shape[0]
                    results[-1]["image_idx"] = np.array([img_idx] *
                                                        num_example,
                                                        dtype=np.int64)

                batch_size = len(data['sample_idx'])
                for _ in range(batch_size):
                    prog_bar.update()

            self._barrier(runner.rank, runner.world_size)
            self.evaluate(runner, results)
        else:
            self._barrier(runner.rank, runner.world_size)