Skip to content

YOLOv5 Model

sahi.models.yolov5

YOLOv5 detection model wrapper for SAHI.

Provides integration with Ultralytics YOLOv5 for object detection.

Classes

Yolov5DetectionModel 

Yolov5DetectionModel(*args: object, **kwargs: object)

Bases: DetectionModel

YOLOv5 object detection model.

Wraps Ultralytics YOLOv5 for fast object detection.

Initialize YOLOv5 detection model.

Source code in sahi/models/yolov5.py 
25
26
27
28
29
def __init__(self, *args: object, **kwargs: object) -> None:
    """Initialize YOLOv5 detection model."""
    existing_packages = getattr(self, "required_packages", None) or []
    self.required_packages = [*list(existing_packages), "yolov5", "torch"]
    super().__init__(*args, **kwargs)  # type: ignore[misc, arg-type]
Attributes
num_categories property 
num_categories: int

Returns number of categories.

has_mask property 
has_mask: bool

Returns if model output contains segmentation mask.

category_names property 
category_names: list

Return category names from model.

Methods:
load_model 
load_model() -> None

Detection model is initialized and set to self.model.

Source code in sahi/models/yolov5.py 
31
32
33
34
35
36
37
38
39
def load_model(self) -> None:
    """Detection model is initialized and set to self.model."""
    import yolov5

    try:
        model = yolov5.load(self.model_path, device=self.device)
        self.set_model(model)
    except Exception as e:
        raise TypeError("model_path is not a valid yolov5 model path: ", e)
set_model 
set_model(model: Any, **kwargs: Any) -> None

Sets the underlying YOLOv5 model.

Parameters:

Name Type Description Default
model Any

Any A YOLOv5 model

 required
**kwargs Any

Any Additional keyword arguments for model setup.

 {}
Source code in sahi/models/yolov5.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def set_model(self, model: Any, **kwargs: Any) -> None:
    """Sets the underlying YOLOv5 model.

    Args:
        model: Any
            A YOLOv5 model
        **kwargs: Any
            Additional keyword arguments for model setup.
    """
    if model.__class__.__module__ not in ["yolov5.models.common", "models.common"]:
        raise Exception(f"Not a yolov5 model: {type(model)}")

    model.conf = self.confidence_threshold  # type: ignore[attr-defined]
    self.model = model

    # set category_mapping
    if not self.category_mapping:
        category_mapping = {str(ind): category_name for ind, category_name in enumerate(self.category_names)}
        self.category_mapping = category_mapping
perform_inference 
perform_inference(image: ndarray) -> None

Prediction is performed using self.model and the prediction result is set to self._original_predictions.

Parameters:

Name Type Description Default
image ndarray

np.ndarray A numpy array that contains the image to be predicted. 3 channel image should be in RGB order.

 required
Source code in sahi/models/yolov5.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def perform_inference(self, image: np.ndarray) -> None:
    """Prediction is performed using self.model and the prediction result is set to self._original_predictions.

    Args:
        image: np.ndarray
            A numpy array that contains the image to be predicted. 3 channel image should be in RGB order.
    """
    # Confirm model is loaded
    if self.model is None:
        raise ValueError("Model is not loaded, load it by calling .load_model()")
    if self.image_size is not None:
        prediction_result = self.model(image, size=self.image_size)
    else:
        prediction_result = self.model(image)

    self._original_predictions = prediction_result

Functions: