diff --git a/yolov5.py b/yolov5.py index 477d9d6..a9f10d6 100644 --- a/yolov5.py +++ b/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