Skip to content

Ultralytics YOLO

tensoris.lib.integrations.ultralytics

Ultralytics YOLO Integration Helper.

Classes

UltralyticsIntegration

Ultralytics YOLO model helper for object detection, segmentation, and classification workflows.

References

Ultralytics Documentation: https://docs.ultralytics.com/

Source code in src/tensoris/lib/integrations/ultralytics.py
class UltralyticsIntegration:
    """Ultralytics YOLO model helper for object detection, segmentation, and classification workflows.

    References:
        Ultralytics Documentation: https://docs.ultralytics.com/
    """

    def __init__(self, model_name: str = "yolov8n.pt") -> None:
        """Initialize Ultralytics YOLO model.

        Args:
            model_name: Pretrained YOLO weight file or config path (e.g. 'yolov8n.pt', 'yolov8x-seg.pt').
        """
        self.model_name = model_name
        self.model = None

    @property
    def is_available(self) -> bool:
        """Check if ultralytics package is installed."""
        return _ULTRALYTICS_AVAILABLE

    def load_model(self) -> Any:
        """Instantiate Ultralytics YOLO model class."""
        if not self.is_available:
            raise RuntimeError(
                "ultralytics package is not installed. Install via `pip install ultralytics`."
            )
        self.model = ultralytics.YOLO(self.model_name)  # type: ignore
        return self.model

    def train(
        self, data_yaml: str | Path, epochs: int = 50, imgsz: int = 640, batch: int = 16
    ) -> Any:
        """Train YOLO model on dataset.

        Args:
            data_yaml: Path to dataset configuration YAML file.
            epochs: Training epoch count.
            imgsz: Target image resolution size.
            batch: Training batch size.

        Returns:
            Training results object.
        """
        if self.model is None:
            self.load_model()
        assert self.model is not None
        return self.model.train(
            data=str(data_yaml),
            epochs=epochs,
            imgsz=imgsz,
            batch=batch,
        )  # type: ignore

    def export(self, format: str = "onnx") -> str:
        """Export trained YOLO model to deployment format (ONNX, TorchScript, Engine).

        Args:
            format: Output target format ('onnx', 'torchscript', 'engine', 'openvino').

        Returns:
            String path of exported model weight file.
        """
        if self.model is None:
            self.load_model()
        assert self.model is not None
        return self.model.export(format=format)  # type: ignore
Attributes
is_available property
is_available

Check if ultralytics package is installed.

Methods:
__init__
__init__(model_name='yolov8n.pt')

Initialize Ultralytics YOLO model.

Parameters:

Name Type Description Default
model_name str

Pretrained YOLO weight file or config path (e.g. 'yolov8n.pt', 'yolov8x-seg.pt').

'yolov8n.pt'
Source code in src/tensoris/lib/integrations/ultralytics.py
def __init__(self, model_name: str = "yolov8n.pt") -> None:
    """Initialize Ultralytics YOLO model.

    Args:
        model_name: Pretrained YOLO weight file or config path (e.g. 'yolov8n.pt', 'yolov8x-seg.pt').
    """
    self.model_name = model_name
    self.model = None
export
export(format='onnx')

Export trained YOLO model to deployment format (ONNX, TorchScript, Engine).

Parameters:

Name Type Description Default
format str

Output target format ('onnx', 'torchscript', 'engine', 'openvino').

'onnx'

Returns:

Type Description
str

String path of exported model weight file.

Source code in src/tensoris/lib/integrations/ultralytics.py
def export(self, format: str = "onnx") -> str:
    """Export trained YOLO model to deployment format (ONNX, TorchScript, Engine).

    Args:
        format: Output target format ('onnx', 'torchscript', 'engine', 'openvino').

    Returns:
        String path of exported model weight file.
    """
    if self.model is None:
        self.load_model()
    assert self.model is not None
    return self.model.export(format=format)  # type: ignore
load_model
load_model()

Instantiate Ultralytics YOLO model class.

Source code in src/tensoris/lib/integrations/ultralytics.py
def load_model(self) -> Any:
    """Instantiate Ultralytics YOLO model class."""
    if not self.is_available:
        raise RuntimeError(
            "ultralytics package is not installed. Install via `pip install ultralytics`."
        )
    self.model = ultralytics.YOLO(self.model_name)  # type: ignore
    return self.model
train
train(data_yaml, epochs=50, imgsz=640, batch=16)

Train YOLO model on dataset.

Parameters:

Name Type Description Default
data_yaml str | Path

Path to dataset configuration YAML file.

required
epochs int

Training epoch count.

50
imgsz int

Target image resolution size.

640
batch int

Training batch size.

16

Returns:

Type Description
Any

Training results object.

Source code in src/tensoris/lib/integrations/ultralytics.py
def train(
    self, data_yaml: str | Path, epochs: int = 50, imgsz: int = 640, batch: int = 16
) -> Any:
    """Train YOLO model on dataset.

    Args:
        data_yaml: Path to dataset configuration YAML file.
        epochs: Training epoch count.
        imgsz: Target image resolution size.
        batch: Training batch size.

    Returns:
        Training results object.
    """
    if self.model is None:
        self.load_model()
    assert self.model is not None
    return self.model.train(
        data=str(data_yaml),
        epochs=epochs,
        imgsz=imgsz,
        batch=batch,
    )  # type: ignore