From 9465d49b2325cfdb896d3c082a61992217059683 Mon Sep 17 00:00:00 2001
From: shiyu22 <shiyu.chen@zilliz.com>
Date: Mon, 11 Apr 2022 18:36:52 +0800
Subject: [PATCH] Add yolov5 python file

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
---
 __init__.py      |  5 +++++
 requirements.txt |  5 +++++
 yolov5.py        | 18 ++++++++++++++++++
 3 files changed, 28 insertions(+)
 create mode 100644 __init__.py
 create mode 100644 requirements.txt
 create mode 100644 yolov5.py

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..64db8e3
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,5 @@
+from yolov5.yolov5 import Yolov5
+
+
+def yolov5():
+    return Yolov5()
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..e27b04f
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,5 @@
+matplotlib>=3.2.2
+opencv-python>=4.1.2
+torch>=1.7.0
+torchvision>=0.8.1
+seaborn
diff --git a/yolov5.py b/yolov5.py
new file mode 100644
index 0000000..b60e2d8
--- /dev/null
+++ b/yolov5.py
@@ -0,0 +1,18 @@
+import torch
+import numpy
+from towhee.operator import NNOperator
+
+@register(outputschema=['boxes', 'classes', 'scores'])
+class Yolov5(NNOperator):
+    def __init__(self, model_name):
+        super().__init__()
+        self._model = torch.hub.load("ultralytics/yolov5", model_name, pretrained=True)
+
+    def __call__(self, img: numpy.ndarray):
+        # Get object detection results with YOLOv5 model
+        results = self._model(img)
+        
+        boxes = results.xyxy[0]
+        classes = list(results.pandas().xyxy[0].name)
+        scores = list(results.pandas().xyxy[0].confidence)
+        return boxes, classes, scores