# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/vision.data.ipynb (unless otherwise specified).


from __future__ import annotations


__all__ = ['PreBatchAsItem', 'PostBatchAsItem', 'ImageCPUBlock', 'MaskCPUBlock']

# Cell
#nbdev_comment from __future__ import annotations

from fastcore.transform import DisplayedTransform

from fastai.data.block import TransformBlock
from fastai.data.transforms import IntToFloatTensor
from fastai.vision.core import PILBase, PILImage, PILMask, TensorPoint, TensorBBox, AddMaskCodes

from ..imports import *

# Cell
class PreBatchAsItem(DisplayedTransform):
    "Converts Tensor from CHW to BCHW by adding a fake B dim"
    order=11
    def encodes(self, x:TensorImage|TensorMask|TensorPoint|TensorBBox):
        return x.unsqueeze(0)
    def decodes(self, x:TensorImage|TensorMask|TensorPoint|TensorBBox):
        return x.squeeze(0)

# Cell
class PostBatchAsItem(DisplayedTransform):
    "Converts Tensor from BCHW to CHW by removing the fake B dim"
    order=98
    def encodes(self, x:TensorImage|TensorMask|TensorPoint|TensorBBox):
        return x.squeeze(0)
    def decodes(self, x:TensorImage|TensorMask|TensorPoint|TensorBBox):
        return x.unsqueeze(0)

# Cell
def ImageCPUBlock(cls:PILBase=PILImage):
    "A `TransformBlock` for images of `cls` for running batch_tfms on CPU"
    return TransformBlock(type_tfms=cls.create,
                          item_tfms=[IntToFloatTensor, PreBatchAsItem, PostBatchAsItem])

# Cell
def MaskCPUBlock(codes:listy|None=None):
    "A `TransformBlock` for segmentation masks, potentially with `codes`, for running batch_tfms on CPU"
    return TransformBlock(type_tfms=PILMask.create,
                          item_tfms=[AddMaskCodes(codes=codes), IntToFloatTensor, PreBatchAsItem, PostBatchAsItem])