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()
-```
-
-
+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()
+```
@@ -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