Skip to content

Reference for ultralytics/models/yolo/model.py

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/model.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!


ultralytics.models.yolo.model.YOLO 

YOLO(model='yolov8n.pt', task=None, verbose=False)

Bases: Model

YOLO (You Only Look Once) object detection model.

Source code in ultralytics/models/yolo/model.py 
14
15
16
17
18
19
20
21
22
23
def __init__(self, model="yolov8n.pt", task=None, verbose=False):
    """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, verbose=verbose)
        self.__class__ = type(new_instance)
        self.__dict__ = new_instance.__dict__
    else:
        # Continue with default YOLO initialization
        super().__init__(model=model, task=task, verbose=verbose)

task_map property 

task_map

Map head to model, trainer, validator, and predictor classes.




ultralytics.models.yolo.model.YOLOWorld 

YOLOWorld(model='yolov8s-world.pt', verbose=False)

Bases: Model

YOLO-World object detection model.

Parameters:

Name Type Description Default
model str | Path

Path to the pre-trained model. Defaults to 'yolov8s-world.pt'.

 'yolov8s-world.pt'
Source code in ultralytics/models/yolo/model.py 
65
66
67
68
69
70
71
72
73
74
75
76
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", verbose=verbose)

    # Assign default COCO class names when there are no custom names
    if not hasattr(self.model, "names"):
        self.model.names = yaml_load(ROOT / "cfg/datasets/coco8.yaml").get("names")

task_map property 

task_map

Map head to model, validator, and predictor classes.

set_classes 

set_classes(classes)

Set classes.

Parameters:

Name Type Description Default
classes List(str

A list of categories i.e. ["person"].

 required
Source code in ultralytics/models/yolo/model.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def set_classes(self, classes):
    """
    Set classes.

    Args:
        classes (List(str)): A list of categories i.e. ["person"].
    """
    self.model.set_classes(classes)
    # Remove background if it's given
    background = " "
    if background in classes:
        classes.remove(background)
    self.model.names = classes

    # Reset method class names
    # self.predictor = None  # reset predictor otherwise old names remain
    if self.predictor:
        self.predictor.model.names = classes





Created 2023-11-12, Updated 2024-07-21
Authors: glenn-jocher (7), sergiuwaxmann (1), Burhan-Q (1)