Spaces:
Sleeping
Sleeping
File size: 2,886 Bytes
20239f9 |
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 |
# Description: This file contains the code for the reversible affine transform
import torchvision.transforms as transforms
import torch
from typing import List, Optional, Tuple, Any
def generate_affine_trans_params(
degrees: List[float],
translate: Optional[List[float]],
scale_ranges: Optional[List[float]],
shears: Optional[List[float]],
img_size: List[int],
) -> Tuple[float, Tuple[int, int], float, Any]:
"""Get parameters for affine transformation
Returns:
params to be passed to the affine transformation
"""
angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item())
if translate is not None:
max_dx = float(translate[0] * img_size[0])
max_dy = float(translate[1] * img_size[1])
tx = int(round(torch.empty(1).uniform_(-max_dx, max_dx).item()))
ty = int(round(torch.empty(1).uniform_(-max_dy, max_dy).item()))
translations = (tx, ty)
else:
translations = (0, 0)
if scale_ranges is not None:
scale = float(torch.empty(1).uniform_(scale_ranges[0], scale_ranges[1]).item())
else:
scale = 1.0
shear_x = shear_y = 0.0
if shears is not None:
shear_x = float(torch.empty(1).uniform_(shears[0], shears[1]).item())
if len(shears) == 4:
shear_y = float(torch.empty(1).uniform_(shears[2], shears[3]).item())
shear = (shear_x, shear_y)
if shear_x == 0.0 and shear_y == 0.0:
shear = 0.0
return angle, translations, scale, shear
def rigid_transform(img, angle, translate, scale, invert=False, shear=0,
interpolation=transforms.InterpolationMode.BILINEAR):
"""
Affine transforms input image
Modified from: https://github.com/robertdvdk/part_detection/blob/eec53f2f40602113f74c6c1f60a2034823b0fcaf/lib.py#L54
Parameters
----------
img: Tensor
Input image
angle: int
Rotation angle between -180 and 180 degrees
translate: [int]
Sequence of horizontal/vertical translations
scale: float
How to scale the image
invert: bool
Whether to invert the transformation
shear: float
Shear angle in degrees
interpolation: InterpolationMode
Interpolation mode to calculate output values
Returns
----------
img: Tensor
Transformed image
"""
if not invert:
img = transforms.functional.affine(img, angle=angle, translate=translate, scale=scale, shear=shear,
interpolation=interpolation)
else:
translate = [-t for t in translate]
img = transforms.functional.affine(img=img, angle=0, translate=translate, scale=1, shear=shear)
img = transforms.functional.affine(img=img, angle=-angle, translate=[0, 0], scale=1 / scale, shear=shear)
return img
|