Spaces:
Configuration error
Configuration error
File size: 10,763 Bytes
0034848 |
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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 |
from __future__ import absolute_import
import random
from copy import deepcopy
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, cast
from warnings import warn
import cv2
import numpy as np
from .serialization import Serializable, get_shortest_class_fullname
from .utils import format_args
__all__ = [
"to_tuple",
"BasicTransform",
"DualTransform",
"ImageOnlyTransform",
"NoOp",
"BoxType",
"KeypointType",
"ImageColorType",
"ScaleFloatType",
"ScaleIntType",
"ImageColorType",
]
NumType = Union[int, float, np.ndarray]
BoxInternalType = Tuple[float, float, float, float]
BoxType = Union[BoxInternalType, Tuple[float, float, float, float, Any]]
KeypointInternalType = Tuple[float, float, float, float]
KeypointType = Union[KeypointInternalType, Tuple[float, float, float, float, Any]]
ImageColorType = Union[float, Sequence[float]]
ScaleFloatType = Union[float, Tuple[float, float]]
ScaleIntType = Union[int, Tuple[int, int]]
FillValueType = Optional[Union[int, float, Sequence[int], Sequence[float]]]
def to_tuple(param, low=None, bias=None):
"""Convert input argument to min-max tuple
Args:
param (scalar, tuple or list of 2+ elements): Input value.
If value is scalar, return value would be (offset - value, offset + value).
If value is tuple, return value would be value + offset (broadcasted).
low: Second element of tuple can be passed as optional argument
bias: An offset factor added to each element
"""
if low is not None and bias is not None:
raise ValueError("Arguments low and bias are mutually exclusive")
if param is None:
return param
if isinstance(param, (int, float)):
if low is None:
param = -param, +param
else:
param = (low, param) if low < param else (param, low)
elif isinstance(param, Sequence):
if len(param) != 2:
raise ValueError("to_tuple expects 1 or 2 values")
param = tuple(param)
else:
raise ValueError("Argument param must be either scalar (int, float) or tuple")
if bias is not None:
return tuple(bias + x for x in param)
return tuple(param)
class BasicTransform(Serializable):
call_backup = None
interpolation: Any
fill_value: Any
mask_fill_value: Any
def __init__(self, always_apply: bool = False, p: float = 0.5):
self.p = p
self.always_apply = always_apply
self._additional_targets: Dict[str, str] = {}
# replay mode params
self.deterministic = False
self.save_key = "replay"
self.params: Dict[Any, Any] = {}
self.replay_mode = False
self.applied_in_replay = False
def __call__(self, *args, force_apply: bool = False, **kwargs) -> Dict[str, Any]:
if args:
raise KeyError("You have to pass data to augmentations as named arguments, for example: aug(image=image)")
if self.replay_mode:
if self.applied_in_replay:
return self.apply_with_params(self.params, **kwargs)
return kwargs
if (random.random() < self.p) or self.always_apply or force_apply:
params = self.get_params()
if self.targets_as_params:
assert all(key in kwargs for key in self.targets_as_params), "{} requires {}".format(
self.__class__.__name__, self.targets_as_params
)
targets_as_params = {k: kwargs[k] for k in self.targets_as_params}
params_dependent_on_targets = self.get_params_dependent_on_targets(targets_as_params)
params.update(params_dependent_on_targets)
if self.deterministic:
if self.targets_as_params:
warn(
self.get_class_fullname() + " could work incorrectly in ReplayMode for other input data"
" because its' params depend on targets."
)
kwargs[self.save_key][id(self)] = deepcopy(params)
return self.apply_with_params(params, **kwargs)
return kwargs
def apply_with_params(self, params: Dict[str, Any], **kwargs) -> Dict[str, Any]: # skipcq: PYL-W0613
if params is None:
return kwargs
params = self.update_params(params, **kwargs)
res = {}
for key, arg in kwargs.items():
if arg is not None:
target_function = self._get_target_function(key)
target_dependencies = {k: kwargs[k] for k in self.target_dependence.get(key, [])}
res[key] = target_function(arg, **dict(params, **target_dependencies))
else:
res[key] = None
return res
def set_deterministic(self, flag: bool, save_key: str = "replay") -> "BasicTransform":
assert save_key != "params", "params save_key is reserved"
self.deterministic = flag
self.save_key = save_key
return self
def __repr__(self) -> str:
state = self.get_base_init_args()
state.update(self.get_transform_init_args())
return "{name}({args})".format(name=self.__class__.__name__, args=format_args(state))
def _get_target_function(self, key: str) -> Callable:
transform_key = key
if key in self._additional_targets:
transform_key = self._additional_targets.get(key, key)
target_function = self.targets.get(transform_key, lambda x, **p: x)
return target_function
def apply(self, img: np.ndarray, **params) -> np.ndarray:
raise NotImplementedError
def get_params(self) -> Dict:
return {}
@property
def targets(self) -> Dict[str, Callable]:
# you must specify targets in subclass
# for example: ('image', 'mask')
# ('image', 'boxes')
raise NotImplementedError
def update_params(self, params: Dict[str, Any], **kwargs) -> Dict[str, Any]:
if hasattr(self, "interpolation"):
params["interpolation"] = self.interpolation
if hasattr(self, "fill_value"):
params["fill_value"] = self.fill_value
if hasattr(self, "mask_fill_value"):
params["mask_fill_value"] = self.mask_fill_value
params.update({"cols": kwargs["image"].shape[1], "rows": kwargs["image"].shape[0]})
return params
@property
def target_dependence(self) -> Dict:
return {}
def add_targets(self, additional_targets: Dict[str, str]):
"""Add targets to transform them the same way as one of existing targets
ex: {'target_image': 'image'}
ex: {'obj1_mask': 'mask', 'obj2_mask': 'mask'}
by the way you must have at least one object with key 'image'
Args:
additional_targets (dict): keys - new target name, values - old target name. ex: {'image2': 'image'}
"""
self._additional_targets = additional_targets
@property
def targets_as_params(self) -> List[str]:
return []
def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, Any]:
raise NotImplementedError(
"Method get_params_dependent_on_targets is not implemented in class " + self.__class__.__name__
)
@classmethod
def get_class_fullname(cls) -> str:
return get_shortest_class_fullname(cls)
@classmethod
def is_serializable(cls):
return True
def get_transform_init_args_names(self) -> Tuple[str, ...]:
raise NotImplementedError(
"Class {name} is not serializable because the `get_transform_init_args_names` method is not "
"implemented".format(name=self.get_class_fullname())
)
def get_base_init_args(self) -> Dict[str, Any]:
return {"always_apply": self.always_apply, "p": self.p}
def get_transform_init_args(self) -> Dict[str, Any]:
return {k: getattr(self, k) for k in self.get_transform_init_args_names()}
def _to_dict(self) -> Dict[str, Any]:
state = {"__class_fullname__": self.get_class_fullname()}
state.update(self.get_base_init_args())
state.update(self.get_transform_init_args())
return state
def get_dict_with_id(self) -> Dict[str, Any]:
d = self._to_dict()
d["id"] = id(self)
return d
class DualTransform(BasicTransform):
"""Transform for segmentation task."""
@property
def targets(self) -> Dict[str, Callable]:
return {
"image": self.apply,
"mask": self.apply_to_mask,
"masks": self.apply_to_masks,
"bboxes": self.apply_to_bboxes,
"keypoints": self.apply_to_keypoints,
}
def apply_to_bbox(self, bbox: BoxInternalType, **params) -> BoxInternalType:
raise NotImplementedError("Method apply_to_bbox is not implemented in class " + self.__class__.__name__)
def apply_to_keypoint(self, keypoint: KeypointInternalType, **params) -> KeypointInternalType:
raise NotImplementedError("Method apply_to_keypoint is not implemented in class " + self.__class__.__name__)
def apply_to_bboxes(self, bboxes: Sequence[BoxType], **params) -> List[BoxType]:
return [self.apply_to_bbox(tuple(bbox[:4]), **params) + tuple(bbox[4:]) for bbox in bboxes] # type: ignore
def apply_to_keypoints(self, keypoints: Sequence[KeypointType], **params) -> List[KeypointType]:
return [ # type: ignore
self.apply_to_keypoint(tuple(keypoint[:4]), **params) + tuple(keypoint[4:]) # type: ignore
for keypoint in keypoints
]
def apply_to_mask(self, img: np.ndarray, **params) -> np.ndarray:
return self.apply(img, **{k: cv2.INTER_NEAREST if k == "interpolation" else v for k, v in params.items()})
def apply_to_masks(self, masks: Sequence[np.ndarray], **params) -> List[np.ndarray]:
return [self.apply_to_mask(mask, **params) for mask in masks]
class ImageOnlyTransform(BasicTransform):
"""Transform applied to image only."""
@property
def targets(self) -> Dict[str, Callable]:
return {"image": self.apply}
class NoOp(DualTransform):
"""Does nothing"""
def apply_to_keypoint(self, keypoint: KeypointInternalType, **params) -> KeypointInternalType:
return keypoint
def apply_to_bbox(self, bbox: BoxInternalType, **params) -> BoxInternalType:
return bbox
def apply(self, img: np.ndarray, **params) -> np.ndarray:
return img
def apply_to_mask(self, img: np.ndarray, **params) -> np.ndarray:
return img
def get_transform_init_args_names(self) -> Tuple:
return ()
|