mnistvit-transformers / image_processing_mnistvit.py

Upload folder using huggingface_hub

66be11f verified 26 days ago

5.34 kB

import numpy as np
import PIL.Image
import PIL.ImageOps
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from transformers.image_transforms import to_pil_image

# MNIST training set statistics, as in mnistvit.preprocess.preprocess_mnist
MNIST_MEAN = 0.1307
MNIST_STD = 0.3081
NORMALIZATIONS = ("image", "mnist", None)


class MnistViTImageProcessor(BaseImageProcessor):
    """Image processor reproducing the preprocessing of mnistvit.
    Converts to grayscale, center crops to a square, resizes and normalizes.  By
    default, it reproduces `mnistvit.preprocess.read_digit_image`: each image is
    standardized with its own mean and standard deviation, and dark-on-bright images are
    inverted so that digits are bright on a dark background as in MNIST.
    Args:
        size (int, optional): Height and width of the output images.  Default: 28.
        normalization (str or None, optional): `'image'` standardizes each image with
            its own mean and standard deviation, which is robust to arbitrary images.
            `'mnist'` normalizes with the MNIST training set mean and standard deviation
            as in `mnistvit.preprocess.preprocess_mnist`, which reproduces the original
            test set evaluation but assumes MNIST-like images.  `None` keeps pixel
            values in [0, 1], as expected by the ONNX model, which standardizes and
            inverts internally.  Default: `'image'`.
        do_auto_invert (bool, optional): Whether to invert images with more bright than
            dark pixels.  Only applies if `normalization` is `'image'`.
            Default: `True`.
    """

    model_input_names = ["pixel_values"]
    # Class defaults let transformers v5 serialize an explicit `normalization=None`
    size = 28
    normalization = "image"
    do_auto_invert = True

    def __init__(
        self,
        size: int = 28,
        normalization: str | None = "image",
        do_auto_invert: bool = True,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.size = size
        self.normalization = normalization
        self.do_auto_invert = do_auto_invert

    def __call__(self, images, **kwargs) -> BatchFeature:
        return self.preprocess(images, **kwargs)

    def preprocess(
        self,
        images,
        size: int | None = None,
        normalization: str | None = "default",
        do_auto_invert: bool | None = None,
        return_tensors=None,
        **kwargs,
    ) -> BatchFeature:
        """Preprocess an image or a list of images.
        Args:
            images (PIL.Image.Image or numpy.ndarray or torch.Tensor or list): The image
                or images to preprocess.
            size (int, optional): Overrides `self.size`.
            normalization (str or None, optional): Overrides `self.normalization`.
            do_auto_invert (bool, optional): Overrides `self.do_auto_invert`.
            return_tensors (str or transformers.TensorType, optional): Type of the
                returned tensors, e.g. `'pt'` or `'np'`.
        Returns:
            transformers.BatchFeature: `pixel_values` of shape
                `(batch_size, 1, size, size)`.
        """
        size = self.size if size is None else size
        # None is a valid normalization, so a sentinel marks the default
        if normalization == "default":
            normalization = self.normalization
        if normalization not in NORMALIZATIONS:
            raise ValueError(
                f"normalization must be one of {NORMALIZATIONS}, got {normalization!r}"
            )
        do_auto_invert = (
            self.do_auto_invert if do_auto_invert is None else do_auto_invert
        )
        if not isinstance(images, (list, tuple)):
            images = [images]
        pixel_values = np.stack(
            [
                self._preprocess_single(image, size, normalization, do_auto_invert)
                for image in images
            ]
        )
        return BatchFeature(
            data={"pixel_values": pixel_values}, tensor_type=return_tensors
        )

    @staticmethod
    def _preprocess_single(
        image, size: int, normalization: str | None, do_auto_invert: bool
    ) -> np.ndarray:
        if isinstance(image, PIL.Image.Image):
            image = PIL.ImageOps.exif_transpose(image)
        else:
            image = to_pil_image(image)
        image = image.convert("L")
        width, height = image.size
        crop_size = min(width, height)
        left = int(round((width - crop_size) / 2.0))
        top = int(round((height - crop_size) / 2.0))
        image = image.crop((left, top, left + crop_size, top + crop_size))
        image = image.resize((size, size), PIL.Image.Resampling.BILINEAR)
        data = np.asarray(image, dtype=np.float32)[np.newaxis] / 255.0
        if normalization == "image":
            # For robustness, standardize with image statistics rather than MNIST
            # statistics
            data = (data - data.mean()) / max(data.std(ddof=1), 1e-7)
            if do_auto_invert and np.count_nonzero(data > 0) > data.size / 2:
                # More bright than dark pixels
                data = -data
        elif normalization == "mnist":
            data = (data - MNIST_MEAN) / MNIST_STD
        return data.astype(np.float32)

	import numpy as np
	import PIL.Image
	import PIL.ImageOps
	from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
	from transformers.image_transforms import to_pil_image

	# MNIST training set statistics, as in mnistvit.preprocess.preprocess_mnist
	MNIST_MEAN = 0.1307
	MNIST_STD = 0.3081
	NORMALIZATIONS = ("image", "mnist", None)


	class MnistViTImageProcessor(BaseImageProcessor):
	"""Image processor reproducing the preprocessing of mnistvit.

	Converts to grayscale, center crops to a square, resizes and normalizes. By
	default, it reproduces `mnistvit.preprocess.read_digit_image`: each image is
	standardized with its own mean and standard deviation, and dark-on-bright images are
	inverted so that digits are bright on a dark background as in MNIST.

	Args:
	size (int, optional): Height and width of the output images. Default: 28.
	normalization (str or None, optional): `'image'` standardizes each image with
	its own mean and standard deviation, which is robust to arbitrary images.
	`'mnist'` normalizes with the MNIST training set mean and standard deviation
	as in `mnistvit.preprocess.preprocess_mnist`, which reproduces the original
	test set evaluation but assumes MNIST-like images. `None` keeps pixel
	values in [0, 1], as expected by the ONNX model, which standardizes and
	inverts internally. Default: `'image'`.
	do_auto_invert (bool, optional): Whether to invert images with more bright than
	dark pixels. Only applies if `normalization` is `'image'`.
	Default: `True`.
	"""

	model_input_names = ["pixel_values"]
	# Class defaults let transformers v5 serialize an explicit `normalization=None`
	size = 28
	normalization = "image"
	do_auto_invert = True

	def __init__(
	self,
	size: int = 28,
	normalization: str \| None = "image",
	do_auto_invert: bool = True,
	**kwargs,
	) -> None:
	super().__init__(**kwargs)
	self.size = size
	self.normalization = normalization
	self.do_auto_invert = do_auto_invert

	def __call__(self, images, **kwargs) -> BatchFeature:
	return self.preprocess(images, **kwargs)

	def preprocess(
	self,
	images,
	size: int \| None = None,
	normalization: str \| None = "default",
	do_auto_invert: bool \| None = None,
	return_tensors=None,
	**kwargs,
	) -> BatchFeature:
	"""Preprocess an image or a list of images.

	Args:
	images (PIL.Image.Image or numpy.ndarray or torch.Tensor or list): The image
	or images to preprocess.
	size (int, optional): Overrides `self.size`.
	normalization (str or None, optional): Overrides `self.normalization`.
	do_auto_invert (bool, optional): Overrides `self.do_auto_invert`.
	return_tensors (str or transformers.TensorType, optional): Type of the
	returned tensors, e.g. `'pt'` or `'np'`.

	Returns:
	transformers.BatchFeature: `pixel_values` of shape
	`(batch_size, 1, size, size)`.
	"""
	size = self.size if size is None else size
	# None is a valid normalization, so a sentinel marks the default
	if normalization == "default":
	normalization = self.normalization
	if normalization not in NORMALIZATIONS:
	raise ValueError(
	f"normalization must be one of {NORMALIZATIONS}, got {normalization!r}"
	)
	do_auto_invert = (
	self.do_auto_invert if do_auto_invert is None else do_auto_invert
	)
	if not isinstance(images, (list, tuple)):
	images = [images]
	pixel_values = np.stack(
	[
	self._preprocess_single(image, size, normalization, do_auto_invert)
	for image in images
	]
	)
	return BatchFeature(
	data={"pixel_values": pixel_values}, tensor_type=return_tensors
	)

	@staticmethod
	def _preprocess_single(
	image, size: int, normalization: str \| None, do_auto_invert: bool
	) -> np.ndarray:
	if isinstance(image, PIL.Image.Image):
	image = PIL.ImageOps.exif_transpose(image)
	else:
	image = to_pil_image(image)
	image = image.convert("L")
	width, height = image.size
	crop_size = min(width, height)
	left = int(round((width - crop_size) / 2.0))
	top = int(round((height - crop_size) / 2.0))
	image = image.crop((left, top, left + crop_size, top + crop_size))
	image = image.resize((size, size), PIL.Image.Resampling.BILINEAR)
	data = np.asarray(image, dtype=np.float32)[np.newaxis] / 255.0
	if normalization == "image":
	# For robustness, standardize with image statistics rather than MNIST
	# statistics
	data = (data - data.mean()) / max(data.std(ddof=1), 1e-7)
	if do_auto_invert and np.count_nonzero(data > 0) > data.size / 2:
	# More bright than dark pixels
	data = -data
	elif normalization == "mnist":
	data = (data - MNIST_MEAN) / MNIST_STD
	return data.astype(np.float32)