Add operator for object detection with yolo

setup.py

@@ -53,15 +53,17 @@ def run(self):
     test_suite="setup.test_suite",
     cmdclass={'install': PostInstallCommand},
     install_requires=[
-        'torch>=1.2.0',
-        'torchvision>=0.4.0',
+        'torch>=1.7.0',
+        'torchvision>=0.8.1',
         'numpy>=1.19.5',
         'pandas>=1.1.5',
         'pyyaml>=5.3.0',
         'requests>=2.12.5',
         'tqdm>=4.59.0',
         'pillow>=8.3.1',
         'scipy>=1.5.3',
+        'matplotlib>=3.2.2',
+        'opencv-python>=4.1.2',
     ],
     packages=find_packages(),
     package_data={'towhee.tests.test_util': ['*.yaml']},

towhee/tests/mock_operators/pytorch_yolo_detector_operator/__init__.py

@@ -0,0 +1,13 @@
+# 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.

towhee/tests/mock_operators/pytorch_yolo_detector_operator/pytorch_yolo_detector_operator.py

@@ -0,0 +1,51 @@
+# 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 torch
+import cv2
+from PIL import Image
+from typing import NamedTuple
+from towhee.operator import Operator
+
+
+class PytorchYoloDetectorOperator(Operator):
+    """
+    Stateful operator
+    """
+
+    def __init__(self, model_name) -> None:
+        super().__init__()
+        self._model = torch.hub.load("ultralytics/yolov5", model_name, pretrained=True)
+
+    def __call__(self, img_path: str) -> NamedTuple("Outputs", [("objs_list", list)]):
+        self.img_path = img_path
+        # Get object detection results in yolov5 model
+        results = self._model([img_path])
+        self.bboxes = results.xyxy[0].tolist()
+        # Get the detected objects list([PIL.Image])
+        objs_list = self.get_obj_image()
+        Outputs = NamedTuple("Outputs", [("objs_list", list)])
+        return Outputs(objs_list)
+
+    def get_obj_image(self):
+        objs_list = []
+        img = cv2.imread(self.img_path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        # Convert image from openCV format to PIL.Image
+        for bbox in self.bboxes:
+            tmp_obj = img[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2])]
+            pil_img = Image.fromarray(cv2.cvtColor(tmp_obj, cv2.COLOR_BGR2RGB))
+            objs_list.append(pil_img)
+        return objs_list

towhee/tests/mock_operators/pytorch_yolo_detector_operator/requirements.txt

@@ -0,0 +1,22 @@
+# pip install -r requirements.txt
+
+# Base ----------------------------------------
+matplotlib>=3.2.2
+numpy>=1.18.5
+opencv-python>=4.1.2
+Pillow>=7.1.2
+PyYAML>=5.3.1
+scipy>=1.4.1
+torch>=1.7.0
+torchvision>=0.8.1
+tqdm>=4.41.0
+
+# Logging -------------------------------------
+tensorboard>=2.4.1
+
+# Plotting ------------------------------------
+pandas
+seaborn>=0.11.0
+
+# Extras --------------------------------------
+thop  # FLOPs computation

towhee/tests/test_util/test_yolo.yaml

@@ -0,0 +1,58 @@
+name: 'test'
+operators:
+    -
+        name: '_start_op'
+        function: 'internal'
+        init_args:
+        inputs:
+            -
+                df: 'img_path'
+                name: 'img_path'
+                col: 0
+        outputs:
+            -
+                df: 'img_path'
+        iter_info:
+            type: map
+    -
+        name: 'object_detection'
+        function: 'mock_operators/pytorch_yolo_detector_operator'
+        init_args:
+            model_name: 'yolov5s'
+        inputs:
+            -
+                df: 'img_path'
+                name: 'img_path'
+                col: 0
+        outputs:
+            -
+                df: 'objs_list'
+        iter_info:
+            type: map
+    -
+        name: '_end_op'
+        function: 'internal'
+        init_args:
+        inputs:
+            -
+                df: 'objs_list'
+                name: 'objs_list'
+                col: 0
+        outputs:
+            -
+                df: 'objs_list'
+        iter_info:
+            type: map
+dataframes:
+    -
+        name: 'img_path'
+        columns:
+            -
+                name: 'img_path'
+                vtype: 'str'
+    -
+        name: 'objs_list'
+        columns:
+            -
+                name: 'objs_list'
+                vtype: 'list'