xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 13:34:23 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

add names when pushing to huggingface hub

Browse Source
This commit is contained in:
wa22 2024-06-03 10:21:43 +08:00
parent 0efdd2f0c4
commit 7a0cce17b6
1 changed files with 12 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
ultralytics/models/yolov10/model.py
View File

@ -6,8 +6,19 @@ from .train import YOLOv10DetectionTrainer
class YOLOv10(Model): 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) 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 @property
def task_map(self): def task_map(self):