Add PostInferenceHook feature initial implementation and notebook

geti_sdk/deployment/deployment.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions
 # and limitations under the License.
-
+import datetime
 import json
 import logging
 import os
@@ -27,6 +27,7 @@
 from geti_sdk.rest_converters import ProjectRESTConverter
 
 from .deployed_model import DeployedModel
+from .inference_hook_interfaces import PostInferenceHookInterface
 from .utils import OVMS_README_PATH, generate_ovms_model_name
 
 
@@ -49,6 +50,7 @@ def __attrs_post_init__(self):
         self._inference_converters: Dict[str, Any] = {}
         self._path_to_temp_resources: Optional[str] = None
         self._requires_resource_cleanup: bool = False
+        self._post_inference_hooks: List[PostInferenceHookInterface] = []
 
     @property
     def is_single_task(self) -> bool:
@@ -99,6 +101,15 @@ def save(self, path_to_folder: Union[str, os.PathLike]) -> bool:
         with open(project_filepath, "w") as project_file:
             json.dump(project_dict, project_file, indent=4)
 
+        # Save post inference hooks, if any
+        if self.post_inference_hooks:
+            hook_config_file = os.path.join(deployment_folder, "hook_config.json")
+            hook_configs: List[Dict[str, Any]] = []
+            for hook in self.post_inference_hooks:
+                hook_configs.append(hook.to_dict())
+            with open(hook_config_file, "w") as file:
+                json.dump({"post_inference_hooks": hook_configs}, file)
+
         # Clean up temp resources if needed
         if self._requires_resource_cleanup:
             self._remove_temporary_resources()
@@ -135,7 +146,23 @@ def from_folder(cls, path_to_folder: Union[str, os.PathLike]) -> "Deployment":
             models.append(
                 DeployedModel.from_folder(os.path.join(deployment_folder, task_folder))
             )
-        return cls(models=models, project=project)
+        deployment = cls(models=models, project=project)
+
+        # Load post inference hooks, if any
+        hook_config_file = os.path.join(deployment_folder, "hook_config.json")
+        if os.path.isfile(hook_config_file):
+            available_hooks = {
+                subcls.__name__: subcls
+                for subcls in PostInferenceHookInterface.__subclasses__()
+            }
+            with open(hook_config_file, "r") as file:
+                hook_dict = json.load(file)
+            for hook_data in hook_dict["post_inference_hooks"]:
+                for hook_name, hook_args in hook_data.items():
+                    target_hook = available_hooks[hook_name]
+                    hook = target_hook.from_dict(hook_args)
+                deployment.add_post_inference_hook(hook)
+        return deployment
 
     def load_inference_models(self, device: str = "CPU"):
         """
@@ -153,13 +180,15 @@ def load_inference_models(self, device: str = "CPU"):
         self._are_models_loaded = True
         logging.info(f"Inference models loaded on device `{device}` successfully.")
 
-    def infer(self, image: np.ndarray) -> Prediction:
+    def infer(self, image: np.ndarray, name: Optional[str] = None) -> Prediction:
         """
         Run inference on an image for the full model chain in the deployment.
 
         :param image: Image to run inference on, as a numpy array containing the pixel
             data. The image is expected to have dimensions [height x width x channels],
             with the channels in RGB order
+        :param name: Optional name for the image, if specified this will be used in
+            any post inference hooks belonging to the deployment.
         :return: inference results
         """
         self._check_models_loaded()
@@ -172,9 +201,12 @@ def infer(self, image: np.ndarray) -> Prediction:
         # Multi-task inference
         else:
             prediction = self._infer_pipeline(image=image, explain=False)
+        self._execute_post_inference_hooks(
+            image=image, prediction=prediction, name=name
+        )
         return prediction
 
-    def explain(self, image: np.ndarray) -> Prediction:
+    def explain(self, image: np.ndarray, name: Optional[str] = None) -> Prediction:
         """
         Run inference on an image for the full model chain in the deployment. The
         resulting prediction will also contain saliency maps and the feature vector
@@ -183,6 +215,8 @@ def explain(self, image: np.ndarray) -> Prediction:
         :param image: Image to run inference on, as a numpy array containing the pixel
             data. The image is expected to have dimensions [height x width x channels],
             with the channels in RGB order
+        :param name: Optional name for the image, if specified this will be used in
+            any post inference hooks belonging to the deployment.
         :return: inference results
         """
         self._check_models_loaded()
@@ -195,6 +229,9 @@ def explain(self, image: np.ndarray) -> Prediction:
         # Multi-task inference
         else:
             prediction = self._infer_pipeline(image=image, explain=True)
+        self._execute_post_inference_hooks(
+            image=image, prediction=prediction, name=name
+        )
         return prediction
 
     def _check_models_loaded(self) -> None:
@@ -437,3 +474,52 @@ def generate_ovms_config(self, output_folder: Union[str, os.PathLike]) -> None:
             f"file with instructions on how to launch OVMS, connect to it and run "
             f"inference. Please follow the instructions outlined there to get started."
         )
+
+    @property
+    def post_inference_hooks(self) -> List[PostInferenceHookInterface]:
+        """
+        Return the currently active post inference hooks for the deployment
+
+        :return: list of PostInferenceHook objects
+        """
+        return self._post_inference_hooks
+
+    def clear_inference_hooks(self) -> None:
+        """
+        Remove all post inference hooks for the deployment
+        """
+        n_hooks = len(self.post_inference_hooks)
+        self._post_inference_hooks = []
+        if n_hooks != 0:
+            logging.info(
+                f"Post inference hooks cleared. {n_hooks} hooks were removed "
+                f"successfully"
+            )
+
+    def add_post_inference_hook(self, hook: PostInferenceHookInterface) -> None:
+        """
+        Add a post inference hook, which will be executed after each call to
+        `Deployment.infer`
+
+        :param hook: PostInferenceHook to be added to the deployment
+        """
+        self._post_inference_hooks.append(hook)
+        logging.info(f"Hook `{hook}` added.")
+        logging.info(
+            f"Deployment now contains {len(self.post_inference_hooks)} "
+            f"post inference hooks."
+        )
+
+    def _execute_post_inference_hooks(
+        self, image: np.ndarray, prediction: Prediction, name: Optional[str] = None
+    ) -> None:
+        """
+        Execute all post inference hooks
+
+        :param image: Numpy image which was inferred
+        :param prediction: Prediction for the image
+        :param name: Optional name for the image
+        """
+        timestamp = datetime.datetime.now()
+        for hook in self._post_inference_hooks:
+            hook.run(image, prediction, name, timestamp)