vkit.mechanism.distortion_policy.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 Optional, Tuple
 15from enum import Enum, auto
 16
 17from numpy.random import Generator as RandomGenerator
 18
 19from vkit.utility import rng_choice_with_size
 20
 21LEVEL_MIN = 1
 22LEVEL_MAX = 10
 23CHANNELS = [0, 1, 2]
 24
 25
 26def sample_channels(rng: RandomGenerator):
 27    num_channels = rng.integers(1, 4)
 28    channels = None
 29    if num_channels < 3:
 30        channels = sorted(rng_choice_with_size(rng, CHANNELS, num_channels, replace=False))
 31    return channels
 32
 33
 34def sample_int(
 35    level: int,
 36    value_min: int,
 37    value_max: int,
 38    prob_negative: Optional[float],
 39    rng: RandomGenerator,
 40    inverse_level: bool = False,
 41):
 42    if inverse_level:
 43        level = LEVEL_MAX + 1 - level
 44
 45    value_range = value_max - value_min
 46    level_value_min = round(value_min + (level - 1) / LEVEL_MAX * value_range)
 47    level_value_max = round(value_min + level / LEVEL_MAX * value_range)
 48
 49    if level == LEVEL_MAX:
 50        # Make sure value_max could be sampled.
 51        level_value_max += 1
 52
 53    value = rng.integers(level_value_min, max(level_value_min + 1, level_value_max))
 54    if prob_negative and rng.random() < prob_negative:
 55        value *= -1
 56
 57    # NOTE: rng.integers returns numpy.int64 instead of int.
 58    # Some opencv function cannot handle numpy type, hence need to cast to int here.
 59    value = int(value)
 60
 61    return value
 62
 63
 64class SampleFloatMode(Enum):
 65    LINEAR = auto()
 66    QUAD = auto()
 67
 68
 69def func_quad(x: float):
 70    return -x**2 + 2 * x
 71
 72
 73def sample_float(
 74    level: int,
 75    value_min: float,
 76    value_max: float,
 77    prob_reciprocal: Optional[float],
 78    rng: RandomGenerator,
 79    mode: SampleFloatMode = SampleFloatMode.LINEAR,
 80    inverse_level: bool = False,
 81):
 82    if inverse_level:
 83        level = LEVEL_MAX + 1 - level
 84
 85    value_range = value_max - value_min
 86
 87    if mode == SampleFloatMode.LINEAR:
 88        level_ratio_min = (level - 1) / LEVEL_MAX
 89        level_ratio_max = level / LEVEL_MAX
 90
 91    elif mode == SampleFloatMode.QUAD:
 92        level_ratio_min = func_quad((level - 1) / LEVEL_MAX)
 93        level_ratio_max = func_quad(level / LEVEL_MAX)
 94
 95    else:
 96        raise NotImplementedError()
 97
 98    level_value_min = value_min + level_ratio_min * value_range
 99    level_value_max = value_min + level_ratio_max * value_range
100    value = rng.uniform(level_value_min, level_value_max)
101
102    if prob_reciprocal and rng.random() < prob_reciprocal:
103        value = 1 / value
104
105    return value
106
107
108def generate_grid_size(
109    grid_size_min: int,
110    grid_size_ratio: float,
111    shape: Tuple[int, int],
112):
113    return max(
114        grid_size_min,
115        int(grid_size_ratio * max(shape)),
116    )