cv2
copied
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Readme
Files and versions
74 lines
2.6 KiB
74 lines
2.6 KiB
# Copyright 2021 Zilliz. 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.
|
|
|
|
|
|
import logging
|
|
from typing import Union
|
|
import cv2
|
|
import requests
|
|
import numpy as np
|
|
|
|
from towhee.operator import PyOperator
|
|
from towhee.types import Image
|
|
|
|
log = logging.getLogger()
|
|
|
|
|
|
class ImageDecodeCV2(PyOperator):
|
|
def __init__(self, mode='BGR'):
|
|
mode = mode.upper()
|
|
if mode not in ['BGR', 'RGB']:
|
|
raise RuntimeError("Mode only support BRG and RGB")
|
|
self._mode = mode
|
|
|
|
@staticmethod
|
|
def _load_from_remote(image_url: str) -> np.ndarray:
|
|
try:
|
|
r = requests.get(image_url, timeout=(20, 20))
|
|
if r.status_code // 100 != 2:
|
|
log.error('Download image from %s failed, error msg: %s, request code: %s ',
|
|
image_url, r.text, r.status_code)
|
|
return None
|
|
arr = np.asarray(bytearray(r.content), dtype=np.uint8)
|
|
return cv2.imdecode(arr, -1)
|
|
except Exception as e:
|
|
log.error('Download image from %s failed, error msg: %s', image_url, str(e))
|
|
return False
|
|
|
|
@staticmethod
|
|
def _load_from_local(image_path: str) -> np.ndarray:
|
|
return cv2.imread(image_path)
|
|
|
|
@staticmethod
|
|
def _from_bytes(image_bytes) -> np.ndarray:
|
|
arr = np.asarray(bytearray(image_bytes), dtype=np.uint8)
|
|
return cv2.imdecode(arr, -1)
|
|
|
|
def __call__(self, image_data: Union[str, bytes]):
|
|
if isinstance(image_data, bytes):
|
|
bgr_cv_image = ImageDecodeCV2._from_bytes(image_data)
|
|
elif image_data.startswith('http'):
|
|
bgr_cv_image = ImageDecodeCV2._load_from_remote(image_data)
|
|
else:
|
|
bgr_cv_image = ImageDecodeCV2._load_from_local(image_data)
|
|
|
|
if bgr_cv_image is None:
|
|
err = 'Read image %s failed' % image_data
|
|
log.error(err)
|
|
raise RuntimeError(err)
|
|
|
|
if self._mode == 'BGR':
|
|
return Image(bgr_cv_image, 'BGR')
|
|
rgb_cv_image = cv2.cvtColor(bgr_cv_image, cv2.COLOR_BGR2RGB)
|
|
return Image(rgb_cv_image, 'RGB')
|