Browse Source
Set log level
Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
1 changed files with
4 additions and
1 deletions
-
yolov5.py
|
|
|
@ -1,14 +1,17 @@ |
|
|
|
import torch |
|
|
|
import numpy |
|
|
|
import logging |
|
|
|
|
|
|
|
from towhee import register |
|
|
|
from towhee.operator import NNOperator |
|
|
|
|
|
|
|
logging.getLogger("yolov5").setLevel(logging.WARNING) |
|
|
|
|
|
|
|
@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) |
|
|
|
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 |
|
|
|
|
