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