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

# %% auto 0
__all__ = ['np2pil', 'pil2np', 'AlignOrientation']

# %% ../../nbs/03f_data.transforms.ipynb 2
from fastcore.all import *
from PIL import Image
import numpy as np
import polvo as pv

# %% ../../nbs/03f_data.transforms.ipynb 3
def np2pil(image): return Image.fromarray(image)

# %% ../../nbs/03f_data.transforms.ipynb 4
def pil2np(image): return np.array(image)

# %% ../../nbs/03f_data.transforms.ipynb 5
class AlignOrientation():
    def __init__(self, orientation='vertical', **kwargs):
        "Keep all images in same orientation"
        self.orientation = orientation
        
    def transform(self, x:np.ndarray):
        h, w = x.shape[:2]
        if (w>h and self.orientation=='vertical') or (h>w and self.orientation=='horizontal'):
            self._rotated = True
            return np.rot90(x)
        self._rotated = False
        return x
    
    def reverse(self, x:np.array):
        if self._rotated: return np.rot90(x)
        return x
