logo
Browse Source

Add yolov5 python file

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
shiyu22 3 years ago
parent
commit
9465d49b23
  1. 5
      __init__.py
  2. 5
      requirements.txt
  3. 18
      yolov5.py

5
__init__.py

@ -0,0 +1,5 @@
from yolov5.yolov5 import Yolov5
def yolov5():
return Yolov5()

5
requirements.txt

@ -0,0 +1,5 @@
matplotlib>=3.2.2
opencv-python>=4.1.2
torch>=1.7.0
torchvision>=0.8.1
seaborn

18
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
Loading…
Cancel
Save