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