VietCat commited on
Commit
77916ce
·
1 Parent(s): 44fc75a

Fix model loading from HuggingFace Hub and improve inference preprocessing

Browse files
Files changed (1) hide show
  1. model.py +8 -11
model.py CHANGED
@@ -23,17 +23,14 @@ class TrafficSignDetector:
23
  model_path = config['model']['path']
24
 
25
  # Handle HuggingFace paths
26
- if '/' in model_path:
27
- if model_path.endswith('.pt'):
28
- # Full path with filename (e.g., VietCat/GTSRB-Model/models/GTSRB.pt)
29
- parts = model_path.rsplit('/', 1)
30
- repo_id = parts[0]
31
- file_path = parts[1]
32
- local_model_path = hf_hub_download(repo_id=repo_id, filename=file_path)
33
- self.model = YOLO(local_model_path)
34
- else:
35
- # Just repo name (e.g., user/model-name)
36
- self.model = YOLO(model_path)
37
  else:
38
  # Local path or direct model path
39
  self.model = YOLO(model_path)
 
23
  model_path = config['model']['path']
24
 
25
  # Handle HuggingFace paths
26
+ if model_path.endswith('.pt'):
27
+ # Full path with filename (e.g., VietCat/GTSRB-Model/models/GTSRB.pt)
28
+ # repo_id can only be namespace/repo_name (2 parts max)
29
+ parts = model_path.split('/')
30
+ repo_id = '/'.join(parts[:2]) # Take first two parts: VietCat/GTSRB-Model
31
+ file_path = '/'.join(parts[2:]) # Take rest: models/GTSRB.pt
32
+ local_model_path = hf_hub_download(repo_id=repo_id, filename=file_path)
33
+ self.model = YOLO(local_model_path)
 
 
 
34
  else:
35
  # Local path or direct model path
36
  self.model = YOLO(model_path)