wucng/custom-resnet18

1import os2 3import albumentations as A4from albumentations.pytorch.transforms import ToTensorV25import PIL.Image6import numpy as np7from functools import partial8from typing import Dict, List, Optional, Union9from datasets import load_dataset, DatasetDict, Image  # pip install datasets10from torch.utils.data import DataLoader11import torch12 13# 自定义 ImageProcessor 为了与 pipeline使用14from transformers import ViTImageProcessor15from transformers.image_utils import PILImageResampling, ChannelDimension16from transformers.image_processing_utils import get_size_dict17 18class ResnetImageProcessor(ViTImageProcessor):19    """20    >>> # tfs = A.Compose([A.Resize(256, 256), A.CenterCrop(224, 224)])21    >>> # 如果传入 参数 tfs=tfs 在调用save_pretrained会报错22    >>> # 本地使用23    >>> mean = [0.485, 0.456, 0.406];std = [0.229, 0.224, 0.225]24    >>> image_processor = ResnetImageProcessor(size=(224, 224), image_mean=mean, image_std=std)25    >>> image_processor.save_pretrained("custom-resnet")26    >>> image_processor = ResnetImageProcessor.from_pretrained("custom-resnet")27 28    >>> # push_to_hub29    >>> # hub登录30    >>> from huggingface_hub import notebook_login;notebook_login()31    >>> # or huggingface-cli login32 33    >>> ResnetImageProcessor.register_for_auto_class()34    >>> mean = [0.485, 0.456, 0.406];std = [0.229, 0.224, 0.225]35    >>> image_processor = ResnetImageProcessor(size=(224, 224), image_mean=mean, image_std=std)36    >>> image_processor.save_pretrained("custom-resnet")37    >>> # image_processor = ResnetImageProcessor.from_pretrained("custom-resnet")38    >>> # 如果要执行 push_to_hub 需要将 custom-resnet/preprocessor_config.json 中的 "image_processor_type" 改成 "ViTImageProcessor"39    >>> # 默认的 ResnetImageProcessor 没有注册到 AutoImageProcessor40    >>> # 否则从 使用 AutoImageProcessor 加载 会报错了41    >>> image_processor.push_to_hub('custom-resnet')42 43    >>> # 从 huggingface hub 加载44    >>> from transformers import AutoImageProcessor45    >>> AutoImageProcessor.register(config_class='wucng/custom-resnet/config.json',image_processor_class=ResnetImageProcessor)46    >>> image_processor = AutoImageProcessor.from_pretrained('wucng/custom-resnet', trust_remote_code=True)47    """48 49    def resize(50            self,51            image: np.ndarray,52            size: Dict[str, int],53            resample: PILImageResampling = PILImageResampling.BILINEAR,54            data_format: Optional[Union[str, ChannelDimension]] = None,55            input_data_format: Optional[Union[str, ChannelDimension]] = None,56            **kwargs,57    ) -> np.ndarray:58        size = get_size_dict(size)59        output_size = (size["height"], size["width"])60        height, width = size["height"], size["width"]61 62        tfs = kwargs.get('tfs', None)63        if tfs is None:64            ratio = 256 / 22465            tfs = A.Compose([A.Resize(int(ratio * height), int(ratio * width)), A.CenterCrop(height, width)])66        return tfs(image=image)['image']67