logo
Browse Source

Update readme and models

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
shiyu22 2 years ago
parent
commit
a14282de72
  1. 21
      README.md
  2. BIN
      models/yolov5s.pt
  3. 6
      yolov5.py

21
README.md

@ -27,17 +27,18 @@ This operator uses [PyTorch.yolov5](https://pytorch.org/hub/ultralytics_yolov5/)
Writing the pipeline in the simplified way Writing the pipeline in the simplified way
```Python ```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 /> <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. 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.

BIN
models/yolov5s.pt (Stored with Git LFS)

Binary file not shown.

6
yolov5.py

@ -6,9 +6,11 @@ from towhee.operator import NNOperator
@register(output_schema=['boxes', 'classes', 'scores']) @register(output_schema=['boxes', 'classes', 'scores'])
class Yolov5(NNOperator): class Yolov5(NNOperator):
def __init__(self, model_name: str ='yolov5s'):
def __init__(self):
super().__init__() 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): def __call__(self, img: numpy.ndarray):
# Get object detection results with YOLOv5 model # Get object detection results with YOLOv5 model

Loading…
Cancel
Save