add detectron2 support

.github/workflows/ci.yml

@@ -77,14 +77,21 @@ jobs:
         if: matrix.operating-system == 'macos-latest'
         run: pip install torch==1.10.0 torchvision==0.11.1
 
-      - name: Install MMDetection(2.19.1), YOLOv5(6.0.6) and Norfair(0.3.1)
+      - name: Install MMDetection(2.19.1) with MMCV(1.4.1)
         run: >
           pip install mmcv-full==1.4.1 -f https://download.openmmlab.com/mmcv/dist/cpu/torch1.10.0/index.html
-          pip install mmdet==2.19.1 yolov5==6.0.6 norfair==0.3.1
+          pip install mmdet==2.19.1
 
-      - name: Test with unittest
+      - name: Install YOLOv5(6.0.6) and Norfair(0.3.1)
+        run: >
+          pip install yolov5==6.0.6 norfair==0.3.1
+
+      - name: Install Detectron2(0.6)
+        run: >
+          python -m pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cpu/torch1.10/index.html
+
+      - name: Unittest for SAHI+YOLOV5/MMDET/Detectron2 on all platforms
         run: |
-          pip install pytest
           python -m unittest
 
       - name: Install pycocotools(2.0.3)
@@ -113,4 +120,4 @@ jobs:
           # coco evaluate
           sahi coco evaluate --dataset_json_path tests/data/coco_evaluate/dataset.json --result_json_path tests/data/coco_evaluate/result.json
           # coco analyse
-          sahi coco analyse --dataset_json_path tests/data/coco_evaluate/dataset.json --result_json_path tests/data/coco_evaluate/result.json --out_dir tests/data/coco_evaluate/
+          sahi coco analyse --dataset_json_path tests/data/coco_evaluate/dataset.json --result_json_path tests/data/coco_evaluate/result.json --out_dir tests/data/coco_evaluate/

.github/workflows/package_testing.yml

@@ -62,10 +62,18 @@ jobs:
         if: matrix.operating-system == 'macos-latest'
         run: pip install torch==1.10.0 torchvision==0.11.1
 
-      - name: Install MMDetection(2.19.1), YOLOv5(6.0.6) and Norfair(0.3.1)
+      - name: Install MMDetection(2.19.1) with MMCV(1.4.1)
         run: >
           pip install mmcv-full==1.4.1 -f https://download.openmmlab.com/mmcv/dist/cpu/torch1.10.0/index.html
-          pip install mmdet==2.19.1 yolov5==6.0.6 norfair==0.3.1
+          pip install mmdet==2.19.1
+
+      - name: Install YOLOv5(6.0.6) and Norfair(0.3.1)
+        run: >
+          pip install yolov5==6.0.6 norfair==0.3.1
+
+      - name: Install Detectron2(0.6)
+        run: >
+          python -m pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cpu/torch1.10/index.html
 
       - name: Install pycocotools(2.0.3)
         run: >
@@ -75,9 +83,8 @@ jobs:
         run: >
           pip install --upgrade --force-reinstall sahi
 
-      - name: Test with unittest
+      - name: Unittest for SAHI+YOLOV5/MMDET/Detectron2 on all platforms
         run: |
-          pip install pytest
           python -m unittest
 
       - name: Test SAHI CLI

requirements.txt

@@ -5,4 +5,4 @@ pillow>=8.2.0
 pyyaml
 fire
 terminaltables
-requests
+requests

sahi/annotation.py

