Commit fe3cbb93 authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): add pillow check

parent e3f72c50

imgutils/preprocess/__init__.py

+2 −1
Original line number Diff line number Diff line 
from .pillow import register_pillow_transform, create_pillow_transforms

from .pillow import register_pillow_transform, create_pillow_transforms, \

    register_pillow_parse, parse_pillow_transforms

from .torchvision import register_torchvision_transform, create_torchvision_transforms, \

    register_torchvision_parse, parse_torchvision_transforms

imgutils/preprocess/pillow.py

+84 −0
Original line number Diff line number Diff line 
import copy

import io

from functools import wraps

from textwrap import indent

from typing import Union, Optional, Tuple, List



import numpy as np

from PIL import Image



from .base import NotParseTarget



# noinspection PyUnresolvedReferences

_INT_TO_PILLOW = {

    0: Image.NEAREST,
@@ -56,6 +59,24 @@ def register_pillow_transform(name: str): 
    return _fn





_PTRANS_PARSERS = {}





def register_pillow_parse(name: str):

    def _fn(func):

        @wraps(func)

        def _new_func(*args, **kwargs):

            return {

                'type': name,

                **func(*args, **kwargs),

            }



        _PTRANS_PARSERS[name] = _new_func

        return _new_func



    return _fn





class PillowResize:

    # noinspection PyUnresolvedReferences

    def __init__(
@@ -135,6 +156,19 @@ def _create_resize(size, interpolation='bilinear', max_size=None, antialias=True 
    )





@register_pillow_parse('resize')

def _parse_resize(obj: PillowResize):

    if not isinstance(obj, PillowResize):

        raise NotParseTarget



    return {

        'size': obj.size,

        'interpolation': _PILLOW_TO_STR[obj.interpolation],

        'max_size': obj.max_size,

        'antialias': obj.antialias,

    }





class PillowCenterCrop:

    def __init__(self, size):

        if isinstance(size, int):
@@ -189,6 +223,16 @@ def _create_center_crop(size): 
    )





@register_pillow_parse('center_crop')

def _parse_center_crop(obj: PillowCenterCrop):

    if not isinstance(obj, PillowCenterCrop):

        raise NotParseTarget



    return {

        'size': list(obj.size),

    }





class PillowToTensor:

    def __call__(self, pic):

        if not isinstance(pic, Image.Image):
@@ -238,6 +282,14 @@ def _create_to_tensor(): 
    return PillowToTensor()





@register_pillow_parse('to_tensor')

def _parse_to_tensor(obj: PillowToTensor):

    if not isinstance(obj, PillowToTensor):

        raise NotParseTarget



    return {}





class PillowMaybeToTensor:

    def __call__(self, image):

        if isinstance(image, np.ndarray):
@@ -254,6 +306,14 @@ def _create_maybe_to_tensor(): 
    return PillowMaybeToTensor()





@register_pillow_parse('maybe_to_tensor')

def _parse_maybe_to_tensor(obj: PillowMaybeToTensor):

    if not isinstance(obj, PillowMaybeToTensor):

        raise NotParseTarget



    return {}





class PillowNormalize:

    def __init__(self, mean, std, inplace=False):

        if isinstance(mean, (list, tuple)):
@@ -300,6 +360,17 @@ def _create_normalize(mean, std, inplace=False): 
    )





@register_pillow_parse('normalize')

def _parse_normalize(obj: PillowNormalize):

    if not isinstance(obj, PillowNormalize):

        raise NotParseTarget



    return {

        'mean': obj.mean.tolist(),

        'std': obj.std.tolist(),

    }





class PillowCompose:

    def __init__(self, transforms):

        self.transforms = transforms
@@ -328,3 +399,16 @@ def create_pillow_transforms(tvalue: Union[list, dict]): 
        return _PTRANS_CREATORS[ttype](**tvalue)

    else:

        raise TypeError(f'Unknown type of transforms - {tvalue!r}.')





def parse_pillow_transforms(value):

    if isinstance(value, PillowCompose):

        return [parse_pillow_transforms(trans) for trans in value.transforms]

    else:

        for key, parser in _PTRANS_PARSERS.items():

            try:

                return parser(value)

            except NotParseTarget:

                pass



        raise TypeError(f'Unknown parse transform - {value!r}.')

test/preprocess/test_pillow.py

+52 −1
Original line number Diff line number Diff line
@@ -6,7 +6,7 @@ from PIL import Image 
from hbutils.testing import tmatrix



from imgutils.preprocess.pillow import PillowResize, _get_pillow_resample, PillowCenterCrop, PillowToTensor, \

    PillowMaybeToTensor, PillowNormalize, create_pillow_transforms

    PillowMaybeToTensor, PillowNormalize, create_pillow_transforms, parse_pillow_transforms, PillowCompose

from imgutils.preprocess.torchvision import _get_interpolation_mode

