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

# %% ../../nbs/ffcv.utils.ipynb 3
from __future__ import annotations

import cv2
from enum import Enum
from pathlib import Path

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

from torch.utils.data import Dataset

from fastai.data.core import Datasets

from ffcv.fields import IntField, FloatField

from .fields import RGBImageField
from .writer import DatasetWriter

# %% auto 0
__all__ = ['LabelField', 'rgb_dataset_to_ffcv']

# %% ../../nbs/ffcv.utils.ipynb 4
class LabelField(Enum):
    int = 'int'
    float = 'float'

# %% ../../nbs/ffcv.utils.ipynb 5
def rgb_dataset_to_ffcv(
    dataset:Dataset|Datasets, # A PyTorch Dataset or single fastai Datasets
    write_path:str|Path, # File name to store dataset in FFCV beton format
    max_resolution:int|None=None, # If maximum side length is greater than `max_resolution`, resize so maximum side length equals `max_resolution`
    min_resolution:int|None=None, # If minimum side length is greater than `min_resolution`, resize so minimum side length equals `min_resolution`
    write_mode:str='raw', # `RGBImageField` write mode: 'raw', 'jpg', 'smart', 'proportion'
    smart_threshold:int|None=None, # If `write_mode='smart'`, JPEG-compress RAW bytes is larger than `smart_threshold`
    compress_probability:float=0.50, # Probability with which image is JPEG-compressed
    jpeg_quality:float=90, # Quality to use for jpeg compression if `write_mode='proportion'`
    interpolation=cv2.INTER_AREA, # OpenCV interpolation flag for resizing images with OpenCV
    resample=LANCZOS, # Pillow resampling filter for resizing images with Pillow
    num_workers:int=-1, # Number of workers to use. Defaults to number of CPUs
    chunk_size:int=100, # Size of chunks processed by each worker
    pillow_resize:bool=True, # Use Pillow to resize images instead of OpenCV
    label_field:LabelField=LabelField.int # Use FFCV `IntField` or `FloatField` for labels
):
    "Writes PyTorch/fastai compatible `dataset` into FFCV format at filepath `write_path`."
    if isinstance(write_path, Path):
        write_path = str(write_path)
    writer = DatasetWriter(write_path, {
        'image': RGBImageField(write_mode=write_mode, max_resolution=max_resolution,
                               min_resolution=min_resolution, smart_threshold=smart_threshold,
                               compress_probability=compress_probability, jpeg_quality=jpeg_quality,
                               interpolation=interpolation, resample=resample, pillow_resize=pillow_resize),
        'label': IntField() if label_field==LabelField.int else FloatField()
    }, num_workers=num_workers)
    writer.from_indexed_dataset(dataset, chunksize=chunk_size)
