diff --git a/hubconf.py b/hubconf.py index 7f897d1..7359a58 100644 --- a/hubconf.py +++ b/hubconf.py @@ -18,7 +18,7 @@ dependencies = ['torch', 'yaml'] check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('tensorboard', 'pycocotools', 'thop')) -def create(name, pretrained, channels=3, classes=80, autoshape=True, verbose=True): +def create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): """Creates a specified YOLOv5 model Arguments: @@ -106,9 +106,8 @@ def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=Tr if __name__ == '__main__': - model = create(name='weights/yolov5s.pt', pretrained=True, channels=3, classes=80, autoshape=True, - verbose=True) # pretrained - # model = custom(path_or_model='path/to/model.pt') # custom + model = create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained + # model = custom(path='path/to/model.pt') # custom # Verify inference import cv2