Esempio n. 1
0
def test_imshow_bbox():
    img = np.zeros((100, 100, 3), dtype=np.uint8)
    bboxes = np.array([[10, 10, 30, 30], [10, 50, 30, 80]], dtype=np.float32)
    labels = ['label 1', 'label 2']
    colors = ['red', 'green']

    with tempfile.TemporaryDirectory() as tmpdir:
        _ = imshow_bboxes(img,
                          bboxes,
                          labels=labels,
                          colors=colors,
                          show=False,
                          out_file=f'{tmpdir}/out.png')

        # test case of empty bboxes
        _ = imshow_bboxes(img,
                          np.zeros((0, 4), dtype=np.float32),
                          labels=None,
                          colors='red',
                          show=False)

        # test unmatched bboxes and labels
        with pytest.raises(AssertionError):
            _ = imshow_bboxes(img,
                              np.zeros((0, 4), dtype=np.float32),
                              labels=labels[:1],
                              colors='red',
                              show=False)
    def _draw_bbox(self, canvas: np.ndarray, input_msg: FrameMessage):
        """Draw object bboxes."""

        if self.must_have_keypoint:
            objects = input_msg.get_objects(
                lambda x: 'bbox' in x and 'keypoints' in x)
        else:
            objects = input_msg.get_objects(lambda x: 'bbox' in x)
        # return if there is no detected objects
        if not objects:
            return canvas

        bboxes = [obj['bbox'] for obj in objects]
        labels = [obj.get('label', None) for obj in objects]
        default_color = (0, 255, 0)

        # Get bbox colors
        if isinstance(self.bbox_color, dict):
            colors = [
                self.bbox_color.get(label, default_color) for label in labels
            ]
        else:
            colors = self.bbox_color

        imshow_bboxes(canvas,
                      np.vstack(bboxes),
                      labels=labels,
                      colors=colors,
                      text_color='white',
                      font_scale=0.5,
                      show=False)

        return canvas
Esempio n. 3
0
    def show_result(self,
                    img,
                    result,
                    skeleton=None,
                    kpt_score_thr=0.3,
                    bbox_color='green',
                    pose_kpt_color=None,
                    pose_link_color=None,
                    text_color='white',
                    radius=4,
                    thickness=1,
                    font_scale=0.5,
                    bbox_thickness=1,
                    win_name='',
                    show=False,
                    show_keypoint_weight=False,
                    wait_time=0,
                    out_file=None):
        """Draw `result` over `img`.

        Args:
            img (str or Tensor): The image to be displayed.
            result (list[dict]): The results to draw over `img`
                (bbox_result, pose_result).
            skeleton (list[list]): The connection of keypoints.
                skeleton is 0-based indexing.
            kpt_score_thr (float, optional): Minimum score of keypoints
                to be shown. Default: 0.3.
            bbox_color (str or tuple or :obj:`Color`): Color of bbox lines.
            pose_kpt_color (np.array[Nx3]`): Color of N keypoints.
                If None, do not draw keypoints.
            pose_link_color (np.array[Mx3]): Color of M links.
                If None, do not draw links.
            text_color (str or tuple or :obj:`Color`): Color of texts.
            radius (int): Radius of circles.
            thickness (int): Thickness of lines.
            font_scale (float): Font scales of texts.
            win_name (str): The window name.
            show (bool): Whether to show the image. Default: False.
            show_keypoint_weight (bool): Whether to change the transparency
                using the predicted confidence scores of keypoints.
            wait_time (int): Value of waitKey param.
                Default: 0.
            out_file (str or None): The filename to write the image.
                Default: None.

        Returns:
            Tensor: Visualized img, only if not `show` or `out_file`.
        """
        img = mmcv.imread(img)
        img = img.copy()

        bbox_result = []
        bbox_labels = []
        pose_result = []
        for res in result:
            if 'bbox' in res:
                bbox_result.append(res['bbox'])
                bbox_labels.append(res.get('label', None))
            pose_result.append(res['keypoints'])

        if bbox_result:
            bboxes = np.vstack(bbox_result)
            # draw bounding boxes
            imshow_bboxes(img,
                          bboxes,
                          labels=bbox_labels,
                          colors=bbox_color,
                          text_color=text_color,
                          thickness=bbox_thickness,
                          font_scale=font_scale,
                          show=False)

        if pose_result:
            imshow_keypoints(img, pose_result, skeleton, kpt_score_thr,
                             pose_kpt_color, pose_link_color, radius,
                             thickness)

        if show:
            imshow(img, win_name, wait_time)

        if out_file is not None:
            imwrite(img, out_file)

        return img
