vkit.element.opt

  1# Copyright 2022 vkit-x Administrator. All Rights Reserved.
  2#
  3# This project (vkit-x/vkit) is dual-licensed under commercial and SSPL licenses.
  4#
  5# The commercial license gives you the full rights to create and distribute software
  6# on your own terms without any SSPL license obligations. For more information,
  7# please see the "LICENSE_COMMERCIAL.txt" file.
  8#
  9# This project is also available under Server Side Public License (SSPL).
 10# The SSPL licensing is ideal for use cases such as open source projects with
 11# SSPL distribution, student/academic purposes, hobby projects, internal research
 12# projects without external distribution, or other projects where all SSPL
 13# obligations can be met. For more information, please see the "LICENSE_SSPL.txt" file.
 14from typing import Union, Tuple, Optional, TypeVar
 15
 16import numpy as np
 17
 18from .type import Shapable
 19
 20_T = TypeVar('_T', float, int)
 21
 22
 23def clip_val(val: _T, size: int) -> _T:
 24    # Hack to work.
 25    return max(0, min(val, size - 1))  # type: ignore
 26
 27
 28def resize_val(val: _T, size: int, resized_size: int):
 29    return clip_val(val * resized_size / size, resized_size)
 30
 31
 32def extract_shape_from_shapable_or_shape(shapable_or_shape: Union[Shapable, Tuple[int, int]],):
 33    if isinstance(shapable_or_shape, Shapable):
 34        shapable = shapable_or_shape
 35        height, width = shapable.shape
 36    else:
 37        height, width = shapable_or_shape
 38    return height, width
 39
 40
 41def generate_resized_shape(
 42    height: int,
 43    width: int,
 44    resized_height: Optional[int] = None,
 45    resized_width: Optional[int] = None,
 46):
 47    if not resized_height and not resized_width:
 48        raise RuntimeError('Missing resized_height or resized_width.')
 49
 50    if resized_height is None:
 51        assert resized_width
 52        resized_height = round(resized_width * height / width)
 53
 54    if resized_width is None:
 55        assert resized_height
 56        resized_width = round(resized_height * width / height)
 57
 58    return resized_height, resized_width
 59
 60
 61def generate_shape_and_resized_shape(
 62    shapable_or_shape: Union[Shapable, Tuple[int, int]],
 63    resized_height: Optional[int] = None,
 64    resized_width: Optional[int] = None,
 65):
 66    height, width = extract_shape_from_shapable_or_shape(shapable_or_shape)
 67    resized_height, resized_width = generate_resized_shape(
 68        height=height,
 69        width=width,
 70        resized_height=resized_height,
 71        resized_width=resized_width,
 72    )
 73    return (
 74        height,
 75        width,
 76        resized_height,
 77        resized_width,
 78    )
 79
 80
 81def expand_np_mask(mat: np.ndarray, np_mask: np.ndarray):
 82    if mat.ndim == 2:
 83        # Do nothing.
 84        pass
 85
 86    elif mat.ndim == 3:
 87        num_channels = mat.shape[2]
 88        np_mask = np.repeat(np.expand_dims(np_mask, axis=-1), num_channels, axis=-1)
 89
 90    else:
 91        raise NotImplementedError()
 92
 93    return np_mask
 94
 95
 96def prep_value(
 97    mat: np.ndarray,
 98    value: Union[np.ndarray, Tuple[float, ...], float],
 99):
100    if not isinstance(value, np.ndarray):
101        if mat.ndim == 3:
102            num_channels = mat.shape[2]
103            if isinstance(value, tuple) and len(value) != num_channels:
104                raise RuntimeError('value is tuple but len(value) != num_channels.')
105
106        value = np.full_like(mat, value)
107
108    else:
109        if mat.shape != value.shape:
110            raise RuntimeError('value is np.ndarray but shape is not matched.')
111
112        if value.dtype != mat.dtype:
113            value = value.astype(mat.dtype)
114
115    return value
116
117
118def fill_np_array(
119    mat: np.ndarray,
120    value: Union[np.ndarray, Tuple[float, ...], float],
121    np_mask: Optional[np.ndarray] = None,
122    alpha: Union[np.ndarray, float] = 1.0,
123    keep_max_value: bool = False,
124    keep_min_value: bool = False,
125):
126    if not isinstance(value, np.ndarray) \
127            and np_mask is not None \
128            and isinstance(alpha, float) \
129            and alpha == 1.0 \
130            and not keep_max_value \
131            and not keep_min_value:
132        # Most common case, hence optimize the performance here.
133        mat[np_mask] = value
134        return
135
136    mat_origin = mat
137    np_value = prep_value(mat, value)
138
139    if np_mask is not None:
140        # NOTE: Boolean indexing makes a copy.
141        mat = mat[np_mask]
142        np_value = np_value[np_mask]
143        if isinstance(alpha, np.ndarray):
144            alpha = alpha[np_mask]
145
146    if isinstance(alpha, float):
147        if alpha < 0.0 or alpha > 1.0:
148            raise RuntimeError(f'alpha={alpha} is invalid.')
149
150        elif alpha == 0.0:
151            return
152
153        elif alpha == 1.0:
154            if np_mask is None and not (keep_max_value or keep_min_value):
155                np.copyto(mat_origin, np_value)
156
157            else:
158                if keep_max_value or keep_min_value:
159                    assert not (keep_max_value and keep_min_value)
160                    if keep_max_value:
161                        np_to_value_mask = (mat < np_value)
162                    else:
163                        np_to_value_mask = (mat > np_value)
164                    np.putmask(mat, np_to_value_mask, np_value)
165                else:
166                    mat = np_value
167
168                if np_mask is not None:
169                    mat_origin[np_mask] = mat
170
171        else:
172            # 0 < alpha < 1.
173            np_value_weight = np.full(
174                (mat.shape[0], mat.shape[1]),
175                alpha,
176                dtype=np.float32,
177            )
178            if np_value_weight.shape != mat.shape:
179                assert np_value_weight.ndim + 1 == mat.ndim
180                np_value_weight = np.expand_dims(np_value_weight, -1)
181
182            np_mat_weight = 1 - np_value_weight
183            np_weighted_sum = (
184                np_mat_weight * mat.astype(np.float32)
185                + np_value_weight * np_value.astype(np.float32)
186            )
187            np_weighted_sum = np_weighted_sum.astype(mat.dtype)
188
189            if np_mask is not None:
190                mat_origin[np_mask] = np_weighted_sum
191            else:
192                np.copyto(mat_origin, np_weighted_sum)
193
194    else:
195        np_value_weight = alpha
196        if np_value_weight.shape != mat.shape:
197            assert np_value_weight.ndim + 1 == mat.ndim
198            np_value_weight = np.expand_dims(np_value_weight, -1)
199
200        np_mat_weight = 1 - np_value_weight
201        np_weighted_sum = (
202            np_mat_weight * mat.astype(np.float32) + np_value_weight * np_value.astype(np.float32)
203        )
204        np_weighted_sum = np_weighted_sum.astype(mat.dtype)
205
206        if np_mask is not None:
207            mat_origin[np_mask] = np_weighted_sum
208        else:
209            np.copyto(mat_origin, np_weighted_sum)