File: depth.py

package info (click to toggle)
python-imgviz 1.2.4%2Bds-1
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 9,268 kB
  • sloc: python: 3,032; makefile: 15
file content (110 lines) | stat: -rw-r--r-- 2,730 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import matplotlib.cm
import numpy as np

from .normalize import normalize


class Depth2RGB(object):

    """Convert depth array to rgb.

    Parameters
    ----------
    min_value: float, optional
        Minimum value for colorizing.
    max_value: float, optional
        Maximum value for colorizing.
    colormap: str, optional
        Colormap, default: 'jet'.

    """

    def __init__(self, min_value=None, max_value=None, colormap="jet"):
        self._min_value = min_value
        self._max_value = max_value
        self._colormap = colormap

    @property
    def min_value(self):
        """Minimum value of depth."""
        return self._min_value

    @property
    def max_value(self):
        """Maximum value of depth."""
        return self._max_value

    def __call__(self, depth, dtype=np.uint8):
        """Convert depth array to rgb.

        Parameters
        ----------
        depth: numpy.ndarray, (H, W), float
            Depth image.
        dtype: numpy.dtype
            Dtype of output image. default: np.uint8

        Returns
        -------
        rgb: numpy.ndarray, (H, W, 3), np.uint8
            Output colorized image.

        """
        assert depth.ndim == 2, "depth image must be 2 dimensional"
        assert np.issubdtype(
            depth.dtype, np.floating
        ), "depth dtype must be float"

        normalized, self._min_value, self._max_value = normalize(
            depth,
            min_value=self._min_value,
            max_value=self._max_value,
            return_minmax=True,
        )

        isnan = np.isnan(normalized)
        normalized[isnan] = 0

        colormap_func = matplotlib.cm.get_cmap(self._colormap)
        rgb = colormap_func(normalized)[:, :, :3]
        rgb[isnan] = (0, 0, 0)

        if dtype == np.uint8:
            rgb = (rgb * 255).round().astype(np.uint8)
        else:
            assert np.issubdtype(dtype, np.floating)
            rgb = rgb.astype(dtype)

        return rgb


def depth2rgb(
    depth,
    dtype=np.uint8,
    min_value=None,
    max_value=None,
    colormap="jet",
):
    # type: (np.ndarray, np.dtype, float, float, str) -> np.ndarray
    """Convert depth to rgb.

    Parameters
    ----------
    depth: numpy.ndarray, (H, W), float
        Depth image.
    dtype: numpy.dtype
        Dtype of output image. default: np.uint8
    min_value: float, optional
        Minimum value for colorizing.
    max_value: float, optional
        Maximum value for colorizing.
    colormap: str, optional
        Colormap, default: 'jet'.

    Returns
    -------
    rgb: numpy.ndarray, (H, W, 3), np.uint8
        Output colorized image.

    """
    return Depth2RGB(min_value, max_value, colormap)(depth, dtype)