From a14282de722aec50400837988b17b13f03a3b4fb Mon Sep 17 00:00:00 2001
From: shiyu22 <shiyu.chen@zilliz.com>
Date: Thu, 29 Dec 2022 13:25:17 +0800
Subject: [PATCH] Update readme and models

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
---
 README.md         | 21 +++++++++++----------
 models/yolov5s.pt |  3 +++
 yolov5.py         |  6 ++++--
 3 files changed, 18 insertions(+), 12 deletions(-)
 create mode 100644 models/yolov5s.pt

diff --git a/README.md b/README.md
index 0b6f175..ae35cb2 100644
--- a/README.md
+++ b/README.md
@@ -27,17 +27,18 @@ This operator uses [PyTorch.yolov5](https://pytorch.org/hub/ultralytics_yolov5/)
 Writing the pipeline in the simplified way
 
 ```Python
-import towhee
-
-towhee.glob('./test.png') \
-	.image_decode() \
-	.object_detection.yolov5() \
-	.show()
-```
-
-<img src="./results1.png" alt="results1" height="40px"/>
+from towhee.dc2 import pipe, ops
 
+p = (
+    pipe.input('url')
+        .map('url', 'img', ops.image_decode.cv2_rgb())
+        .flat_map('img', ('boxes', 'class', 'score'), ops.object_detection.yolo())
+        .output('class', 'score')
+    )
 
+data = 'https://towhee.io/object-detection/yolov5/raw/branch/main/test.png'
+res = p(data).get()
+```
 
 <br />
 
@@ -67,7 +68,7 @@ The operator takes an image as input. It first detects the objects appeared in t
 
 
 
-**Return**: List[List[(int, int, int, int)], ...], List[str], List[float]
+**Return**: List[List[(int, int, int, int)], ...], List[str], List[float]]
 
 The return value is a tuple of (boxes, classes, scores). The *boxes* is a list of bounding boxes. Each bounding box is represented by the top-left and the bottom right points, i.e. (x1, y1, x2, y2). The *classes* is a list of prediction labels. The *scores* is a list of the confidence scores.
 
diff --git a/models/yolov5s.pt b/models/yolov5s.pt
new file mode 100644
index 0000000..cd69a5e
--- /dev/null
+++ b/models/yolov5s.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b3b748c1e592ddd8868022e8732fde20025197328490623cc16c6f24d0782ee
+size 14808437
diff --git a/yolov5.py b/yolov5.py
index f35679b..fed9d58 100644
--- a/yolov5.py
+++ b/yolov5.py
@@ -6,9 +6,11 @@ from towhee.operator import NNOperator
 
 @register(output_schema=['boxes', 'classes', 'scores'])
 class Yolov5(NNOperator):
-    def __init__(self, model_name: str ='yolov5s'):
+    def __init__(self):
         super().__init__()
-        self._model = torch.hub.load("ultralytics/yolov5", model_name, pretrained=True)
+        model_path = str(Path(__file__).parent / 'models/yolov5s.pt')
+        self._model = torch.hub.load('ultralytics/yolov5', 'custom', model_path)
+        # self._model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
 
     def __call__(self, img: numpy.ndarray):
         # Get object detection results with YOLOv5 model