"""
Random translate
"""
import numpy as np
from numpy.random import randint
from typing import Callable, Optional, Tuple
from dataclasses import replace
from ..pipeline.allocation_query import AllocationQuery
from ..pipeline.operation import Operation
from ..pipeline.state import State
from ..pipeline.compiler import Compiler
[docs]class RandomTranslate(Operation):
"""Translate each image randomly in vertical and horizontal directions
up to specified number of pixels.
Parameters
----------
padding : int
Max number of pixels to translate in any direction.
fill : tuple
An RGB color ((0, 0, 0) by default) to fill the area outside the shifted image.
"""
def __init__(self, padding: int, fill: Tuple[int, int, int] = (0, 0, 0)):
super().__init__()
self.padding = padding
self.fill = np.array(fill)
[docs] def generate_code(self) -> Callable:
my_range = Compiler.get_iterator()
pad = self.padding
fill = self.fill
def translate(images, dst):
n, h, w, _ = images.shape
dst[:] = fill
dst[:, pad:pad+h, pad:pad+w] = images
for i in my_range(n):
y_coord = randint(low=0, high=2 * pad + 1)
x_coord = randint(low=0, high=2 * pad + 1)
images[i] = dst[i, y_coord:y_coord+h, x_coord:x_coord+w]
return images
translate.is_parallel = True
return translate
[docs] def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]:
h, w, c = previous_state.shape
return (replace(previous_state, jit_mode=True), \
AllocationQuery((h + 2 * self.padding, w + 2 * self.padding, c), previous_state.dtype))