wehub-resource-sync
180 lines
5.2 KiB
Python
180 lines
5.2 KiB
Python
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Any, Literal, TypeAlias
|
|
|
|
from PIL import Image
|
|
|
|
from paddle.utils import try_import
|
|
|
|
if TYPE_CHECKING:
|
|
import numpy.typing as npt
|
|
from PIL.Image import Image as PILImage
|
|
|
|
from paddle import Tensor
|
|
|
|
_ImageBackend: TypeAlias = Literal["pil", "cv2", "tensor"]
|
|
_ImageDataType: TypeAlias = Tensor | PILImage | npt.NDArray[Any]
|
|
|
|
__all__ = []
|
|
|
|
_image_backend: _ImageBackend = 'pil'
|
|
|
|
|
|
def set_image_backend(backend: _ImageBackend) -> None:
|
|
"""
|
|
Specifies the backend used to load images in class :ref:`api_paddle_datasets_ImageFolder`
|
|
and :ref:`api_paddle_datasets_DatasetFolder` . Now support backends are pillow and opencv.
|
|
If backend not set, will use 'pil' as default.
|
|
|
|
Args:
|
|
backend (str): Name of the image load backend, should be one of {'pil', 'cv2'}.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import os
|
|
>>> import shutil
|
|
>>> import tempfile
|
|
>>> import numpy as np
|
|
>>> from PIL import Image
|
|
|
|
>>> from paddle.vision import DatasetFolder
|
|
>>> from paddle.vision import set_image_backend
|
|
|
|
>>> set_image_backend('pil')
|
|
|
|
>>> def make_fake_dir():
|
|
... data_dir = tempfile.mkdtemp()
|
|
...
|
|
... for i in range(2):
|
|
... sub_dir = os.path.join(data_dir, 'class_' + str(i))
|
|
... if not os.path.exists(sub_dir):
|
|
... os.makedirs(sub_dir)
|
|
... for j in range(2):
|
|
... fake_img = Image.fromarray((np.random.random((32, 32, 3)) * 255).astype('uint8'))
|
|
... fake_img.save(os.path.join(sub_dir, str(j) + '.png'))
|
|
... return data_dir
|
|
|
|
>>> temp_dir = make_fake_dir()
|
|
|
|
>>> pil_data_folder = DatasetFolder(temp_dir)
|
|
|
|
>>> for items in pil_data_folder:
|
|
... break
|
|
|
|
>>> print(type(items[0]))
|
|
<class 'PIL.Image.Image'>
|
|
|
|
>>> # use opencv as backend
|
|
>>> set_image_backend('cv2')
|
|
|
|
>>> cv2_data_folder = DatasetFolder(temp_dir)
|
|
|
|
>>> for items in cv2_data_folder:
|
|
... break
|
|
|
|
>>> print(type(items[0]))
|
|
<class 'numpy.ndarray'>
|
|
|
|
>>> shutil.rmtree(temp_dir)
|
|
"""
|
|
global _image_backend
|
|
if backend not in ['pil', 'cv2', 'tensor']:
|
|
raise ValueError(
|
|
f"Expected backend are one of ['pil', 'cv2', 'tensor'], but got {backend}"
|
|
)
|
|
_image_backend = backend
|
|
|
|
|
|
def get_image_backend() -> _ImageBackend:
|
|
"""
|
|
Gets the name of the package used to load images
|
|
|
|
Returns:
|
|
str: backend of image load.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> from paddle.vision import get_image_backend
|
|
|
|
>>> backend = get_image_backend()
|
|
>>> print(backend)
|
|
pil
|
|
|
|
"""
|
|
return _image_backend
|
|
|
|
|
|
def image_load(
|
|
path: str, backend: _ImageBackend | None = None
|
|
) -> _ImageDataType | None:
|
|
"""Load an image.
|
|
|
|
Args:
|
|
path (str): Path of the image.
|
|
backend (str, optional): The image decoding backend type. Options are
|
|
`cv2`, `pil`, `None`. If backend is None, the global _imread_backend
|
|
specified by :ref:`api_paddle_vision_set_image_backend` will be used. Default: None.
|
|
|
|
Returns:
|
|
PIL.Image or np.array: Loaded image.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import numpy as np
|
|
>>> from PIL import Image
|
|
>>> from paddle.vision import image_load, set_image_backend
|
|
|
|
>>> fake_img = Image.fromarray((np.random.random((32, 32, 3)) * 255).astype('uint8'))
|
|
|
|
>>> path = 'temp.png'
|
|
>>> fake_img.save(path)
|
|
|
|
>>> set_image_backend('pil')
|
|
|
|
>>> pil_img = image_load(path).convert('RGB') # type: ignore
|
|
|
|
>>> print(type(pil_img))
|
|
<class 'PIL.Image.Image'>
|
|
|
|
>>> # use opencv as backend
|
|
>>> set_image_backend('cv2')
|
|
|
|
>>> np_img = image_load(path)
|
|
>>> print(type(np_img))
|
|
<class 'numpy.ndarray'>
|
|
|
|
"""
|
|
|
|
if backend is None:
|
|
backend = _image_backend
|
|
if backend not in ['pil', 'cv2', 'tensor']:
|
|
raise ValueError(
|
|
f"Expected backend are one of ['pil', 'cv2', 'tensor'], but got {backend}"
|
|
)
|
|
|
|
if backend == 'pil':
|
|
return Image.open(path)
|
|
elif backend == 'cv2':
|
|
cv2 = try_import('cv2')
|
|
return cv2.imread(path)
|