From 1f6d1fc7e6bcf10bf61e63ba1d863bd5f1eacddb Mon Sep 17 00:00:00 2001 From: wa22 Date: Tue, 28 May 2024 21:13:14 +0800 Subject: [PATCH] fix YOLO init --- ultralytics/models/yolo/model.py | 5 +++++ ultralytics/models/yolov10/model.py | 8 ++++++-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/ultralytics/models/yolo/model.py b/ultralytics/models/yolo/model.py index 1bd38d3b..f10dc977 100644 --- a/ultralytics/models/yolo/model.py +++ b/ultralytics/models/yolo/model.py @@ -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) diff --git a/ultralytics/models/yolov10/model.py b/ultralytics/models/yolov10/model.py index e1c3e28c..a8881395 100644 --- a/ultralytics/models/yolov10/model.py +++ b/ultralytics/models/yolov10/model.py @@ -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."""