from test.testings import get_testfile
@@ -816,3 +816,54 @@ class TestPreprocessPillow: 
    def test_normalize_repr(self, mean, std, repr_text):

        pnormalize = PillowNormalize(mean, std)

        assert repr(pnormalize) == repr_text



    def test_parse_pillow_transforms(self):

        assert parse_pillow_transforms(PillowCompose([

            PillowResize(size=384, interpolation=Image.BICUBIC, max_size=None, antialias=True),

            PillowCenterCrop(size=[384, 384]),

            PillowMaybeToTensor(),

            PillowNormalize(mean=[0.5000, 0.5000, 0.5000], std=[0.5000, 0.5000, 0.5000]),

        ])) == [

                   {'antialias': True,

                    'interpolation': 'bicubic',

                    'max_size': None,

                    'size': 384,

                    'type': 'resize'},

                   {'size': [384, 384], 'type': 'center_crop'},

                   {'type': 'maybe_to_tensor'},

                   {'mean': [0.5, 0.5, 0.5], 'std': [0.5, 0.5, 0.5], 'type': 'normalize'}

               ]



        assert parse_pillow_transforms(PillowCompose([

            PillowResize(size=384, interpolation=Image.BICUBIC, max_size=None, antialias=True),

            PillowCenterCrop(size=[384, 384]),

            PillowToTensor(),

            PillowNormalize(mean=[0.5000, 0.5000, 0.5000], std=[0.5000, 0.5000, 0.5000]),

        ])) == [

                   {'antialias': True,

                    'interpolation': 'bicubic',

                    'max_size': None,

                    'size': 384,

                    'type': 'resize'},

                   {'size': [384, 384], 'type': 'center_crop'},

                   {'type': 'to_tensor'},

                   {'mean': [0.5, 0.5, 0.5], 'std': [0.5, 0.5, 0.5], 'type': 'normalize'}

               ]



        assert parse_pillow_transforms(

            PillowResize(size=384, interpolation=Image.BICUBIC, max_size=None, antialias=True)) \

               == {'antialias': True,

                   'interpolation': 'bicubic',

                   'max_size': None,

                   'size': 384,

                   'type': 'resize'}

        assert parse_pillow_transforms(PillowCenterCrop(size=[384, 384])) == {'size': [384, 384], 'type': 'center_crop'}

        assert parse_pillow_transforms(PillowToTensor()) == {'type': 'to_tensor'}

        assert parse_pillow_transforms(

            PillowNormalize(mean=[0.5000, 0.5000, 0.5000], std=[0.5000, 0.5000, 0.5000])) \

               == {'mean': [0.5, 0.5, 0.5], 'std': [0.5, 0.5, 0.5], 'type': 'normalize'}



        with pytest.raises(TypeError):

            _ = parse_pillow_transforms(None)

        with pytest.raises(TypeError):

            _ = parse_pillow_transforms(23344)

test/preprocess/test_pillow_compose.py

+51 −3
Original line number Diff line number Diff line
@@ -7,8 +7,8 @@ from PIL import Image 
from hbutils.testing import tmatrix



from imgutils.preprocess.pillow import PillowNormalize, PillowCompose, PillowResize, PillowMaybeToTensor, \

    PillowCenterCrop, create_pillow_transforms

from imgutils.preprocess.torchvision import create_torchvision_transforms

    PillowCenterCrop, create_pillow_transforms, parse_pillow_transforms

from imgutils.preprocess.torchvision import create_torchvision_transforms, parse_torchvision_transforms

from test.testings import get_testfile



try:
@@ -216,7 +216,6 @@ PillowCompose( 
            ttrans(image).numpy(),

        )





    def test_create_transform_invalid(self):

        with pytest.raises(TypeError):

            _ = create_pillow_transforms(None)
@@ -225,3 +224,52 @@ PillowCompose( 
        with pytest.raises(TypeError):

            _ = create_pillow_transforms('str')



    @skipUnless(_TORCHVISION_AVAILABLE, 'Torchvision required.')

    @pytest.mark.parametrize(*tmatrix({

        'src_image': [

            'png_640.png',

            'png_640_m90.png',

        ],

        'meta_name': [

            'mobilenetv4_conv_large.e600_r384_in1k',

            'caformer_s36.sail_in1k_384',

            'beit_base_patch16_384.in22k_ft_in22k_in1k',

            'resnet101d.ra2_in1k',

        ]

    }))

    def test_compose_parse_and_create_alignment_p_t(self, src_image, meta_name, meta_collect):

        image = Image.open(get_testfile(src_image))

        meta = meta_collect[meta_name]



        ptrans = create_pillow_transforms(meta)

        parsed_meta = parse_pillow_transforms(ptrans)

        ttrans = create_torchvision_transforms(parsed_meta)

        np.testing.assert_array_almost_equal(

            ptrans(image),

            ttrans(image).numpy(),

        )



    @skipUnless(_TORCHVISION_AVAILABLE, 'Torchvision required.')

    @pytest.mark.parametrize(*tmatrix({

        'src_image': [

            'png_640.png',

            'png_640_m90.png',

        ],

        'meta_name': [

            'mobilenetv4_conv_large.e600_r384_in1k',

            'caformer_s36.sail_in1k_384',

            'beit_base_patch16_384.in22k_ft_in22k_in1k',

            'resnet101d.ra2_in1k',

        ]

    }))

    def test_compose_parse_and_create_alignment_t_p(self, src_image, meta_name, meta_collect):

        image = Image.open(get_testfile(src_image))

        meta = meta_collect[meta_name]



        ttrans = create_torchvision_transforms(meta)

        parsed_meta = parse_torchvision_transforms(ttrans)

        ptrans = create_pillow_transforms(parsed_meta)

        np.testing.assert_array_almost_equal(

            ptrans(image),

            ttrans(image).numpy(),

        )