Source code for tensorpack.dataflow.imgaug.geometry

# -*- coding: utf-8 -*-
# File: geometry.py


import math
import numpy as np
import cv2

from .base import ImageAugmentor
from .transform import WarpAffineTransform, CropTransform, TransformList

__all__ = ['Shift', 'Rotation', 'RotationAndCropValid', 'Affine']


[docs]class Shift(ImageAugmentor): """ Random horizontal and vertical shifts """
[docs] def __init__(self, horiz_frac=0, vert_frac=0, border=cv2.BORDER_REPLICATE, border_value=0): """ Args: horiz_frac (float): max abs fraction for horizontal shift vert_frac (float): max abs fraction for horizontal shift border: cv2 border method border_value: cv2 border value for border=cv2.BORDER_CONSTANT """ assert horiz_frac < 1.0 and vert_frac < 1.0 super(Shift, self).__init__() self._init(locals())
[docs] def get_transform(self, img): max_dx = self.horiz_frac * img.shape[1] max_dy = self.vert_frac * img.shape[0] dx = np.round(self._rand_range(-max_dx, max_dx)) dy = np.round(self._rand_range(-max_dy, max_dy)) mat = np.array([[1, 0, dx], [0, 1, dy]], dtype='float32') return WarpAffineTransform( mat, img.shape[1::-1], borderMode=self.border, borderValue=self.border_value)
[docs]class Rotation(ImageAugmentor): """ Random rotate the image w.r.t a random center"""
[docs] def __init__(self, max_deg, center_range=(0, 1), interp=cv2.INTER_LINEAR, border=cv2.BORDER_REPLICATE, step_deg=None, border_value=0): """ Args: max_deg (float): max abs value of the rotation angle (in degree). center_range (tuple): (min, max) range of the random rotation center. interp: cv2 interpolation method border: cv2 border method step_deg (float): if not None, the stepping of the rotation angle. The rotation angle will be a multiple of step_deg. This option requires ``max_deg==180`` and step_deg has to be a divisor of 180) border_value: cv2 border value for border=cv2.BORDER_CONSTANT """ assert step_deg is None or (max_deg == 180 and max_deg % step_deg == 0) super(Rotation, self).__init__() self._init(locals())
[docs] def get_transform(self, img): center = img.shape[1::-1] * self._rand_range( self.center_range[0], self.center_range[1], (2,)) deg = self._rand_range(-self.max_deg, self.max_deg) if self.step_deg: deg = deg // self.step_deg * self.step_deg """ The correct center is shape*0.5-0.5. This can be verified by: SHAPE = 7 arr = np.random.rand(SHAPE, SHAPE) orig = arr c = SHAPE * 0.5 - 0.5 c = (c, c) for k in range(4): mat = cv2.getRotationMatrix2D(c, 90, 1) arr = cv2.warpAffine(arr, mat, arr.shape) assert np.all(arr == orig) """ mat = cv2.getRotationMatrix2D(tuple(center - 0.5), float(deg), 1) return WarpAffineTransform( mat, img.shape[1::-1], interp=self.interp, borderMode=self.border, borderValue=self.border_value)
[docs]class RotationAndCropValid(ImageAugmentor): """ Random rotate and then crop the largest possible rectangle. Note that this will produce images of different shapes. """
[docs] def __init__(self, max_deg, interp=cv2.INTER_LINEAR, step_deg=None): """ Args: max_deg, interp, step_deg: same as :class:`Rotation` """ assert step_deg is None or (max_deg == 180 and max_deg % step_deg == 0) super(RotationAndCropValid, self).__init__() self._init(locals())
def _get_deg(self, img): deg = self._rand_range(-self.max_deg, self.max_deg) if self.step_deg: deg = deg // self.step_deg * self.step_deg return float(deg)
[docs] def get_transform(self, img): deg = self._get_deg(img) h, w = img.shape[:2] center = (img.shape[1] * 0.5, img.shape[0] * 0.5) rot_m = cv2.getRotationMatrix2D((center[0] - 0.5, center[1] - 0.5), deg, 1) tfm = WarpAffineTransform(rot_m, (w, h), interp=self.interp) neww, newh = RotationAndCropValid.largest_rotated_rect(w, h, deg) neww = min(neww, w) newh = min(newh, h) newx = int(center[0] - neww * 0.5) newy = int(center[1] - newh * 0.5) tfm2 = CropTransform(newy, newx, newh, neww) return TransformList([tfm, tfm2])
[docs] @staticmethod def largest_rotated_rect(w, h, angle): """ Get largest rectangle after rotation. http://stackoverflow.com/questions/16702966/rotate-image-and-crop-out-black-borders """ angle = angle / 180.0 * math.pi if w <= 0 or h <= 0: return 0, 0 width_is_longer = w >= h side_long, side_short = (w, h) if width_is_longer else (h, w) # since the solutions for angle, -angle and 180-angle are all the same, # if suffices to look at the first quadrant and the absolute values of sin,cos: sin_a, cos_a = abs(math.sin(angle)), abs(math.cos(angle)) if side_short <= 2. * sin_a * cos_a * side_long: # half constrained case: two crop corners touch the longer side, # the other two corners are on the mid-line parallel to the longer line x = 0.5 * side_short wr, hr = (x / sin_a, x / cos_a) if width_is_longer else (x / cos_a, x / sin_a) else: # fully constrained case: crop touches all 4 sides cos_2a = cos_a * cos_a - sin_a * sin_a wr, hr = (w * cos_a - h * sin_a) / cos_2a, (h * cos_a - w * sin_a) / cos_2a return int(np.round(wr)), int(np.round(hr))
[docs]class Affine(ImageAugmentor): """ Random affine transform of the image w.r.t to the image center. Transformations involve: - Translation ("move" image on the x-/y-axis) - Rotation - Scaling ("zoom" in/out) - Shear (move one side of the image, turning a square into a trapezoid) """
[docs] def __init__(self, scale=None, translate_frac=None, rotate_max_deg=0.0, shear=0.0, interp=cv2.INTER_LINEAR, border=cv2.BORDER_REPLICATE, border_value=0): """ Args: scale (tuple of 2 floats): scaling factor interval, e.g (a, b), then scale is randomly sampled from the range a <= scale <= b. Will keep original scale by default. translate_frac (tuple of 2 floats): tuple of max abs fraction for horizontal and vertical translation. For example translate_frac=(a, b), then horizontal shift is randomly sampled in the range 0 < dx < img_width * a and vertical shift is randomly sampled in the range 0 < dy < img_height * b. Will not translate by default. shear (float): max abs shear value in degrees between 0 to 180 interp: cv2 interpolation method border: cv2 border method border_value: cv2 border value for border=cv2.BORDER_CONSTANT """ if scale is not None: assert isinstance(scale, tuple) and len(scale) == 2, \ "Argument scale should be a tuple of two floats, e.g (a, b)" if translate_frac is not None: assert isinstance(translate_frac, tuple) and len(translate_frac) == 2, \ "Argument translate_frac should be a tuple of two floats, e.g (a, b)" assert shear >= 0.0, "Argument shear should be between 0.0 and 180.0" super(Affine, self).__init__() self._init(locals())
[docs] def get_transform(self, img): if self.scale is not None: scale = self._rand_range(self.scale[0], self.scale[1]) else: scale = 1.0 if self.translate_frac is not None: max_dx = self.translate_frac[0] * img.shape[1] max_dy = self.translate_frac[1] * img.shape[0] dx = np.round(self._rand_range(-max_dx, max_dx)) dy = np.round(self._rand_range(-max_dy, max_dy)) else: dx = 0 dy = 0 if self.shear > 0.0: shear = self._rand_range(-self.shear, self.shear) sin_shear = math.sin(math.radians(shear)) cos_shear = math.cos(math.radians(shear)) else: sin_shear = 0.0 cos_shear = 1.0 center = (img.shape[1::-1] * np.array((0.5, 0.5))) - 0.5 deg = self._rand_range(-self.rotate_max_deg, self.rotate_max_deg) transform_matrix = cv2.getRotationMatrix2D(tuple(center), deg, scale) # Apply shear : if self.shear > 0.0: m00 = transform_matrix[0, 0] m01 = transform_matrix[0, 1] m10 = transform_matrix[1, 0] m11 = transform_matrix[1, 1] transform_matrix[0, 1] = m01 * cos_shear + m00 * sin_shear transform_matrix[1, 1] = m11 * cos_shear + m10 * sin_shear # Add correction term to keep the center unchanged tx = center[0] * (1.0 - m00) - center[1] * transform_matrix[0, 1] ty = -center[0] * m10 + center[1] * (1.0 - transform_matrix[1, 1]) transform_matrix[0, 2] = tx transform_matrix[1, 2] = ty # Apply shift : transform_matrix[0, 2] += dx transform_matrix[1, 2] += dy return WarpAffineTransform(transform_matrix, img.shape[1::-1], self.interp, self.border, self.border_value)