yolov5
copied
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Readme
Files and versions
19 lines
627 B
19 lines
627 B
3 years ago
|
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
|