yolov5
copied
shiyu22
3 years ago
3 changed files with 28 additions and 0 deletions
@ -0,0 +1,5 @@ |
|||
from yolov5.yolov5 import Yolov5 |
|||
|
|||
|
|||
def yolov5(): |
|||
return Yolov5() |
@ -0,0 +1,5 @@ |
|||
matplotlib>=3.2.2 |
|||
opencv-python>=4.1.2 |
|||
torch>=1.7.0 |
|||
torchvision>=0.8.1 |
|||
seaborn |
@ -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…
Reference in new issue