spynet(二)：光流估计中的数据增强函数

7.F.rotate()

使用前需要：

import torchvision.transforms.functional as F  

和RandomRotation功能类似
The RandomRotation transform (see also rotate()) rotates an image with random angle.

rotater = T.RandomRotation(degrees=(0, 180))
rotated_imgs = [rotater(orig_img) for _ in range(4)]
plot(rotated_imgs)

得到结果：

8. 数据增强相关操作

可以参考文章

import random
from PIL import Image
import numpy as np
import torch

# 图像处理函数
import torchvision.transforms.functional as F
from skimage import transform


# 用于注释参数
from typing import Union, Tuple 
ImageOrTensor = Union['Image', torch.Tensor]            
Transformed = Tuple[Tuple[ImageOrTensor, ImageOrTensor], 
                    Union[np.ndarray, torch.Tensor]]


class Compose(object):

    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, *args):
        o = args
        for t in self.transforms:
            o = t(*o)
        return o


class Resize(object): # 对图像光流都resize同样角度

    def __init__(self, height: int, width: int) -> None:
        self.height = height
        self.width = width
    
    def __call__(self, 
                 frames: Tuple['Image', 'Image'], 
                 optical_flow: np.ndarray) -> Transformed:
        
        frame1 = F.resize(frames[0], (self.height, self.width))
        frame2 = F.resize(frames[1], (self.height, self.width))
        optical_flow = transform.resize(optical_flow,
                                        (self.height, self.width))

        return (frame1, frame2), optical_flow


class RandomRotate(object): # 对图像光流都旋转同样角度

    def __init__(self, minmax: Union[Tuple[int, int], int]) -> None:
        self.minmax = minmax
        if isinstance(minmax, int):
            self.minmax = (-minmax, minmax)
    
    def __call__(self, 
                 frames: Tuple['Image', 'Image'], 
                 optical_flow: np.ndarray) -> Transformed:
        angle = random.randint(*self.minmax)
        frame1 = F.rotate(frames[0], angle)
        frame2 = F.rotate(frames[1], angle)
        optical_flow = transform.rotate(optical_flow, angle)

        return (frame1, frame2), optical_flow


class Normalize(object):  # # 对图像 标准化，减去均值，除以标准差

    def __init__(self, 
                 mean: Tuple[float, float, float],
                 std: Tuple[float, float, float]) -> None:
        self.mean = mean
        self.std = std
    
    def __call__(self, 
                 frames: Tuple[torch.Tensor, torch.Tensor], 
                 optical_flow: torch.Tensor) -> Transformed:

        frame1 = F.normalize(frames[0], self.mean, self.std)
        frame2 = F.normalize(frames[1], self.mean, self.std)

        return (frame1, frame2), optical_flow


class ToTensor(object): # 对图像和光流都转化为tensor

    def __call__(self, 
                 frames: Tuple['Image', 'Image'], 
                 optical_flow: np.ndarray) -> Transformed:
        
        return ((F.to_tensor(frames[0]), F.to_tensor(frames[1])), 
                 torch.from_numpy(optical_flow).permute(2, 0, 1).float())