Esempio n. 4
0
    def show_result(self,
                    result,
                    img=None,
                    skeleton=None,
                    pose_kpt_color=None,
                    pose_link_color=None,
                    radius=8,
                    thickness=2,
                    vis_height=400,
                    num_instances=-1,
                    axis_azimuth=70,
                    win_name='',
                    show=False,
                    wait_time=0,
                    out_file=None):
        """Visualize 3D pose estimation results.

        Args:
            result (list[dict]): The pose estimation results containing:

                - "keypoints_3d" ([K,4]): 3D keypoints
                - "keypoints" ([K,3] or [T,K,3]): Optional for visualizing
                    2D inputs. If a sequence is given, only the last frame
                    will be used for visualization
                - "bbox" ([4,] or [T,4]): Optional for visualizing 2D inputs
                - "title" (str): title for the subplot
            img (str or Tensor): Optional. The image to visualize 2D inputs on.
            skeleton (list of [idx_i,idx_j]): Skeleton described by a list of
                links, each is a pair of joint indices.
            pose_kpt_color (np.array[Nx3]`): Color of N keypoints.
                If None, do not draw keypoints.
            pose_link_color (np.array[Mx3]): Color of M links.
                If None, do not draw links.
            radius (int): Radius of circles.
            thickness (int): Thickness of lines.
            vis_height (int): The image height of the visualization. The width
                will be N*vis_height depending on the number of visualized
                items.
            num_instances (int): Number of instances to be shown in 3D. If
                smaller than 0, all the instances in the result will be shown.
                Otherwise, pad or truncate the result to a length of
                num_instances.
            axis_azimuth (float): axis azimuth angle for 3D visualizations.
            win_name (str): The window name.
            show (bool): Whether to directly show the visualization.
            wait_time (int): Value of waitKey param.
                Default: 0.
            out_file (str or None): The filename to write the image.
                Default: None.

        Returns:
            Tensor: Visualized img, only if not `show` or `out_file`.
        """
        if num_instances < 0:
            assert len(result) > 0
        result = sorted(result, key=lambda x: x.get('track_id', 1e4))

        # draw image and input 2d poses
        if img is not None:
            img = mmcv.imread(img)

            bbox_result = []
            pose_input_2d = []
            for res in result:
                if 'bbox' in res:
                    bbox = np.array(res['bbox'])
                    if bbox.ndim != 1:
                        assert bbox.ndim == 2
                        bbox = bbox[-1]  # Get bbox from the last frame
                    bbox_result.append(bbox)
                if 'keypoints' in res:
                    kpts = np.array(res['keypoints'])
                    if kpts.ndim != 2:
                        assert kpts.ndim == 3
                        kpts = kpts[-1]  # Get 2D keypoints from the last frame
                    pose_input_2d.append(kpts)

            if len(bbox_result) > 0:
                bboxes = np.vstack(bbox_result)
                imshow_bboxes(
                    img,
                    bboxes,
                    colors='green',
                    thickness=thickness,
                    show=False)
            if len(pose_input_2d) > 0:
                imshow_keypoints(
                    img,
                    pose_input_2d,
                    skeleton,
                    kpt_score_thr=0.3,
                    pose_kpt_color=pose_kpt_color,
                    pose_link_color=pose_link_color,
                    radius=radius,
                    thickness=thickness)
            img = mmcv.imrescale(img, scale=vis_height / img.shape[0])

        img_vis = imshow_keypoints_3d(
            result,
            img,
            skeleton,
            pose_kpt_color,
            pose_link_color,
            vis_height,
            num_instances=num_instances,
            axis_azimuth=axis_azimuth,
        )

        if show:
            mmcv.visualization.imshow(img_vis, win_name, wait_time)

        if out_file is not None:
            mmcv.imwrite(img_vis, out_file)

        return img_vis