import types
from typing import Any, List, Tuple
import cv2
import numpy
from PIL.ImageTransform import Transform
from numpy import random
__all__ = [
"CV2Compose",
"Lambda",
"ConvertFromInts",
"SubtractMeans",
"CV2ToAbsoluteCoords",
"CV2ToPercentCoords",
"CV2Resize",
"CV2RandomSaturation",
"CV2RandomHue",
"CV2RandomLightingNoise",
"CV2ConvertColor",
"CV2RandomContrast",
"CV2RandomBrightness",
"CV2RandomSampleCrop",
"CV2Expand",
"CV2RandomMirror",
"CV2SwapChannels",
"CV2PhotometricDistort",
]
from draugr.opencv_utilities.bounding_boxes.evaluation import (
jaccard_overlap_numpy,
remove_null_boxes,
)
from warg import TripleNumber
[docs]class CV2Compose(object):
"""Composes several augmentations together.
Args:
transforms (List[Transform]): list of transforms to compose.
Example:
>>> augmentations.Compose([
>>> transforms.CenterCrop(10),
>>> transforms.ToTensor(),
>>> ])"""
[docs] def __init__(self, transforms: List[Transform]):
self._transforms = transforms
def __call__(self, img, boxes=None, labels=None) -> Tuple:
for t in self._transforms:
img, boxes, labels = t(img, boxes, labels)
if boxes is not None:
boxes, labels = remove_null_boxes(boxes, labels)
return img, boxes, labels
[docs]class Lambda(object):
"""Applies a lambda as a transform."""
[docs] def __init__(self, lambd: callable):
assert isinstance(lambd, types.LambdaType)
self.lambd = lambd
def __call__(self, img: Any, boxes: Any = None, labels: Any = None):
return self.lambd(img, boxes, labels)
[docs]class ConvertFromInts(object):
def __call__(self, image: Any, boxes: Any = None, labels: Any = None) -> Tuple:
return image.astype(numpy.float32), boxes, labels
[docs]class SubtractMeans(object):
"""description"""
[docs] def __init__(self, mean: Tuple):
self.mean = numpy.array(mean, dtype=numpy.float32)
def __call__(self, image: Any, boxes: Any = None, labels: Any = None) -> Tuple:
image = image.astype(numpy.float32)
image -= self.mean
return image.astype(numpy.float32), boxes, labels
[docs]class CV2ToAbsoluteCoords(object):
def __call__(self, image: Any, boxes: Any = None, labels: Any = None) -> Tuple:
height, width, channels = image.shape
boxes[:, 0] *= width
boxes[:, 2] *= width
boxes[:, 1] *= height
boxes[:, 3] *= height
return image, boxes, labels
[docs]class CV2ToPercentCoords(object):
def __call__(
self,
image: numpy.ndarray,
boxes: numpy.ndarray = None,
labels: numpy.ndarray = None,
) -> Tuple:
height, width, channels = image.shape
boxes[:, 0] /= width
boxes[:, 2] /= width
boxes[:, 1] /= height
boxes[:, 3] /= height
return image, boxes, labels
[docs]class CV2Resize(object):
"""description"""
[docs] def __init__(self, size: int = 300):
self._size = size
def __call__(
self,
image: numpy.ndarray,
boxes: numpy.ndarray = None,
labels: numpy.ndarray = None,
) -> Tuple:
image = cv2.resize(image, (self._size, self._size))
return image, boxes, labels
[docs]class CV2RandomSaturation(object):
"""description"""
[docs] def __init__(self, lower: float = 0.5, upper: float = 1.5):
self.lower = lower
self.upper = upper
assert self.upper >= self.lower, "contrast upper must be >= lower."
assert self.lower >= 0, "contrast lower must be non-negative."
def __call__(
self,
image: numpy.ndarray,
boxes: numpy.ndarray = None,
labels: numpy.ndarray = None,
) -> Tuple:
if random.randint(2):
image[..., 1] *= random.uniform(self.lower, self.upper)
return image, boxes, labels
[docs]class CV2RandomHue(object):
"""description"""
[docs] def __init__(self, delta: float = 18.0):
assert 0.0 <= delta <= 360.0
self.delta = delta
def __call__(
self,
image: numpy.ndarray,
boxes: numpy.ndarray = None,
labels: numpy.ndarray = None,
) -> Tuple:
if random.randint(2):
image[..., 0] += random.uniform(-self.delta, self.delta)
image[..., 0][image[..., 0] > 360.0] -= 360.0
image[..., 0][image[..., 0] < 0.0] += 360.0
return image, boxes, labels
[docs]class CV2RandomLightingNoise(object):
"""description"""
[docs] def __init__(self):
self.perms = ((0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0))
def __call__(
self,
image: numpy.ndarray,
boxes: numpy.ndarray = None,
labels: numpy.ndarray = None,
) -> Tuple:
if random.randint(2):
swap = self.perms[random.randint(len(self.perms))]
shuffle = CV2SwapChannels(swap) # shuffle channels
image = shuffle(image)
return image, boxes, labels
[docs]class CV2ConvertColor(object):
"""description"""
[docs] def __init__(self, current, transform):
self.transform = transform
self.current = current
def __call__(
self,
image: numpy.ndarray,
boxes: numpy.ndarray = None,
labels: numpy.ndarray = None,
) -> Tuple:
if self.current == "BGR" and self.transform == "HSV":
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
elif self.current == "RGB" and self.transform == "HSV":
image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
elif self.current == "BGR" and self.transform == "RGB":
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
elif self.current == "HSV" and self.transform == "BGR":
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
elif self.current == "HSV" and self.transform == "RGB":
image = cv2.cvtColor(image, cv2.COLOR_HSV2RGB)
else:
raise NotImplementedError
return image, boxes, labels
[docs]class CV2RandomContrast(object):
"""description"""
[docs] def __init__(self, lower: float = 0.5, upper: float = 1.5):
self.lower = lower
self.upper = upper
assert self.upper >= self.lower, "contrast upper must be >= lower."
assert self.lower >= 0, "contrast lower must be non-negative."
# expects float image
def __call__(
self,
image: numpy.ndarray,
boxes: numpy.ndarray = None,
labels: numpy.ndarray = None,
) -> Tuple:
if random.randint(2):
alpha = random.uniform(self.lower, self.upper)
image *= alpha
return image, boxes, labels
[docs]class CV2RandomBrightness(object):
"""description"""
[docs] def __init__(self, delta: float = 32):
assert delta >= 0.0
assert delta <= 255.0
self.delta = delta
def __call__(
self,
image: numpy.ndarray,
boxes: numpy.ndarray = None,
labels: numpy.ndarray = None,
) -> Tuple:
if random.randint(2):
delta = random.uniform(-self.delta, self.delta)
image += delta
return image, boxes, labels
[docs]class CV2RandomSampleCrop(object):
"""Crop
Arguments:
img (Image): the image being input during training
boxes (Tensor): the original bounding boxes in pt form
labels (Tensor): the class labels for each bbox
mode (float tuple): the min and max jaccard overlaps
Return:
(img, boxes, classes)
img (Image): the cropped image
boxes (Tensor): the adjusted bounding boxes in pt form
labels (Tensor): the class labels for each bbox"""
[docs] def __init__(self):
self.sample_options = (
# using entire original input image
None,
# sample a patch s.t. MIN jaccard w/ obj in .1,.3,.4,.7,.9
(0.1, None),
(0.3, None),
(0.7, None),
(0.9, None),
# randomly sample a patch
(None, None),
)
def __call__(
self,
image: numpy.ndarray,
boxes: numpy.ndarray = None,
labels: numpy.ndarray = None,
) -> Tuple:
# guard against no boxes
if boxes is not None and boxes.shape[0] == 0:
return image, boxes, labels
height, width, _ = image.shape
while True:
# randomly choose a mode
mode = random.choice(self.sample_options)
if mode is None:
return image, boxes, labels
min_iou, max_iou = mode
if min_iou is None:
min_iou = float("-inf")
if max_iou is None:
max_iou = float("inf")
# max trails (50)
for _ in range(50):
current_image = image
w = random.uniform(0.3 * width, width)
h = random.uniform(0.3 * height, height)
# aspect ratio constraint b/t .5 & 2
if h / w < 0.5 or h / w > 2:
continue
left = random.uniform(width - w)
top = random.uniform(height - h)
# convert to integer rect x1,y1,x2,y2
rect = numpy.array([int(left), int(top), int(left + w), int(top + h)])
# calculate IoU (jaccard overlap) b/t the cropped and gt boxes
overlap = jaccard_overlap_numpy(boxes, rect)
# is min and max overlap constraint satisfied? if not try again
if overlap.max() < min_iou or overlap.min() > max_iou:
continue
# cut the crop from the image
current_image = current_image[rect[1] : rect[3], rect[0] : rect[2], :]
# keep overlap with gt box IF center in sampled patch
centers = (boxes[:, :2] + boxes[:, 2:]) / 2.0
# mask in all gt boxes that above and to the left of centers
m1 = (rect[0] < centers[:, 0]) * (rect[1] < centers[:, 1])
# mask in all gt boxes that under and to the right of centers
m2 = (rect[2] > centers[:, 0]) * (rect[3] > centers[:, 1])
# mask in that both m1 and m2 are true
mask = m1 * m2
# have any valid boxes? try again if not
if not mask.any():
continue
# take only matching gt boxes
current_boxes = boxes[mask, :].copy()
# take only matching gt labels
current_labels = labels[mask]
# should we use the box left and top corner or the crop's
current_boxes[:, :2] = numpy.maximum(current_boxes[:, :2], rect[:2])
# adjust to crop (by substracting crop's left,top)
current_boxes[:, :2] -= rect[:2]
current_boxes[:, 2:] = numpy.minimum(current_boxes[:, 2:], rect[2:])
# adjust to crop (by substracting crop's left,top)
current_boxes[:, 2:] -= rect[:2]
return current_image, current_boxes, current_labels
[docs]class CV2Expand(object):
"""description"""
[docs] def __init__(self, mean: float):
self.mean = mean
def __call__(
self, image: numpy.ndarray, boxes: numpy.ndarray, labels: numpy.ndarray
) -> Tuple:
if random.randint(2):
return image, boxes, labels
height, width, depth = image.shape
ratio = random.uniform(1, 4)
left = random.uniform(0, width * ratio - width)
top = random.uniform(0, height * ratio - height)
expand_image = numpy.zeros(
(int(height * ratio), int(width * ratio), depth), dtype=image.dtype
)
expand_image[..., :] = self.mean
expand_image[
int(top) : int(top + height), int(left) : int(left + width)
] = image
image = expand_image
boxes = boxes.copy()
boxes[:, :2] += (int(left), int(top))
boxes[:, 2:] += (int(left), int(top))
return image, boxes, labels
[docs]class CV2RandomMirror(object):
def __call__(
self, image: numpy.ndarray, boxes: numpy.ndarray, classes: numpy.ndarray
) -> Tuple:
_, width, _ = image.shape
if random.randint(2):
image = image[:, ::-1]
boxes = boxes.copy()
boxes[:, 0::2] = width - boxes[:, 2::-2]
return image, boxes, classes
[docs]class CV2SwapChannels(object):
"""Transforms a tensorized image by swapping the channels in the order
specified in the swap tuple.
Args:
swaps (int triple): final order of channels
eg: (2, 1, 0)"""
[docs] def __init__(self, swaps: TripleNumber):
self.swaps = swaps
def __call__(self, image: numpy.ndarray) -> numpy.ndarray:
"""
Args:
image (Tensor): image tensor to be transformed
Return:
a tensor with channels swapped according to swap"""
# if torch.is_tensor(image):
# image = image.data.cpu().numpy()
# else:
# image = numpy.array(image)
image = image[..., self.swaps]
return image
[docs]class CV2PhotometricDistort(object):
"""description"""
[docs] def __init__(self):
self.pd = [
CV2RandomContrast(), # RGB
CV2ConvertColor(current="RGB", transform="HSV"), # HSV
CV2RandomSaturation(), # HSV
CV2RandomHue(), # HSV
CV2ConvertColor(current="HSV", transform="RGB"), # RGB
CV2RandomContrast(), # RGB
]
self.rand_brightness = CV2RandomBrightness()
self.rand_light_noise = CV2RandomLightingNoise()
def __call__(
self, image: numpy.ndarray, boxes: numpy.ndarray, labels: numpy.ndarray
) -> Tuple:
im = image.copy()
im, boxes, labels = self.rand_brightness(im, boxes, labels)
if random.randint(2):
distort = CV2Compose(self.pd[:-1])
else:
distort = CV2Compose(self.pd[1:])
im, boxes, labels = distort(im, boxes, labels)
return self.rand_light_noise(im, boxes, labels)