@@ -535,6 +535,8 @@ def __init__(
             # https://github.com/obss/sahi/issues/235
             if bbox_from_bool_mask is not None:
                 bbox = bbox_from_bool_mask
+            else:
+                raise ValueError("Invalid boolean mask.")
         else:
             self.mask = None
 

sahi/model.py

@@ -2,14 +2,14 @@
 # Code written by Fatih C Akyon, 2020.
 
 import logging
-import os
 import warnings
 from typing import Dict, List, Optional, Union
 
 import numpy as np
 
 from sahi.prediction import ObjectPrediction
 from sahi.utils.compatibility import fix_full_shape_list, fix_shift_amount_list
+from sahi.utils.cv import get_bbox_from_bool_mask
 from sahi.utils.torch import cuda_is_available, empty_cuda_cache
 
 logger = logging.getLogger(__name__)
@@ -89,7 +89,6 @@ def perform_inference(self, image: np.ndarray, image_size: int = None):
         """
         This function should be implemented in a way that prediction should be
         performed using self.model and the prediction result should be set to self._original_predictions.
-
         Args:
             image: np.ndarray
                 A numpy array that contains the image to be predicted.
@@ -107,7 +106,6 @@ def _create_object_prediction_list_from_original_predictions(
         This function should be implemented in a way that self._original_predictions should
         be converted to a list of prediction.ObjectPrediction and set to
         self._object_prediction_list. self.mask_threshold can also be utilized.
-
         Args:
             shift_amount: list
                 To shift the box and mask predictions from sliced image to full sized image, should be in the form of [shift_x, shift_y]
@@ -137,7 +135,6 @@ def convert_original_predictions(
         """
         Converts original predictions of the detection model to a list of
         prediction.ObjectPrediction object. Should be called after perform_inference().
-
         Args:
             shift_amount: list
                 To shift the box and mask predictions from sliced image to full sized image, should be in the form of [shift_x, shift_y]
@@ -200,7 +197,6 @@ def load_model(self):
     def perform_inference(self, image: np.ndarray, image_size: int = 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.
@@ -271,7 +267,6 @@ def _create_object_prediction_list_from_original_predictions(
         """
         self._original_predictions is converted to a list of prediction.ObjectPrediction and set to
         self._object_prediction_list_per_image.
-
         Args:
             shift_amount_list: list of list
                 To shift the box and mask predictions from sliced image to full sized image, should
@@ -384,7 +379,6 @@ def load_model(self):
     def perform_inference(self, image: np.ndarray, image_size: int = 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.
@@ -436,7 +430,6 @@ def _create_object_prediction_list_from_original_predictions(
         """
         self._original_predictions is converted to a list of prediction.ObjectPrediction and set to
         self._object_prediction_list_per_image.
-
         Args:
             shift_amount_list: list of list
                 To shift the box and mask predictions from sliced image to full sized image, should
@@ -495,3 +488,178 @@ def _create_object_prediction_list_from_original_predictions(
             object_prediction_list_per_image.append(object_prediction_list)
 
         self._object_prediction_list_per_image = object_prediction_list_per_image
+
+
+class Detectron2DetectionModel(DetectionModel):
+    def load_model(self):
+        try:
+            import detectron2
+        except ImportError:
+            raise ImportError(
+                "Please install detectron2. Check "
+                "`https://detectron2.readthedocs.io/en/latest/tutorials/install.html` "
+                "for instalattion details."
+            )
+
+        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()
+        cfg.MODEL.DEVICE = self.device
+
+        try:  # try to load from model zoo
+            config_file = model_zoo.get_config_file(self.config_path)
+            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.merge_from_file(self.config_path)
+            cfg.MODEL.WEIGHTS = self.model_path
+        # 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())
+
+    def perform_inference(self, image: np.ndarray, image_size: int = 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.
+        """
+        try:
+            import detectron2
+        except ImportError:
+            raise ImportError("Please install detectron2 via `pip install detectron2`")
+
+        # confirm image_size is not provided
+        if image_size is not None:
+            warnings.warn("Set 'image_size' at DetectionModel init.")
+
+        # 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
+
+    @property
+    def num_categories(self):
+        """
+        Returns number of categories
+        """
+        num_categories = len(self.category_mapping)
+        return num_categories
+
+    def _create_object_prediction_list_from_original_predictions(
+        self,
+        shift_amount_list: Optional[List[List[int]]] = [[0, 0]],
+        full_shape_list: Optional[List[List[int]]] = None,
+    ):
+        """
+        self._original_predictions is converted to a list of prediction.ObjectPrediction and set to
+        self._object_prediction_list_per_image.
+        Args:
+            shift_amount_list: list of list
+                To shift the box and mask predictions from sliced image to full sized image, should
+                be in the form of List[[shift_x, shift_y],[shift_x, shift_y],...]
+            full_shape_list: list of list
+                Size of the full image after shifting, should be in the form of
+                List[[height, width],[height, width],...]
+        """
+        original_predictions = self._original_predictions
+        category_mapping = self.category_mapping
+
+        # compatilibty for sahi v0.8.15
+        if isinstance(shift_amount_list[0], int):
+            shift_amount_list = [shift_amount_list]
+        if full_shape_list is not None and isinstance(full_shape_list[0], int):
+            full_shape_list = [full_shape_list]
+
+        # parse boxes, masks, scores, category_ids from predictions
+        boxes = original_predictions["instances"].pred_boxes.tensor.tolist()
+        scores = original_predictions["instances"].scores.tolist()
+        category_ids = original_predictions["instances"].pred_classes.tolist()
+        try:
+            masks = original_predictions["instances"].pred_masks.tolist()
+        except AttributeError:
+            masks = None
+
+        # create object_prediction_list
+        num_categories = self.num_categories
+        object_prediction_list_per_image = []
+        object_prediction_list = []
+
+        # detectron2 DefaultPredictor supports single image
+        shift_amount = shift_amount_list[0]
+        full_shape = None if full_shape_list is None else full_shape_list[0]
+
+        for ind in range(len(boxes)):
+            score = scores[ind]
+            if score < self.confidence_threshold:
+                continue
+
+            category_id = category_ids[ind]
+
+            if masks is None:
+                bbox = boxes[ind]
+                mask = None
+            else:
+                mask = np.array(masks[ind])
+
+                # check if mask is valid
+                if get_bbox_from_bool_mask(mask) is not None:
+                    bbox = None
+                else:
+                    continue
+
+            object_prediction = ObjectPrediction(
+                bbox=bbox,
+                bool_mask=mask,
+                category_id=category_id,
+                category_name=self.category_mapping[str(category_id)],
+                shift_amount=shift_amount,
+                score=score,
+                full_shape=full_shape,
+            )
+            object_prediction_list.append(object_prediction)
+
+        # detectron2 DefaultPredictor supports single image
+        object_prediction_list_per_image = [object_prediction_list]
+
+        self._object_prediction_list_per_image = object_prediction_list_per_image

sahi/predict.py

@@ -27,6 +27,7 @@
 MODEL_TYPE_TO_MODEL_CLASS_NAME = {
     "mmdet": "MmdetDetectionModel",
     "yolov5": "Yolov5DetectionModel",
+    "detectron2": "Detectron2DetectionModel",
 }
 
 LOW_MODEL_CONFIDENCE = 0.1

sahi/utils/cv.py

@@ -434,8 +434,11 @@ def get_bbox_from_bool_mask(bool_mask):
 
     ymin, ymax = np.where(rows)[0][[0, -1]]
     xmin, xmax = np.where(cols)[0][[0, -1]]
-    # width = xmax - xmin
-    # height = ymax - ymin
+    width = xmax - xmin
+    height = ymax - ymin
+
+    if width == 0 or height == 0:
+        return None
 
     return [xmin, ymin, xmax, ymax]
 

sahi/utils/detectron2.py

@@ -0,0 +1,21 @@
+from pathlib import Path
+
+
+class Detectron2TestConstants:
+    FASTERCNN_MODEL_ZOO_NAME = "COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"
+    RETINANET_MODEL_ZOO_NAME = "COCO-Detection/retinanet_R_50_FPN_3x.yaml"
+    MASKRCNN_MODEL_ZOO_NAME = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
+
+
+def export_cfg_as_yaml(cfg, export_path: str = "config.yaml"):
+    """
+    Exports Detectron2 config object in yaml format so that it can be used later.
+    Args:
+        cfg (detectron2.config.CfgNode): Detectron2 config object.
+        export_path (str): Path to export the Detectron2 config.
+    Related Detectron2 doc: https://detectron2.readthedocs.io/en/stable/modules/config.html#detectron2.config.CfgNode.dump
+    """
+    Path(export_path).parent.mkdir(exist_ok=True, parents=True)
+
+    with open(export_path, "w") as f:
+        f.write(cfg.dump())