File size: 1,642 Bytes
36c95ba | 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 | from enum import Enum, EnumMeta
from typing import cast, TypeVar, Union
import torch
__all__ = ['pi', 'Resample', 'BorderType', 'SamplePadding']
pi = torch.tensor(3.14159265358979323846)
T = TypeVar('T', bound='ConstantBase')
class ConstantBase:
@classmethod
def get(cls, value: Union[str, int, T]) -> T: # type: ignore
if type(value) is str:
return cls[value.upper()] # type: ignore
if type(value) is int:
return cls(value) # type: ignore
if type(value) is cls:
return value # type: ignore
raise TypeError()
class EnumMetaFlags(EnumMeta):
def __contains__(self, other: Union[str, int, T]) -> bool: # type: ignore
if type(other) is str:
other = cast(str, other)
return any(val.name == other.upper() for val in self) # type: ignore
if type(other) is int:
return any(val.value == other for val in self) # type: ignore
return any(val == other for val in self) # type: ignore
def __repr__(self):
return ' | '.join(f"{self.__name__}.{val.name}" for val in self)
class Resample(ConstantBase, Enum, metaclass=EnumMetaFlags):
NEAREST = 0
BILINEAR = 1
BICUBIC = 2
class BorderType(ConstantBase, Enum, metaclass=EnumMetaFlags):
CONSTANT = 0
REFLECT = 1
REPLICATE = 2
CIRCULAR = 3
class SamplePadding(ConstantBase, Enum, metaclass=EnumMetaFlags):
ZEROS = 0
BORDER = 1
REFLECTION = 2
class DataKey(ConstantBase, Enum, metaclass=EnumMetaFlags):
INPUT = 0
MASK = 1
BBOX = 2
BBOX_XYXY = 3
BBOX_XYHW = 4
KEYPOINTS = 5
|