fix YOLO init

This commit is contained in:
wa22 2024-05-28 21:13:14 +08:00
parent d6c084c33f
commit 1f6d1fc7e6
2 changed files with 11 additions and 2 deletions

View File

@ -18,6 +18,11 @@ class YOLO(Model):
new_instance = YOLOWorld(path) new_instance = YOLOWorld(path)
self.__class__ = type(new_instance) self.__class__ = type(new_instance)
self.__dict__ = new_instance.__dict__ self.__dict__ = new_instance.__dict__
elif "yolov10" in path.stem:
from ultralytics import YOLOv10
new_instance = YOLOv10(path)
self.__class__ = type(new_instance)
self.__dict__ = new_instance.__dict__
else: else:
# Continue with default YOLO initialization # Continue with default YOLO initialization
super().__init__(model=model, task=task, verbose=verbose) super().__init__(model=model, task=task, verbose=verbose)

View File

@ -1,10 +1,14 @@
from ..yolo import YOLO from ultralytics.engine.model import Model
from ultralytics.nn.tasks import YOLOv10DetectionModel from ultralytics.nn.tasks import YOLOv10DetectionModel
from .val import YOLOv10DetectionValidator from .val import YOLOv10DetectionValidator
from .predict import YOLOv10DetectionPredictor from .predict import YOLOv10DetectionPredictor
from .train import YOLOv10DetectionTrainer from .train import YOLOv10DetectionTrainer
class YOLOv10(YOLO): class YOLOv10(Model):
def __init__(self, model="yolov10n.pt", task=None, verbose=False):
super().__init__(model=model, task=task, verbose=verbose)
@property @property
def task_map(self): def task_map(self):
"""Map head to model, trainer, validator, and predictor classes.""" """Map head to model, trainer, validator, and predictor classes."""