Commit adbfb52c authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): align to_tensor

parent de45a333

imgutils/preprocess/pillow.py

+6 −12
Original line number Diff line number Diff line
@@ -190,27 +190,23 @@ def _create_center_crop(size): 




class PillowToTensor:

    def __init__(self):



        self.pil_modes = {'L', 'LA', 'P', 'I', 'F', 'RGB', 'YCbCr', 'RGBA', 'CMYK', '1'}



    def __call__(self, pic):

        if not isinstance(pic, Image.Image):

            raise TypeError('pic should be PIL Image. Got {}'.format(type(pic)))



        if pic.mode == 'I':

            # 32-bit signed integer pixels

            return np.array(pic, np.int32, copy=True)

            return np.array(pic, np.int32, copy=True)[None, ...]

        elif pic.mode == 'I;16':

            # 16-bit signed integer pixels

            return np.array(pic, np.int16, copy=True)

            return np.array(pic, np.int16, copy=True)[None, ...]

        elif pic.mode == 'F':

            # 32-bit floating point pixels

            return np.array(pic, np.float32, copy=True)

            return np.array(pic, np.float32, copy=True)[None, ...]



        img = np.array(pic, copy=True)

        if pic.mode == '1':

            return img.astype(np.float32)

            return img.astype(np.float32)[None, ...]

        elif pic.mode == 'L':

            img = img.reshape((1,) + img.shape)

            return img.astype(np.float32) / 255
@@ -219,9 +215,7 @@ class PillowToTensor: 
            img_a = img[..., 1].reshape((1,) + img.shape[:2])

            return np.concatenate((img_l, img_a), axis=0).astype(np.float32) / 255

        elif pic.mode == 'P':

            pic = pic.convert('RGB')

            img = np.array(pic, copy=True)

            img = img.transpose((2, 0, 1))

            img = np.array(pic, copy=True)[None, ...]

            return img.astype(np.float32) / 255

        elif pic.mode in ('RGB', 'YCbCr'):

            img = img.transpose((2, 0, 1))
@@ -233,7 +227,7 @@ class PillowToTensor: 
            img = img.transpose((2, 0, 1))

            return img.astype(np.float32) / 255



        raise ValueError(f"Unsupported PIL image mode: {pic.mode}")

        raise ValueError(f"Unsupported PIL image mode: {pic.mode}")  # pragma: no cover



    def __repr__(self) -> str:

        return f"{self.__class__.__name__}()"

test/preprocess/test_pillow.py

+29 −1
Original line number Diff line number Diff line
@@ -5,7 +5,7 @@ import pytest 
from PIL import Image

from hbutils.testing import tmatrix



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

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

from imgutils.preprocess.torchvision import _get_interpolation_mode

from test.testings import get_testfile
@@ -385,3 +385,31 @@ class TestPreprocessPillow: 
    ])

    def test_center_crop_repr(self, size, repr_text):

        assert repr(PillowCenterCrop(size=size)) == repr_text



    @pytest.mark.parametrize(*tmatrix({

        'src_image': [

            'png_640.png',

            'png_640_m90.png',

        ],

        'mode': [

            'I', 'I;16', 'F',

            '1', 'L', 'LA', 'P',

            'RGB', 'YCbCr', 'RGBA', 'CMYK',

        ]

    }))

    def test_to_tensor(self, src_image, mode):

        from torchvision.transforms import ToTensor

        image = Image.open(get_testfile(src_image))

        image = image.convert(mode)

        assert image.mode == mode

        ptotensor = PillowToTensor()

        ttotensor = ToTensor()

        np.testing.assert_array_almost_equal(ptotensor(image), ttotensor(image).numpy())



    def test_to_tensor_invalid(self):

        ptotensor = PillowToTensor()

        with pytest.raises(TypeError):

            _ = ptotensor(np.random.randn(3, 384, 384))



    def test_to_tensor_repr(self):

        return repr(PillowToTensor()) == 'PillowToTensor()'