# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/03f_data.transforms.ipynb.

# %% auto 0
__all__ = ['FlipVert', 'AlignOrientation']

# %% ../../nbs/03f_data.transforms.ipynb 4
from fastai.vision.all import *
import polvo as pv
import fastpolvo as fpv

# %% ../../nbs/03f_data.transforms.ipynb 5
class FlipVert(RandTransform):
    order = 30
    @delegates(RandTransform.__init__)
    def __init__(self, p=0.5, **kwargs):
        "Batch transform for flipping images vertically"
        super().__init__(p=p, **kwargs)
        
    def encodes(self, x:TensorImage): return torch.flip(x, [2]) # (BS, C, H, W)
    def encodes(self, x:TensorMask): return torch.flip(x, [1]) # only 3 channels (BS, H, W)

# %% ../../nbs/03f_data.transforms.ipynb 8
class AlignOrientation(Transform):
    @delegates(Transform.__init__)
    def __init__(self, orientation='vertical', **kwargs):
        "Keep all images in same orientation"
        self.orientation = orientation
        super().__init__(**kwargs)
        
    def encodes(self, x:Image.Image):
        w, h = x.size
        if w>h and self.orientation=='vertical': x = x.rotate(90, expand=True)
        elif h>w and self.orientation=='horizontal': x = x.rotate(90, expand=True)
        return x
