support loading models from hf hub (#172)

* add support to load models from huggingface hub

* add suppor to load models from hf hub

* some fixes

* improve tests

* fix ci

* update readme

* update package version

.github/workflows/ci.yml

@@ -63,6 +63,10 @@ jobs:
       - name: Test scripts
         run: |
           pip install -e .
+          # train (dl base model from hf hub)
+          python yolov5/train.py --img 128 --batch 16 --weights fcakyon/yolov5n-v7.0 --epochs 1 --device cpu
+          yolov5 train --img 128 --batch 16 --weights fcakyon/yolov5n-v7.0 --epochs 1 --device cpu --freeze 10
+          yolov5 train --img 128 --batch 16 --weights fcakyon/yolov5n-v7.0 --epochs 1 --device cpu --evolve 2
           # train
           python yolov5/train.py --img 128 --batch 16 --weights yolov5/weights/yolov5n.pt --epochs 1 --device cpu
           yolov5 train --img 128 --batch 16 --weights yolov5/weights/yolov5n.pt --epochs 1 --device cpu --freeze 10
@@ -84,10 +88,12 @@ jobs:
           # benckmarks
           yolov5 benchmarks --weights yolov5n.pt --img 128 --pt-only --device cpu
           # classify
-          yolov5 classify train --img 128 --data mnist2560 --weights yolov5n-cls.pt --epochs 1 --device cpu
+          yolov5 classify train --img 128 --data mnist2560 --model yolov5n-cls.pt --epochs 1 --device cpu
+          yolov5 classify train --img 128 --data mnist2560 --model fcakyon/yolov5n-cls-v7.0 --epochs 1 --device cpu
           yolov5 classify val --img 128 --data datasets/mnist2560 --weights yolov5n-cls.pt --device cpu
           yolov5 classify predict --img 128 --weights yolov5n-cls.pt --device cpu  
           # segment
           yolov5 segment train --img 128 --weights yolov5n-seg.pt --epochs 1 --device cpu
+          yolov5 segment train --img 128 --weights fcakyon/yolov5n-seg-v7.0 --epochs 1 --device cpu
           # yolov5 segment val --img 128 --weights yolov5n-seg.pt --device cpu
           yolov5 segment predict --img 128 --weights yolov5n-seg.pt --device cpu

README.md

@@ -25,7 +25,7 @@ You can finally install <a href="https://github.com/ultralytics/yolov5">YOLOv5 o
 </div>
 
 <br>
-This yolov5 package contains everything from ultralytics/yolov5 <a href="https://github.com/ultralytics/yolov5/commit/357cde9ee7da13ba3095995488c5a23631467f1a">at this commit</a> plus:
+This yolov5 package contains everything from ultralytics/yolov5 <a href="https://github.com/ultralytics/yolov5/tree/357cde9ee7da13ba3095995488c5a23631467f1a">at this commit</a> plus:
 <br>
 1. Easy installation via pip: `pip install yolov5`
 <br>

requirements.txt

@@ -54,3 +54,5 @@ fire
 boto3>=1.19.1
 # coco to yolov5 conversion
 sahi>=0.10.5
+# huggingface
+huggingface-hub>=0.11.1

tests/test_utils.py

@@ -1,4 +1,7 @@
 class TestConstants:
+    YOLOV5N_HUB_ID = "fcakyon/yolov5n-v7.0"
+    YOLOV5S_HUB_ID = "fcakyon/yolov5s-v7.0"
+    YOLOV5N_MODEL_PATH = "yolov5/weights/yolov5n.pt"
     YOLOV5S_MODEL_PATH = "yolov5/weights/yolov5s.pt"
     YOLOV5L_MODEL_PATH = "yolov5/weights/yolov5l.pt"
     ZIDANE_IMAGE_PATH = "tests/data/zidane.jpg"

tests/test_yolov5.py

@@ -4,7 +4,7 @@
 
 DEVICE = "cpu"
 
-class TestYolov5(unittest.TestCase):
+class TestYolov5FromUltralytics(unittest.TestCase):
     def test_load_model(self):
         from yolov5 import YOLOv5
 
@@ -182,6 +182,169 @@ def test_hublike_predict(self):
         self.assertEqual(len(results.pred[1]), 5)
 
 
+class TestYolov5FromHuggingface(unittest.TestCase):
+    def test_load_model(self):
+        from yolov5 import YOLOv5
+
+        # init model
+        model_path = TestConstants.YOLOV5S_HUB_ID
+        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=False)
+        yolov5.load_model()
+
+        # check if loaded
+        self.assertNotEqual(yolov5.model, None)
+
+    def test_load_model_on_init(self):
+        from yolov5 import YOLOv5
+
+        # init model
+        model_path = TestConstants.YOLOV5S_HUB_ID
+        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
+
+        # check if loaded
+        self.assertNotEqual(yolov5.model, None)
+
+    def test_predict(self):
+        from PIL import Image
+        from yolov5 import YOLOv5
+
+        # init yolov5s model
+        model_path = TestConstants.YOLOV5S_HUB_ID
+        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
+
+        # prepare image
+        image_path = TestConstants.ZIDANE_IMAGE_PATH
+        image = Image.open(image_path)
+
+        # perform inference
+        results = yolov5.predict(image, size=640, augment=False)
+
+        # compare
+        self.assertEqual(results.n, 1)
+        self.assertEqual(len(results.names), 80)
+        self.assertEqual(len(results.pred[0]), 4)
+
+        # prepare image
+        image = image_path
+
+        # perform inference
+        results = yolov5.predict(image, size=640, augment=True)
+
+        # compare
+        self.assertEqual(results.n, 1)
+        self.assertEqual(len(results.names), 80)
+        self.assertEqual(len(results.pred[0]), 5)
+
+        # init yolov5l model
+        model_path = TestConstants.YOLOV5L_MODEL_PATH
+        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
+
+        # prepare image
+        image_path = TestConstants.BUS_IMAGE_PATH
+        image = Image.open(image_path)
+        # perform inference
+        results = yolov5.predict(image, size=1280, augment=False)
+
+        # compare
+        self.assertEqual(results.n, 1)
+        self.assertEqual(len(results.names), 80)
+        self.assertEqual(len(results.pred[0]), 6)
+
+        # prepare image
+        image = image_path
+
+        # perform inference
+        results = yolov5.predict(image, size=1280, augment=False)
+
+        # compare
+        self.assertEqual(results.n, 1)
+        self.assertEqual(len(results.names), 80)
+        self.assertEqual(len(results.pred[0]), 6)
+
+        # init yolov5s model
+        model_path = TestConstants.YOLOV5S_HUB_ID
+        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
+
+        # prepare images
+        image_path1 = TestConstants.ZIDANE_IMAGE_PATH
+        image_path2 = TestConstants.BUS_IMAGE_PATH
+        image1 = Image.open(image_path1)
+        image2 = Image.open(image_path2)
+
+        # perform inference with multiple images and test augmentation
+        results = yolov5.predict([image1, image2], size=1280, augment=True)
+
+        # compare
+        self.assertEqual(results.n, 2)
+        self.assertEqual(len(results.names), 80)
+        self.assertEqual(len(results.pred[0]), 4)
+        self.assertEqual(len(results.pred[1]), 5)
+
+        # prepare image
+        image_path1 = TestConstants.ZIDANE_IMAGE_PATH
+        image_path2 = TestConstants.BUS_IMAGE_PATH
+        image1 = image_path1
+        image2 = image_path2
+
+        # perform inference
+        results = yolov5.predict([image1, image2], size=1280, augment=True)
+
+        # compare
+        self.assertEqual(results.n, 2)
+        self.assertEqual(len(results.names), 80)
+        self.assertEqual(len(results.pred[0]), 4)
+        self.assertEqual(len(results.pred[1]), 5)
+
+    def test_hublike_load_model(self):
+        import yolov5
+
+        # init model
+        model_path = TestConstants.YOLOV5S_HUB_ID
+        model = yolov5.load(model_path, device=DEVICE)
+
+        # check if loaded
+        self.assertNotEqual(model, None)
+
+    def test_hublike_predict(self):
+        import yolov5
+        from PIL import Image
+
+        # init yolov5s model
+        model_path = TestConstants.YOLOV5S_HUB_ID
+        model = yolov5.load(model_path, device=DEVICE)
+
+        # prepare image
+        image_path = TestConstants.ZIDANE_IMAGE_PATH
+        image = Image.open(image_path)
+
+        # perform inference
+        results = model(image, size=640, augment=False)
+
+        # compare
+        self.assertEqual(results.n, 1)
+        self.assertEqual(len(results.names), 80)
+        self.assertEqual(len(results.pred[0]), 4)
+
+        # init yolov5s model
+        model_path = TestConstants.YOLOV5S_HUB_ID
+        model = yolov5.load(model_path, device=DEVICE)
+
+        # prepare images
+        image_path1 = TestConstants.ZIDANE_IMAGE_PATH
+        image_path2 = TestConstants.BUS_IMAGE_PATH
+        image1 = Image.open(image_path1)
+        image2 = Image.open(image_path2)
+
+        # perform inference with multiple images and test augmentation
+        results = model([image1, image2], size=1280, augment=True)
+
+        # compare
+        self.assertEqual(results.n, 2)
+        self.assertEqual(len(results.names), 80)
+        self.assertEqual(len(results.pred[0]), 4)
+        self.assertEqual(len(results.pred[1]), 5)
+
+
 
 if __name__ == "__main__":
     unittest.main()

