vkit.mechanism.distortion_policy.type

  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 (
 15    Mapping,
 16    Any,
 17    Generic,
 18    TypeVar,
 19    Type,
 20    Iterable,
 21    Union,
 22    Tuple,
 23    Optional,
 24)
 25
 26from numpy.random import Generator as RandomGenerator
 27
 28from vkit.utility import (
 29    dyn_structure,
 30    get_generic_classes,
 31    PathType,
 32)
 33from vkit.element import (
 34    Shapable,
 35    Image,
 36    Point,
 37    PointList,
 38    PointTuple,
 39    Polygon,
 40    Mask,
 41    ScoreMap,
 42)
 43from ..distortion.interface import DistortionConfig, DistortionState, Distortion
 44
 45_T_GENERATOR_CONFIG = TypeVar('_T_GENERATOR_CONFIG')
 46_T_CONFIG = TypeVar('_T_CONFIG', bound=DistortionConfig)
 47_T_STATE = TypeVar('_T_STATE', bound=DistortionState)
 48
 49
 50class DistortionConfigGenerator(Generic[_T_GENERATOR_CONFIG, _T_CONFIG]):
 51
 52    @classmethod
 53    def get_generator_config_cls(cls) -> Type[_T_GENERATOR_CONFIG]:
 54        return get_generic_classes(cls)[0]  # type: ignore
 55
 56    @classmethod
 57    def get_config_cls(cls) -> Type[_T_CONFIG]:
 58        return get_generic_classes(cls)[1]  # type: ignore
 59
 60    def __init__(self, config: _T_GENERATOR_CONFIG, level: int) -> None:
 61        self.config = config
 62        assert 1 <= level <= 10
 63        self.level = level
 64
 65    def __call__(self, shape: Tuple[int, int], rng: RandomGenerator) -> _T_CONFIG:
 66        raise NotImplementedError()
 67
 68
 69class DistortionPolicy(Generic[_T_GENERATOR_CONFIG, _T_CONFIG, _T_STATE]):
 70
 71    def __init__(
 72        self,
 73        distortion: Distortion[_T_CONFIG, _T_STATE],
 74        config_for_config_generator: _T_GENERATOR_CONFIG,
 75        config_generator_cls: Type[DistortionConfigGenerator[_T_GENERATOR_CONFIG, _T_CONFIG]],
 76    ):
 77        self.distortion = distortion
 78        self.config_for_config_generator = config_for_config_generator
 79        self.config_generator_cls = config_generator_cls
 80
 81    def distort(
 82        self,
 83        level: int,
 84        shapable_or_shape: Optional[Union[Shapable, Tuple[int, int]]] = None,
 85        image: Optional[Image] = None,
 86        mask: Optional[Mask] = None,
 87        score_map: Optional[ScoreMap] = None,
 88        point: Optional[Point] = None,
 89        points: Optional[Union[PointList, PointTuple, Iterable[Point]]] = None,
 90        corner_points: Optional[Union[PointList, PointTuple, Iterable[Point]]] = None,
 91        polygon: Optional[Polygon] = None,
 92        polygons: Optional[Iterable[Polygon]] = None,
 93        rng: Optional[RandomGenerator] = None,
 94        enable_debug: bool = False,
 95    ):
 96        config_generator = self.config_generator_cls(
 97            self.config_for_config_generator,
 98            level,
 99        )
100        return self.distortion.distort(
101            config_or_config_generator=config_generator,
102            shapable_or_shape=shapable_or_shape,
103            image=image,
104            mask=mask,
105            score_map=score_map,
106            point=point,
107            points=points,
108            corner_points=corner_points,
109            polygon=polygon,
110            polygons=polygons,
111            rng=rng,
112            get_config=enable_debug,
113            get_state=enable_debug,
114        )
115
116    @property
117    def name(self):
118        return self.config_generator_cls.get_config_cls().get_name()
119
120    def __repr__(self):
121        return f'DistortionPolicy({self.name})'
122
123
124class DistortionPolicyFactory(Generic[_T_GENERATOR_CONFIG, _T_CONFIG, _T_STATE]):
125
126    def __init__(
127        self,
128        distortion: Distortion[_T_CONFIG, _T_STATE],
129        config_generator_cls: Type[DistortionConfigGenerator[_T_GENERATOR_CONFIG, _T_CONFIG]],
130    ):
131        self.distortion = distortion
132        self.config_generator_cls = config_generator_cls
133
134    def create(
135        self,
136        config: Optional[Union[Mapping[str, Any], PathType, _T_GENERATOR_CONFIG]] = None,
137    ):
138        config = dyn_structure(
139            config,
140            self.config_generator_cls.get_generator_config_cls(),
141            support_path_type=True,
142            support_none_type=True,
143        )
144        return DistortionPolicy(
145            self.distortion,
146            config,
147            self.config_generator_cls,
148        )
149
150    @property
151    def name(self):
152        return self.config_generator_cls.get_config_cls().get_name()