mirror of https://github.com/alibaba/EasyCV.git
110 lines
3.3 KiB
Python
110 lines
3.3 KiB
Python
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import io
|
|
import logging
|
|
import time
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from cv2 import IMREAD_COLOR
|
|
from PIL import Image
|
|
|
|
from easycv import file
|
|
from easycv.framework.errors import IOError, KeyError, ValueError
|
|
from easycv.utils.constant import MAX_READ_IMAGE_TRY_TIMES
|
|
from .utils import is_oss_path, is_url_path
|
|
|
|
try:
|
|
from turbojpeg import TurboJPEG, TJCS_RGB, TJPF_BGR
|
|
turbo_jpeg = TurboJPEG()
|
|
turbo_jpeg_mode = {'RGB': TJCS_RGB, 'BGR': TJPF_BGR}
|
|
except:
|
|
turbo_jpeg = None
|
|
turbo_jpeg_mode = None
|
|
|
|
|
|
def load_image_with_pillow(content, mode='BGR', dtype=np.uint8):
|
|
with io.BytesIO(content) as buff:
|
|
image = Image.open(buff)
|
|
|
|
if mode.upper() != 'BGR':
|
|
if image.mode.upper() != mode.upper():
|
|
image = image.convert(mode.upper())
|
|
img = np.asarray(image, dtype=dtype)
|
|
else:
|
|
if image.mode.upper() != 'RGB':
|
|
image = image.convert('RGB')
|
|
img = np.asarray(image, dtype=dtype)
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
|
|
|
return img
|
|
|
|
|
|
def load_image_with_turbojpeg(content, mode='BGR', dtype=np.uint8):
|
|
assert mode.upper() in turbo_jpeg_mode
|
|
if turbo_jpeg is None or turbo_jpeg_mode is None:
|
|
raise ValueError(
|
|
'Please install turbojpeg by "pip install PyTurboJPEG" !')
|
|
|
|
img = turbo_jpeg.decode(
|
|
content, pixel_format=turbo_jpeg_mode[mode.upper()])
|
|
|
|
if img.dtype != dtype:
|
|
img = img.astype(dtype)
|
|
|
|
return img
|
|
|
|
|
|
def load_image_with_cv2(content, mode='BGR', dtype=np.uint8):
|
|
assert mode.upper() in ['BGR', 'RGB']
|
|
|
|
img_np = np.frombuffer(content, np.uint8)
|
|
img = cv2.imdecode(img_np, flags=IMREAD_COLOR)
|
|
|
|
if mode.upper() == 'RGB':
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
if img.dtype != dtype:
|
|
img = img.astype(dtype)
|
|
|
|
return img
|
|
|
|
|
|
def _load_image(fp, mode='BGR', dtype=np.uint8, backend='pillow'):
|
|
if backend == 'pillow':
|
|
img = load_image_with_pillow(fp, mode=mode, dtype=dtype)
|
|
elif backend == 'turbojpeg':
|
|
img = load_image_with_turbojpeg(fp, mode=mode, dtype=dtype)
|
|
elif backend == 'cv2':
|
|
img = load_image_with_cv2(fp, mode=mode, dtype=dtype)
|
|
else:
|
|
raise KeyError(
|
|
'Only support backend in ["pillow", "turbojpeg", "cv2"]')
|
|
return img
|
|
|
|
|
|
def load_image(img_path,
|
|
mode='BGR',
|
|
dtype=np.uint8,
|
|
backend='pillow',
|
|
max_try_times=MAX_READ_IMAGE_TRY_TIMES):
|
|
"""Load image file, return np.ndarray.
|
|
|
|
Args:
|
|
img_path (str): Image file path.
|
|
mode (str): Order of channel, candidates are `bgr` and `rgb`.
|
|
dtype : Output data type.
|
|
backend (str): The image decoding backend type. Options are `cv2`, `pillow`, `turbojpeg`.
|
|
"""
|
|
# TODO: functions of multi tries should be in the `io.open`
|
|
img = None
|
|
if is_url_path(img_path):
|
|
from mmcv.fileio.file_client import HTTPBackend
|
|
client = HTTPBackend()
|
|
img_bytes = client.get(img_path)
|
|
img = _load_image(img_bytes, mode=mode, dtype=dtype, backend=backend)
|
|
else:
|
|
with file.io.open(img_path, 'rb') as infile:
|
|
img = _load_image(
|
|
infile.read(), mode=mode, dtype=dtype, backend=backend)
|
|
|
|
return img
|