yolov5/__init__.py

@@ -1,4 +1,4 @@
 from yolov5.helpers import YOLOv5
 from yolov5.helpers import load_model as load
 
-__version__ = "7.0.0"
+__version__ = "7.0.1"

yolov5/classify/train.py

@@ -30,6 +30,8 @@
 from torch.cuda import amp
 from tqdm import tqdm
 
+from yolov5.utils.downloads import attempt_donwload_from_hub
+
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[1]  # YOLOv5 root directory
 if str(ROOT) not in sys.path:
@@ -106,6 +108,10 @@ def train(opt, device):
                                                       workers=nw)
 
     # Model
+    # try to download from hf hub
+    result = attempt_donwload_from_hub(opt.model, hf_token=None)
+    if result is not None:
+        opt.model = result
     with torch_distributed_zero_first(LOCAL_RANK), WorkingDirectory(ROOT):
         if Path(opt.model).is_file() or opt.model.endswith('.pt'):
             model = attempt_load(opt.model, device='cpu', fuse=False)

yolov5/helpers.py

@@ -7,7 +7,7 @@
 from yolov5.utils.torch_utils import select_device
 
 
-def load_model(model_path, device=None, autoshape=True, verbose=False):
+def load_model(model_path, device=None, autoshape=True, verbose=False, hf_token: str = None):
     """
     Creates a specified YOLOv5 model
 
@@ -17,6 +17,7 @@ def load_model(model_path, device=None, autoshape=True, verbose=False):
         pretrained (bool): load pretrained weights into the model
         autoshape (bool): make model ready for inference
         verbose (bool): if False, yolov5 logs will be silent
+        hf_token (str): huggingface read token for private models
 
     Returns:
         pytorch model
@@ -31,7 +32,7 @@ def load_model(model_path, device=None, autoshape=True, verbose=False):
     device = select_device(device)
 
     try:
-        model = DetectMultiBackend(model_path, device=device, fuse=autoshape)  # detection model
+        model = DetectMultiBackend(model_path, device=device, fuse=autoshape, hf_token=hf_token)  # detection model
         if autoshape:
             if model.pt and isinstance(model.model, ClassificationModel):
                 LOGGER.warning('WARNING ⚠️ YOLOv5 ClassificationModel is not yet AutoShape compatible. '
@@ -81,6 +82,7 @@ def predict(self, image_list, size=640, augment=False):
         results = self.model(ims=image_list, size=size, augment=augment)
         return results
 
+
 if __name__ == "__main__":
     model_path = "yolov5/weights/yolov5s.pt"
     device = "cuda:0"

yolov5/models/common.py

@@ -27,6 +27,7 @@
 
 from yolov5.utils import TryExcept
 from yolov5.utils.dataloaders import exif_transpose, letterbox
+from yolov5.utils.downloads import attempt_donwload_from_hub
 from yolov5.utils.general import (LOGGER, ROOT, Profile, check_requirements, check_suffix, check_version, colorstr,
                            increment_path, is_notebook, make_divisible, non_max_suppression, scale_boxes, xywh2xyxy,
                            xyxy2xywh, yaml_load)
@@ -315,7 +316,7 @@ def forward(self, x):
 
 class DetectMultiBackend(nn.Module):
     # YOLOv5 MultiBackend class for python inference on various backends
-    def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True):
+    def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True, hf_token=None):
         # Usage:
         #   PyTorch:              weights = *.pt
         #   TorchScript:                    *.torchscript
@@ -332,7 +333,14 @@ def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False,
         from yolov5.models.experimental import attempt_download, attempt_load  # scoped to avoid circular import
 
         super().__init__()
+
         w = str(weights[0] if isinstance(weights, list) else weights)
+
+        # try to dowload from hf hub
+        result = attempt_donwload_from_hub(w, hf_token=hf_token)
+        if result is not None:
+            w = result
+
         pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, triton = self._model_type(w)
         fp16 &= pt or jit or onnx or engine  # FP16
         nhwc = coreml or saved_model or pb or tflite or edgetpu  # BHWC formats (vs torch BCWH)

yolov5/segment/train.py

@@ -45,7 +45,7 @@
 from yolov5.utils.autoanchor import check_anchors
 from yolov5.utils.autobatch import check_train_batch_size
 from yolov5.utils.callbacks import Callbacks
-from yolov5.utils.downloads import attempt_download, is_url
+from yolov5.utils.downloads import attempt_donwload_from_hub, attempt_download, is_url
 from yolov5.utils.general import (LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_info,
                            check_git_status, check_img_size, check_requirements, check_suffix, check_yaml, colorstr,
                            get_latest_run, increment_path, init_seeds, intersect_dicts, labels_to_class_weights,
@@ -106,6 +106,10 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt')  # COCO dataset
 
     # Model
+    # try to download from hf hub
+    result = attempt_donwload_from_hub(weights, hf_token=None)
+    if result is not None:
+        weights = result
     check_suffix(weights, '.pt')  # check weights
     pretrained = weights.endswith('.pt')
     if pretrained:

yolov5/train.py

@@ -47,7 +47,7 @@
 from yolov5.utils.autobatch import check_train_batch_size
 from yolov5.utils.callbacks import Callbacks
 from yolov5.utils.dataloaders import create_dataloader
-from yolov5.utils.downloads import attempt_download, is_url
+from yolov5.utils.downloads import attempt_donwload_from_hub, attempt_download, is_url
 from yolov5.utils.general import (LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_info,
                            check_git_status, check_img_size, check_requirements, check_suffix, check_yaml, colorstr,
                            get_latest_run, increment_path, init_seeds, intersect_dicts, labels_to_class_weights,
@@ -180,7 +180,12 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     # upload dataset to s3
     if opt.upload_dataset and opt.s3_upload_dir:
         upload_to_s3(opt, data, save_dir)
+
     # Model
+    # try to download from hf hub
+    result = attempt_donwload_from_hub(weights, hf_token=None)
+    if result is not None:
+        weights = result
     check_suffix(weights, '.pt')  # check weights
     pretrained = weights.endswith('.pt')
     if pretrained: