add Tensorflow Hub detector support

.github/workflows/ci.yml

@@ -91,6 +91,11 @@ jobs:
       - name: Install Layer (0.10.*)
         run: >
           pip install -U "layer>=0.10.0,<0.11.0"
+
+      - name: Install tensorflow(2.9.1) and tensorflow_hub(0.12.0)
+        run: >
+          pip install  tensorflow==2.9.1
+          pip install tensorflow_hub==0.12.0
 
       - name: Unittest for SAHI+YOLOV5/MMDET/Detectron2 on all platforms
         run: |

.github/workflows/ci_torch1.10.yml

@@ -100,6 +100,13 @@ jobs:
           matrix.python-version != '3.6'
         run: >
           pip install -U "layer>=0.10.0,<0.11.0"
+
+      - name: Install tensorflow(2.9.1) and tensorflow_hub(0.12.0)
+        if: >
+            matrix.python-version != '3.6'
+        run: >
+          pip install  tensorflow==2.9.1
+          pip install tensorflow_hub==0.12.0  
 
       - name: Unittest for SAHI+YOLOV5/MMDET/Detectron2 on all platforms
         run: |

.github/workflows/package_testing.yml

@@ -91,6 +91,13 @@ jobs:
         run: >
           pip install pycocotools==2.0.4
 
+      - name: Install tensorflow(2.9.1) and tensorflow_hub(0.12.0)
+        if: >
+            matrix.python-version != '3.6'
+        run: >
+          pip install  tensorflow==2.9.1
+          pip install tensorflow_hub==0.12.0
+
       - name: Install latest SAHI package
         run: >
           pip install --upgrade --force-reinstall sahi

sahi/auto_model.py

@@ -9,6 +9,7 @@
     "detectron2": "Detectron2DetectionModel",
     "huggingface": "HuggingfaceDetectionModel",
     "torchvision": "TorchVisionDetectionModel",
+    "tensorflow": "TensorflowhubDetectionModel",
 }
 
 

sahi/model.py

@@ -2,7 +2,6 @@
 # Code written by Fatih C Akyon, 2020.
 
 import logging
-import warnings
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -93,6 +92,15 @@ def set_model(self, model: Any, **kwargs):
         """
         raise NotImplementedError()
 
+    def set_device(self):
+        """
+        Sets the device for the model.
+        """
+        import tensorflow as tf
+
+        if not (self.device):
+            self.device = "cuda:0" if is_torch_cuda_available() else "cpu"
+
     def unload_model(self):
         """
         Unloads the model from CPU/GPU.
