From 7a0cce17b64133c03ee276f250168622c062f8be Mon Sep 17 00:00:00 2001 From: wa22 Date: Mon, 3 Jun 2024 10:21:43 +0800 Subject: [PATCH] add names when pushing to huggingface hub --- ultralytics/models/yolov10/model.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/ultralytics/models/yolov10/model.py b/ultralytics/models/yolov10/model.py index a8881395..33fbd845 100644 --- a/ultralytics/models/yolov10/model.py +++ b/ultralytics/models/yolov10/model.py @@ -6,8 +6,19 @@ from .train import YOLOv10DetectionTrainer class YOLOv10(Model): - def __init__(self, model="yolov10n.pt", task=None, verbose=False): + def __init__(self, model="yolov10n.pt", task=None, verbose=False, + names=None): super().__init__(model=model, task=task, verbose=verbose) + if names is not None: + setattr(self.model, 'names', names) + + def push_to_hub(self, repo_name, **kwargs): + config = kwargs.get('config', {}) + config['names'] = self.names + config['model'] = self.model.yaml['yaml_file'] + config['task'] = self.task + kwargs['config'] = config + super().push_to_hub(repo_name, **kwargs) @property def task_map(self):