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):