Download model from HF Hub at runtime

Amp-Thread-ID: https://ampcode.com/threads/T-22e83726-d77c-475f-b185-6d55f37c5979
Co-authored-by: Amp <amp@ampcode.com>

config.yaml

@@ -1,5 +1,5 @@
 model:
-  path: 'VietCat/GTSRB-Model/models/GTSRB.pt' # Path to the YOLO model on Hugging Face Hub
+  path: 'VietCat/GTSRB-Model/models/GTSRB.pt' # Path to the YOLO model on Hugging Face Hub (will be downloaded automatically)
   confidence_threshold: 0.5 # Minimum confidence for detections
 
 inference:

model.py

@@ -2,13 +2,24 @@ import cv2
 import numpy as np
 from ultralytics import YOLO
 import yaml
+from huggingface_hub import hf_hub_download
+import os
 
 class TrafficSignDetector:
     def __init__(self, config_path):
         with open(config_path, 'r') as f:
             config = yaml.safe_load(f)
-        
-        self.model = YOLO(config['model']['path'])
+
+        # Download model from HF Hub if not local
+        model_path = config['model']['path']
+        if model_path.startswith('VietCat/'):
+            # Extract repo and file path
+            repo_id, file_path = model_path.split('/', 1)
+            local_model_path = hf_hub_download(repo_id=repo_id, filename=file_path)
+            self.model = YOLO(local_model_path)
+        else:
+            self.model = YOLO(model_path)
+
         self.conf_threshold = config['model']['confidence_threshold']
         self.box_color = config['inference']['box_color']
         self.text_color = config['inference']['text_color']