| import copy | |
| from typing import Dict, Any | |
| from transformers.image_processing_utils import ImageProcessingMixin, BatchFeature | |
| from timm.data.transforms_factory import create_transform | |
| class ClassificationImageProprocessor(ImageProcessingMixin): | |
| def __init__(self, data_config, **kwargs): | |
| super().__init__(**kwargs) | |
| self.data_config = data_config | |
| self._transform = create_transform(**data_config) | |
| def __call__(self, images, **kwargs) -> BatchFeature: | |
| """Preprocess an image or a batch of images.""" | |
| return self.preprocess(images, **kwargs) | |
| def preprocess(self, images, return_tensors=None, **kwargs) -> BatchFeature: | |
| images = [self._transform(image) for image in images] | |
| data = {"pixel_values": images} | |
| return BatchFeature(data=data, tensor_type=return_tensors) | |
| def to_dict(self) -> Dict[str, Any]: | |
| """ | |
| Serializes this instance to a Python dictionary. | |
| Returns: | |
| `Dict[str, Any]`: Dictionary of all the attributes that make up this image processor instance. | |
| """ | |
| output = copy.deepcopy(self.__dict__) | |
| output.pop("_transform", None) | |
| output["image_processor_type"] = self.__class__.__name__ | |
| return output |