YOLO-World: Fix verbose (#11686)

This commit is contained in:
Laughing 2024-05-06 21:14:33 +08:00 committed by GitHub
parent 9d48190e6d
commit be6bb2aca7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -15,7 +15,7 @@ class YOLO(Model):
"""Initialize YOLO model, switching to YOLOWorld if model filename contains '-world'."""
path = Path(model)
if "-world" in path.stem and path.suffix in {".pt", ".yaml", ".yml"}: # if YOLOWorld PyTorch model
new_instance = YOLOWorld(path)
new_instance = YOLOWorld(path, verbose=verbose)
self.__class__ = type(new_instance)
self.__dict__ = new_instance.__dict__
else:
@ -62,14 +62,14 @@ class YOLO(Model):
class YOLOWorld(Model):
"""YOLO-World object detection model."""
def __init__(self, model="yolov8s-world.pt") -> None:
def __init__(self, model="yolov8s-world.pt", verbose=False) -> None:
"""
Initializes the YOLOv8-World model with the given pre-trained model file. Supports *.pt and *.yaml formats.
Args:
model (str | Path): Path to the pre-trained model. Defaults to 'yolov8s-world.pt'.
"""
super().__init__(model=model, task="detect")
super().__init__(model=model, task="detect", verbose=verbose)
# Assign default COCO class names when there are no custom names
if not hasattr(self.model, "names"):