@@ -1062,3 +1070,114 @@ 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 TensorflowhubDetectionModel(DetectionModel):
+    def set_device(self):
+        import tensorflow as tf
+
+        if not (self.device):
+            self.device = "/gpu:0" if tf.test.is_gpu_available() else "/cpu:0"
+
+    def load_model(self):
+        check_requirements(["tensorflow", "tensorflow_hub"])
+        import tensorflow as tf
+        import tensorflow_hub as hub
+
+        if "tfhub.dev/tensorflow" in self.model_path:
+            with tf.device(self.device):
+                self.model = hub.load(self.model_path)
+        else:
+            raise ValueError(
+                "Check 'https://tfhub.dev/tensorflow/collections/object_detection/' for supported TF Hub models."
+            )
+
+        if self.category_mapping is None:
+            from sahi.utils.tensorflow import COCO_CLASSES
+
+            category_mapping = {str(i): COCO_CLASSES[i] for i in range(len(COCO_CLASSES))}
+            self.category_mapping = category_mapping
+
+    def perform_inference(self, image: np.ndarray):
+        check_requirements(["tensorflow"])
+        from sahi.utils.tensorflow import resize, to_float_tensor
+
+        if self.image_size is not None:
+            img = to_float_tensor(image)
+            img = resize(img, self.image_size)
+            prediction_result = self.model(img)
+
+        else:
+            img = to_float_tensor(image)
+            prediction_result = self.model(img)
+
+        self._original_predictions = prediction_result
+        # TODO: add support for multiple image prediction
+        self.image_height, self.image_width = image.shape[0], image.shape[1]
+
+    @property
+    def num_categories(self):
+        num_categories = len(self.category_mapping)
+        return num_categories
+
+    @property
+    def has_mask(self):
+        # TODO: check if model output contains segmentation mask
+        return False
+
+    @property
+    def category_names(self):
+        return list(self.category_mapping.values())
+
+    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,
+    ):
+        import tensorflow as tf
+
+        original_predictions = self._original_predictions
+
+        # compatilibty for sahi v0.8.20
+        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]
+
+        shift_amount = shift_amount_list[0]
+        full_shape = None if full_shape_list is None else full_shape_list[0]
+
+        boxes = original_predictions["detection_boxes"][0].numpy()
+        scores = original_predictions["detection_scores"][0].numpy()
+        category_ids = original_predictions["detection_classes"][0].numpy()
+
+        # create object_prediction_list
+        object_prediction_list = []
+        object_prediction_list_per_image = []
+        with tf.device(self.device):
+            for i in range(min(boxes.shape[0], 100)):
+                if scores[i] >= self.confidence_threshold:
+                    score = float(scores[i])
+                    category_id = int(category_ids[i])
+                    category_names = self.category_mapping[str(category_id)]
+                    box = [float(box) for box in boxes[i]]
+                    x1, y1, x2, y2 = (
+                        int(box[1] * self.image_width),
+                        int(box[0] * self.image_height),
+                        int(box[3] * self.image_width),
+                        int(box[2] * self.image_height),
+                    )
+                    bbox = [x1, y1, x2, y2]
+
+                    object_prediction = ObjectPrediction(
+                        bbox=bbox,
+                        bool_mask=None,
+                        category_id=category_id,
+                        category_name=category_names,
+                        shift_amount=shift_amount,
+                        score=score,
+                        full_shape=full_shape,
+                    )
+                    object_prediction_list.append(object_prediction)
+            object_prediction_list_per_image.append(object_prediction_list)
+            self._object_prediction_list_per_image = object_prediction_list_per_image

sahi/prediction.py

@@ -8,7 +8,7 @@
 from PIL import Image
 
 from sahi.annotation import ObjectAnnotation
-from sahi.utils.coco import CocoAnnotation, CocoPrediction
+from sahi.utils.coco import CocoPrediction
 from sahi.utils.cv import read_image_as_pil, visualize_object_predictions
 from sahi.utils.file import Path
 

sahi/utils/tensorflow.py

@@ -0,0 +1,102 @@
+import numpy as np
+
+
+def to_float_tensor(image: np.ndarray):
+    import tensorflow as tf
+
+    float_image = np.asarray(image, np.float32)
+    if image.max() <= 1:
+        float_image = float_image * 255.0
+    image_tensor = tf.convert_to_tensor([np.asarray(float_image, np.uint8)], tf.uint8)
+    return image_tensor
+
+
+def resize(array, size):
+    import tensorflow as tf
+
+    return tf.image.resize(array, [size, size]).numpy()
+
+
+COCO_CLASSES = (
+    "background",
+    "person",
+    "bicycle",
+    "car",
+    "motorcycle",
+    "airplane",
+    "bus",
+    "train",
+    "truck",
+    "boat",
+    "traffic light",
+    "fire hydrant",
+    "stop sign",
+    "parking meter",
+    "bench",
+    "bird",
+    "cat",
+    "dog",
+    "horse",
+    "sheep",
+    "cow",
+    "elephant",
+    "bear",
+    "zebra",
+    "giraffe",
+    "backpack",
+    "umbrella",
+    "handbag",
+    "tie",
+    "suitcase",
+    "frisbee",
+    "skis",
+    "snowboard",
+    "sports ball",
+    "kite",
+    "baseball bat",
+    "baseball glove",
+    "skateboard",
+    "surfboard",
+    "tennis racket",
+    "bottle",
+    "wine glass",
+    "cup",
+    "fork",
+    "knife",
+    "spoon",
+    "bowl",
+    "banana",
+    "apple",
+    "sandwich",
+    "orange",
+    "broccoli",
+    "carrot",
+    "hot dog",
+    "pizza",
+    "donut",
+    "cake",
+    "chair",
+    "couch",
+    "potted plant",
+    "bed",
+    "dining table",
+    "toilet",
+    "tv",
+    "laptop",
+    "mouse",
+    "remote",
+    "keyboard",
+    "cell phone",
+    "microwave",
+    "oven",
+    "toaster",
+    "sink",
+    "refrigerator",
+    "book",
+    "clock",
+    "vase",
+    "scissors",
+    "teddy bear",
+    "hair drier",
+    "toothbrush",
+)

tests/test_tfhubmodel.py

@@ -0,0 +1,97 @@
+# OBSS SAHI Tool
+# Code written by Kadir Nar, 2022.
+
+import sys
+import unittest
+
+from sahi.model import TensorflowhubDetectionModel
+from sahi.utils.cv import read_image
+
+MODEL_DEVICE = "cpu"
+CONFIDENCE_THRESHOLD = 0.3
+IMAGE_SIZE = 320
+EFFICIENTDET_URL = "https://tfhub.dev/tensorflow/efficientdet/d0/1"
+
+# Tensorflow is not available for python<3.7
+if sys.version_info >= (3, 7):
+
+    class TestTensorflowhubDetectionModel(unittest.TestCase):
+        def test_load_model(self):
+
+            tensorflow_hub_model = TensorflowhubDetectionModel(
+                model_path=EFFICIENTDET_URL,
+                confidence_threshold=CONFIDENCE_THRESHOLD,
+                device=MODEL_DEVICE,
+                category_remapping=None,
+                load_at_init=True,
+            )
+            self.assertNotEqual(tensorflow_hub_model.model, None)
+
+        def test_perform_inference(self):
+
+            tensorflow_hub_model = TensorflowhubDetectionModel(
+                model_path=EFFICIENTDET_URL,
+                confidence_threshold=CONFIDENCE_THRESHOLD,
+                device=MODEL_DEVICE,
+                category_remapping=None,
+                load_at_init=True,
+                image_size=IMAGE_SIZE,
+            )
+
+            # prepare image
+            image_path = "tests/data/small-vehicles1.jpeg"
+            image = read_image(image_path)
+            image_height, image_width = image.shape[0], image.shape[1]
+            # perform inference
+
+            tensorflow_hub_model.perform_inference(image)
+            original_prediction = tensorflow_hub_model.original_predictions
+
+            boxes = original_prediction["detection_boxes"][0]
+            box = [float(box) for box in boxes[0].numpy()]
+            x1, y1, x2, y2 = (
+                int(box[1] * image_width),
+                int(box[0] * image_height),
+                int(box[3] * image_width),
+                int(box[2] * image_height),
+            )
+            bbox = [x1, y1, x2, y2]
+            # compare
+            desidred_bbox = [317, 324, 381, 364]
+            predicted_bbox = [x1, y1, x2, y2]
+            self.assertEqual(desidred_bbox, predicted_bbox)
+
+        def test_convert_original_predictions(self):
+
+            tensorflow_hub_model = TensorflowhubDetectionModel(
+                model_path=EFFICIENTDET_URL,
+                confidence_threshold=CONFIDENCE_THRESHOLD,
+                device=MODEL_DEVICE,
+                category_remapping=None,
+                load_at_init=True,
+                image_size=IMAGE_SIZE,
+            )
+
+            # prepare image
+            image_path = "tests/data/small-vehicles1.jpeg"
+            image = read_image(image_path)
+            image_height, image_width = image.shape[0], image.shape[1]
+
+            # perform inference
+            tensorflow_hub_model.perform_inference(image)
+
+            # convert predictions to ObjectPrediction list
+            tensorflow_hub_model.convert_original_predictions()
+            object_prediction_list = tensorflow_hub_model.object_prediction_list
+
+            # compare
+            self.assertEqual(len(object_prediction_list), 5)
+            self.assertEqual(object_prediction_list[0].category.id, 3)
+            self.assertEqual(object_prediction_list[0].category.name, "car")
+            desidred_bbox = [317, 324, 64, 40]
+            predicted_bbox = object_prediction_list[0].bbox.to_coco_bbox()
+            self.assertEqual(desidred_bbox, predicted_bbox)
+
+
+if __name__ == "__main__":
+    unittest.main()