# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/ffcv.inference.ipynb.

# %% ../../nbs/ffcv.inference.ipynb 1
# Contains code from:
# fastai - Apache License 2.0 - Copyright (c) 2023 fast.ai
# FFCV - Apache License 2.0 - Copyright (c) 2022 FFCV Team

# %% ../../nbs/ffcv.inference.ipynb 4
from __future__ import annotations

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

from fastcore.basics import Path, merge
from fastcore.transform import DisplayedTransform

from fastai.data.block import TransformBlock
from fastai.data.transforms import IntToFloatTensor, ToTensor
from fastai.vision.core import image2tensor

from ..utils import pil_to_numpy
from ..imports import *

# %% auto 0
__all__ = ['FFAIBase', 'FFAIImage', 'FFAIImageBW', 'FFAIMask', 'encodes', 'FFAIImageBlock', 'FFAICenterCrop']

# %% ../../nbs/ffcv.inference.ipynb 7
def load_image(fn:str|Path, mode=None):
    "Open and load an image with OpenCV, optionally converting to `mode` and RGB order"
    if isinstance(fn, Path):
        fn = str(fn)
    if mode is None:
        return cv2.imread(fn)
    else:
        if mode == cv2.IMREAD_COLOR:
            return cv2.cvtColor(cv2.imread(fn, mode), cv2.COLOR_BGR2RGB)
        else:
            return cv2.imread(fn, mode)

# %% ../../nbs/ffcv.inference.ipynb 8
class FFAIBase(np.ndarray):
    "Base class for a FFCV image that can show itself and convert to a Tensor"
    # _bypass_type = np.ndarray
    _show_args = {'cmap':'viridis'}
    _open_args = {'mode': cv2.IMREAD_COLOR }

    @classmethod
    def create(cls, fn:Path|str|Tensor|np.ndarray|Image.Image, **kwargs):
        "Return an Image from `fn`"
        if isinstance(fn, TensorImage):
            fn = fn.permute(1,2,0).type(torch.uint8)
        if isinstance(fn, TensorMask):
            fn = fn.type(torch.uint8)
        if isinstance(fn, Tensor):
            fn = fn.numpy()
        if isinstance(fn, np.ndarray):
            return cls(fn)
        if isinstance(fn, Image.Image):
            return cls(pil_to_numpy(fn))
        return cls(load_image(fn, **merge(cls._open_args, kwargs)))

    def show(self, ctx=None, **kwargs):
        "Show image using `merge(self._show_args, kwargs)`"
        return show_image(self, ctx=ctx, **merge(self._show_args, kwargs))

    def __new__(cls, input_array):
        obj = np.asarray(input_array).view(cls)
        return obj

    def __array_finalize__(self, obj):
        if obj is None: return

    def __repr__(self):
        return f'{self.__class__.__name__} shape={"x".join([str(d) for d in self.shape])}'

# %% ../../nbs/ffcv.inference.ipynb 9
class FFAIImage(FFAIBase):
    "A RGB FFCV image that can show itself and converts to `TensorImage`"
    _tensor_cls = TensorImage

# %% ../../nbs/ffcv.inference.ipynb 11
class FFAIImageBW(FFAIBase):
    "A BW FFCV image that can show itself and converts to `TensorImageBW`"
    _show_args = {'cmap':'Greys'}
    _open_args = {'mode': cv2.IMREAD_GRAYSCALE}
    _tensor_cls = TensorImageBW

# %% ../../nbs/ffcv.inference.ipynb 15
class FFAIMask(FFAIBase):
    "A FFCV image mask that can show itself and converts to `TensorMask`"
    _show_args = {'cmap':'Greys'}
    _open_args = {'mode': cv2.IMREAD_GRAYSCALE}
    _tensor_cls = TensorMask

# %% ../../nbs/ffcv.inference.ipynb 17
@ToTensor
def encodes(self, o:FFAIBase):
    return o._tensor_cls(image2tensor(o))

@ToTensor
def encodes(self, o:FFAIMask):
    return o._tensor_cls(image2tensor(o)[0])

# %% ../../nbs/ffcv.inference.ipynb 21
def FFAIImageBlock(cls:FFAIBase=FFAIImage):
    "A `TransformBlock` for images of `cls`"
    return TransformBlock(type_tfms=cls.create, batch_tfms=IntToFloatTensor)

# %% ../../nbs/ffcv.inference.ipynb 23
def get_center_crop(height, width, _, ratio):
    s = min(height, width)
    c = int(ratio * s)
    delta_h = (height - c) // 2
    delta_w = (width - c) // 2

    return delta_h, delta_w, c

# %% ../../nbs/ffcv.inference.ipynb 24
class FFAICenterCrop(DisplayedTransform):
    split_idx,order = None,1
    "Resize `FFAIBase` image to `output_size` using `ratio`"
    def __init__(self,
        output_size:int, # Square size to resize and crop image to
        ratio:float, # Ratio of (crop size) / (min side length)
        interpolation=(cv2.INTER_AREA, cv2.INTER_NEAREST), # OpenCV interpolation mode, resamples[1] for mask
    ):
        super().__init__()
        assert isinstance(output_size, int), f"{output_size=} must be an int"
        self.size = (output_size, output_size)
        self.ratio = ratio
        self.interp,self.mask_interp = interpolation

    def encodes(self, x:FFAIBase):
        height, width, _ = x.shape
        s = min(height, width)
        c = int(self.ratio * s)
        h = (height - c) // 2
        w = (width - c) // 2
        return cv2.resize(x[h:h+c, w:w+c], self.size, interpolation=self.mask_interp if isinstance(x,FFAIMask) else self.interp)
