yolo
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
22 lines
772 B
22 lines
772 B
2 years ago
|
import torch
|
||
|
|
import numpy
|
||
|
|
|
||
|
|
from towhee import register
|
||
|
|
from towhee.operator import NNOperator
|
||
|
|
|
||
|
|
@register(output_schema=['boxes', 'classes', 'scores'])
|
||
|
|
class Yolov5(NNOperator):
|
||
|
|
def __init__(self, model_name: str ='yolov5s'):
|
||
|
|
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 = [re[0:4] for re in results.xyxy[0]]
|
||
|
|
boxes = [list(map(int, box)) for box in boxes]
|
||
|
|
classes = list(results.pandas().xyxy[0].name)
|
||
|
|
scores = list(results.pandas().xyxy[0].confidence)
|
||
|
|
return [(b, c, s) for b, c, s in zip(boxes, classes, scores)]
|