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)
self.__class__ = type(new_instance)
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:
# Continue with default YOLO initialization
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 .val import YOLOv10DetectionValidator
from .predict import YOLOv10DetectionPredictor
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
def task_map(self):
"""Map head to model, trainer, validator, and predictor classes."""