Skip to content

Torchvision Model

sahi.models.torchvision

TorchVision detection model wrapper for SAHI.

Provides integration with PyTorch's TorchVision library for object detection and instance segmentation models.

Classes

TorchVisionDetectionModel 

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

Bases: DetectionModel

TorchVision object detection model.

Supports various TorchVision detection models like Faster R-CNN, Mask R-CNN, etc.

Initialize TorchVision detection model.

Source code in sahi/models/torchvision.py 
28
29
30
31
32
def __init__(self, *args: object, **kwargs: object) -> None:
    """Initialize TorchVision detection model."""
    existing_packages = getattr(self, "required_packages", None) or []
    self.required_packages = [*list(existing_packages), "torch", "torchvision"]
    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 mapping.

Methods:
load_model 
load_model() -> None

Load TorchVision model from config and weights.

Source code in sahi/models/torchvision.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def load_model(self) -> None:
    """Load TorchVision model from config and weights."""
    import torch

    # read config params
    model_name = None
    num_classes = None
    if self.config_path is not None:
        with open(self.config_path) as stream:
            try:
                config = yaml.safe_load(stream)
            except yaml.YAMLError as exc:
                raise RuntimeError(exc)

        model_name = config.get("model_name", None)
        num_classes = config.get("num_classes", None)

    # complete params if not provided in config
    if not model_name:
        model_name = "fasterrcnn_resnet50_fpn"
        logger.warning(f"model_name not provided in config, using default model_type: {model_name}'")
    if num_classes is None:
        logger.warning("num_classes not provided in config, using default num_classes: 91")
        num_classes = 91
    if self.model_path is None:
        logger.warning("model_path not provided in config, using pretrained weights and default num_classes: 91.")
        weights = "DEFAULT"
        num_classes = 91
    else:
        weights = None

    # load model
    # Note: torchvision >= 0.13 is required for the 'weights' parameter
    model = MODEL_NAME_TO_CONSTRUCTOR[model_name](num_classes=num_classes, weights=weights)
    if self.model_path:
        try:
            model.load_state_dict(torch.load(self.model_path))
        except Exception as e:
            logger.error(f"Invalid {self.model_path=}")
            raise TypeError("model_path is not a valid torchvision model path: ", e)

    self.set_model(model)
set_model 
set_model(model: Any, **kwargs: Any) -> None

Sets the underlying TorchVision model.

Parameters:

Name Type Description Default
model Any

Any A TorchVision model

 required
**kwargs Any

Any Additional keyword arguments for model setup.

 {}
Source code in sahi/models/torchvision.py 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def set_model(self, model: Any, **kwargs: Any) -> None:
    """Sets the underlying TorchVision model.

    Args:
        model: Any
            A TorchVision model
        **kwargs: Any
            Additional keyword arguments for model setup.
    """
    model.eval()  # type: ignore[attr-defined]
    self.model = model.to(self.device)  # type: ignore[attr-defined]

    # set category_mapping

    if self.category_mapping is None:
        category_names = {str(i): COCO_CLASSES[i] for i in range(len(COCO_CLASSES))}
        self.category_mapping = category_names
perform_inference 
perform_inference(
    image: ndarray, image_size: int | None = None
) -> 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
image_size int | None

int Inference input size.

 None
Source code in sahi/models/torchvision.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def perform_inference(self, image: np.ndarray, image_size: int | None = None) -> 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.
        image_size: int
            Inference input size.
    """
    from sahi.utils.torch_utils import to_float_tensor

    # arrange model input size
    assert self.model is not None
    if self.image_size is not None:
        # get min and max of image height and width
        min_shape, max_shape = min(image.shape[:2]), max(image.shape[:2])
        # torchvision resize transform scales the shorter dimension to the target size
        # we want to scale the longer dimension to the target size
        image_size = self.image_size * min_shape / max_shape
        self.model.transform.min_size = (image_size,)  # default is (800,)
        self.model.transform.max_size = image_size  # default is 1333

    image_tensor = to_float_tensor(image)
    image_tensor = image_tensor.to(self.device)
    prediction_result = self.model([image_tensor])

    self._original_predictions = prediction_result

Functions: