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

# %% auto 0
__all__ = ['RGBImageField']

# %% ../../nbs/ffcv.fields.ipynb 1
# Contains code from:
# FFCV - Apache License 2.0 - Copyright (c) 2022 FFCV Team

# %% ../../nbs/ffcv.fields.ipynb 4
import cv2
import numpy as np
from typing import Type
from numba.typed import Dict
from PIL import Image

try:
    LANCZOS = Image.Resampling.LANCZOS
except AttributeError:
    from PIL.Image import LANCZOS

from ffcv.fields.base import Field, ARG_TYPE
from ffcv.fields.rgb_image import SimpleRGBImageDecoder, encode_jpeg, IMAGE_MODES
from ffcv.fields.rgb_image import RGBImageField as _RGBImageField
from ffcv.pipeline.operation import Operation
from ..utils import pil_to_numpy

# %% ../../nbs/ffcv.fields.ipynb 28
def resizer(image, max_resolution, min_resolution, interpolation=(cv2.INTER_AREA, LANCZOS)):
    pillow_resize = isinstance(image, Image.Image)
    if max_resolution is None and min_resolution is None:
        return pil_to_numpy(image) if pillow_resize else image

    if pillow_resize:
        original_size = np.array([image.size[0], image.size[1]])
    else:
        original_size = np.array([image.shape[1], image.shape[0]])

    if max_resolution is not None:
        ratio = max_resolution / original_size.max()
    elif min_resolution is not None:
        ratio = min_resolution / original_size.min()
    else:
        ratio = 1

    if ratio < 1:
        new_size = (ratio * original_size).astype(int)
        if pillow_resize:
            image = image.resize(new_size, resample=interpolation[1])
        else:
            image = cv2.resize(image, tuple(new_size), interpolation=interpolation[0])
    if pillow_resize:
        return pil_to_numpy(image)
    else:
        return image

# %% ../../nbs/ffcv.fields.ipynb 30
class RGBImageField(_RGBImageField):
    """
    A subclass of :class:`~ffcv.fields.Field` supporting RGB image data.

    Parameters
    ----------
    write_mode : str, optional
        How to write the image data to the dataset file. Should be either 'raw'
        (``uint8`` pixel values), 'jpg' (compress to JPEG format), 'smart'
        (decide between saving pixel values and JPEG compressing based on image
        size), and 'proportion' (JPEG compress a random subset of the data with
        size specified by the ``compress_probability`` argument). By default: 'raw'.
    max_resolution : int, optional
        If specified, resize images to have maximum side length equal to this 
        value if maximum side length is larger. By default: None
    min_resolution : int, optional
        If specified, resize images to have minimum side length equal to this 
        value if minimum side length is larger. By default: None
    smart_threshold : int, optional
        When `write_mode='smart`, will compress an image if RAW byte size is
        larger than `smart_threshold`.
    jpeg_quality : int, optional
        The quality parameter for JPEG encoding (ignored for ``write_mode='raw'``). 
        By default 90
    compress_probability : float, optional
        Ignored unless ``write_mode='proportion'``; in the latter case it is the
        probability with which image is JPEG-compressed. By default 0.5.
    interpolation : optional
        The OpenCV interpolation flag for resizing images with OpenCV. 
        By default INTER_AREA.
    resample : optional
        The Pillow resampling filter for resizing images with Pillow. By default LANCZOS.
    pillow_resize : bool, optional
        Use Pillow to resize images instead of OpenCV. By default False (OpenCV).
    """
    def __init__(self, write_mode='raw', max_resolution: int = None,
                min_resolution: int = None, smart_threshold: int = None,
                jpeg_quality: int = 90, compress_probability: float = 0.5,
                interpolation = cv2.INTER_AREA, resample = LANCZOS,
                pillow_resize:bool = False) -> None:
        self.write_mode = write_mode
        self.smart_threshold = smart_threshold
        self.max_resolution = max_resolution
        self.min_resolution = min_resolution
        self.jpeg_quality = int(jpeg_quality)
        self.proportion = compress_probability
        self.interpolation = interpolation
        self.resample = resample
        self.pillow_resize = pillow_resize
        if max_resolution is not None and min_resolution is not None:
            raise ValueError(f'Can only set one of {max_resolution=} or {min_resolution=}')


    def encode(self, destination, image, malloc):
        if not isinstance(image, np.ndarray) and not isinstance(image, Image.Image):
            raise TypeError(f"Unsupported image type {type(image)}")

        if self.pillow_resize:
            if isinstance(image, np.ndarray):
                image = Image.fromarray(image)
        else:
            if isinstance(image, Image.Image):
                image = pil_to_numpy(image)

        image = resizer(image, self.max_resolution, self.min_resolution,
                        (self.interpolation, self.resample))

        if len(image.shape) > 2 and image.shape[2] != 3:
            raise ValueError(f"Invalid shape for rgb image: {image.shape}")

        if image.dtype != np.uint8:
            raise ValueError("Image type has to be uint8")

        write_mode = self.write_mode
        as_jpg = None

        if write_mode == 'smart':
            as_jpg = encode_jpeg(image, self.jpeg_quality)
            write_mode = 'raw'
            if self.smart_threshold is not None:
                if image.nbytes > self.smart_threshold:
                    write_mode = 'jpg'
        elif write_mode == 'proportion':
            if np.random.rand() < self.proportion:
                write_mode = 'jpg'
            else:
                write_mode = 'raw'

        destination['mode'] = IMAGE_MODES[write_mode]
        destination['height'], destination['width'] = image.shape[:2]

        if write_mode == 'jpg':
            if as_jpg is None:
                as_jpg = encode_jpeg(image, self.jpeg_quality)
            destination['data_ptr'], storage = malloc(as_jpg.nbytes)
            storage[:] = as_jpg
        elif write_mode == 'raw':
            image_bytes = np.ascontiguousarray(image).view('<u1').reshape(-1)
            destination['data_ptr'], storage = malloc(image.nbytes)
            storage[:] = image_bytes
        else:
            raise ValueError(f"Unsupported write mode {self.write_mode}")
