Skip to content

Detectron2 Model

sahi.models.detectron2

Detectron2 detection model wrapper for SAHI.

Provides integration with Facebook's Detectron2 framework for object detection and instance segmentation.

Classes

Detectron2DetectionModel 

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

Bases: DetectionModel

Detectron2 object detection model.

Wraps Detectron2's DefaultPredictor for detection and instance segmentation.

Initialize Detectron2 detection model.

Parameters:

Name Type Description Default
*args
object

Variable length argument list passed to DetectionModel.

 ()
**kwargs
object

Arbitrary keyword arguments passed to DetectionModel.

 {}
Source code in sahi/models/detectron2.py 
23
24
25
26
27
28
29
30
31
32
def __init__(self, *args: object, **kwargs: object) -> None:
    """Initialize Detectron2 detection model.

    Args:
        *args: Variable length argument list passed to DetectionModel.
        **kwargs: Arbitrary keyword arguments passed to DetectionModel.
    """
    existing_packages = getattr(self, "required_packages", None) or []
    self.required_packages = [*list(existing_packages), "torch", "detectron2"]
    super().__init__(*args, **kwargs)  # type: ignore[misc, arg-type]
Attributes
num_categories property 
num_categories: int

Returns number of categories.

Methods:
load_model 
load_model() -> None

Load Detectron2 model from configuration.

Source code in sahi/models/detectron2.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
76
77
78
79
80
81
82
83
84
85
86
87
def load_model(self) -> None:
    """Load Detectron2 model from configuration."""
    from detectron2.config import get_cfg
    from detectron2.data import MetadataCatalog
    from detectron2.engine import DefaultPredictor
    from detectron2.model_zoo import model_zoo

    cfg = get_cfg()

    try:  # try to load from model zoo
        config_file = model_zoo.get_config_file(self.config_path)
        cfg.set_new_allowed(True)
        cfg.merge_from_file(config_file)
        cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(self.config_path)
    except Exception as e:  # try to load from local
        print(e)
        if self.config_path is not None:
            cfg.set_new_allowed(True)
            cfg.merge_from_file(self.config_path)
        cfg.MODEL.WEIGHTS = self.model_path

    # set model device
    cfg.MODEL.DEVICE = self.device.type if hasattr(self.device, "type") else str(self.device)  # type: ignore[union-attr]
    # set input image size
    if self.image_size is not None:
        cfg.INPUT.MIN_SIZE_TEST = self.image_size
        cfg.INPUT.MAX_SIZE_TEST = self.image_size
    # init predictor
    model = DefaultPredictor(cfg)

    self.model = model

    # detectron2 category mapping
    if self.category_mapping is None:
        try:  # try to parse category names from metadata
            metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0])
            category_names = metadata.thing_classes
            self.category_names = category_names
            self.category_mapping = {
                str(ind): category_name for ind, category_name in enumerate(self.category_names)
            }
        except Exception as e:
            logger.warning(e)
            # https://detectron2.readthedocs.io/en/latest/tutorials/datasets.html#update-the-config-for-new-datasets
            if cfg.MODEL.META_ARCHITECTURE == "RetinaNet":
                num_categories = cfg.MODEL.RETINANET.NUM_CLASSES
            else:  # fasterrcnn/maskrcnn etc
                num_categories = cfg.MODEL.ROI_HEADS.NUM_CLASSES
            self.category_names = [str(category_id) for category_id in range(num_categories)]
            self.category_mapping = {
                str(ind): category_name for ind, category_name in enumerate(self.category_names)
            }
    else:
        self.category_names = list(self.category_mapping.values())
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/detectron2.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
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 RuntimeError("Model is not loaded, load it by calling .load_model()")

    if isinstance(image, np.ndarray) and self.model.input_format == "BGR":
        # convert RGB image to BGR format
        image = image[:, :, ::-1]

    prediction_result = self.model(image)

    self._original_predictions = prediction_result

Functions: