logo
Browse Source

Add models

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
shiyu22 2 years ago
parent
commit
1732826220
  1. BIN
      models/yolov5s.pt
  2. 7
      yolov5.py

BIN
models/yolov5s.pt (Stored with Git LFS)

Binary file not shown.

7
yolov5.py

@ -1,6 +1,7 @@
import torch
import numpy
import logging
from pathlib import Path
from towhee import register
from towhee.operator import NNOperator
@ -11,9 +12,11 @@ logging.getLogger("yolov5").setLevel(logging.WARNING)
@register(output_schema=['boxes', 'classes', 'scores'])
class Yolov5(NNOperator):
def __init__(self, model_name: str ='yolov5s'):
def __init__(self):
super().__init__()
self._model = torch.hub.load("ultralytics/yolov5", model_name, pretrained=True, verbose=False)
model_path = str(Path(__file__).parent / 'models/yolov5s')
self._model = torch.hub.load('ultralytics/yolov5', 'custom', model_path)
# self._model = torch.hub.load("ultralytics/yolov5", model_name, pretrained=True, verbose=False)
def __call__(self, img: numpy.ndarray):
# Get object detection results with YOLOv5 model

Loading…
Cancel
Save