support decode from bytes

Signed-off-by: junjiejiangjjj <junjie.jiang@zilliz.com>

__init__.py

@@ -16,5 +16,5 @@ from .image_decode_cv2 import ImageDecodeCV2
 
 
 # The factory method
-def cv2():
-    return ImageDecodeCV2()
+def cv2(*args, **kwargs):
+    return ImageDecodeCV2(*args, **kwargs)

image_decode_cv2.py

@@ -14,22 +14,23 @@
 
 
 import logging
+from typing import Union
 import cv2
 import requests
 import numpy as np
 
-from towhee import register
-from towhee.operator import PyOperator, OperatorFlag
-from towhee._types import Image
+from towhee.operator import PyOperator
+from towhee.types import Image
 
 log = logging.getLogger()
 
 
-@register(output_schema=['img'],
-          flag=OperatorFlag.STATELESS | OperatorFlag.REUSEABLE)
 class ImageDecodeCV2(PyOperator):
-    def __init__(self):
-        pass
+    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:
@@ -49,14 +50,25 @@ class ImageDecodeCV2(PyOperator):
     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)
+    @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_path)
+            bgr_cv_image = ImageDecodeCV2._load_from_local(image_data)
+
         if bgr_cv_image is None:
-            err = 'Read image %s failed' % image_path
+            err = 'Read image %s failed' % image_data
             log.error(err)
             raise RuntimeError(err)
 
-        return Image(bgr_cv_image, 'BGR')
+        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')