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