Add models

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>

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