VietCat Amp commited on
Commit
e1c596d
·
1 Parent(s): 98a8397

Fix repo_id parsing for HF Hub download

Browse files

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

Files changed (1) hide show
  1. model.py +4 -3
model.py CHANGED
@@ -12,9 +12,10 @@ class TrafficSignDetector:
12
 
13
  # Download model from HF Hub if not local
14
  model_path = config['model']['path']
15
- if model_path.startswith('VietCat/'):
16
- # Extract repo and file path
17
- repo_id, file_path = model_path.split('/', 1)
 
18
  local_model_path = hf_hub_download(repo_id=repo_id, filename=file_path)
19
  self.model = YOLO(local_model_path)
20
  else:
 
12
 
13
  # Download model from HF Hub if not local
14
  model_path = config['model']['path']
15
+ if model_path.startswith('VietCat/GTSRB-Model/'):
16
+ # Extract repo and file path for GTSRB-Model
17
+ repo_id = 'VietCat/GTSRB-Model'
18
+ file_path = model_path.replace('VietCat/GTSRB-Model/', '')
19
  local_model_path = hf_hub_download(repo_id=repo_id, filename=file_path)
20
  self.model = YOLO(local_model_path)
21
  else: