cv2-rgb
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
61 lines
2.1 KiB
61 lines
2.1 KiB
2 years ago
|
# 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
|
||
|
|
import cv2
|
||
|
|
import requests
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from towhee import register
|
||
|
|
from towhee.operator import PyOperator, OperatorFlag
|
||
|
|
from towhee._types import Image
|
||
|
|
|
||
|
|
log = logging.getLogger()
|
||
|
|
|
||
|
|
|
||
|
|
@register(output_schema=['img'],
|
||
|
|
flag=OperatorFlag.STATELESS | OperatorFlag.REUSEABLE)
|
||
|
|
class ImageDecodeCV2(PyOperator):
|
||
|
|
@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)
|
||
|
|
|
||
|
|
def __call__(self, image_path: str):
|
||
|
|
if image_path.startswith('http'):
|
||
|
|
bgr_cv_image = ImageDecodeCV2._load_from_remote(image_path)
|
||
|
|
else:
|
||
|
|
bgr_cv_image = ImageDecodeCV2._load_from_local(image_path)
|
||
|
|
if bgr_cv_image is None:
|
||
|
|
err = 'Read image %s failed' % image_path
|
||
|
|
log.error(err)
|
||
|
|
raise RuntimeError(err)
|
||
|
|
|
||
|
|
rgb = cv2.cvtColor(bgr_cv_image, cv2.COLOR_BGR2RGB)
|
||
|
|
return Image(rgb, 'RGB')
|