From e2e1bc4cd457ce5a847d363c75ac3cab10061398 Mon Sep 17 00:00:00 2001 From: shiyu22 Date: Wed, 28 Dec 2022 19:41:52 +0800 Subject: [PATCH] Update the custom path Signed-off-by: shiyu22 --- yolov5.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/yolov5.py b/yolov5.py index 4c8d1cd..a1af818 100644 --- a/yolov5.py +++ b/yolov5.py @@ -14,9 +14,9 @@ logging.getLogger("yolov5").setLevel(logging.WARNING) class Yolov5(NNOperator): def __init__(self): super().__init__() - model_path = str(Path(__file__).parent / 'models/yolov5s') + model_path = str(Path(__file__).parent / 'models/yolov5s.pt') self._model = torch.hub.load('ultralytics/yolov5', 'custom', model_path) - # self._model = torch.hub.load("ultralytics/yolov5", model_name, pretrained=True, verbose=False) + # self._model = torch.hub.load("ultralytics/yolov5", 'yolov5s', pretrained=True, verbose=False) def __call__(self, img: numpy.ndarray): # Get object detection results with YOLOv5 model