From 48e4aeb3dc71a02aa05f880153c14429aac2b575 Mon Sep 17 00:00:00 2001
From: Fabio Manganiello <blacklight86@gmail.com>
Date: Fri, 27 Sep 2019 11:46:50 +0200
Subject: [PATCH] Added machine learning plugin based on OpenCV cv2.dnn module

---
 platypush/plugins/ml/__init__.py |  0
 platypush/plugins/ml/cv.py       | 86 ++++++++++++++++++++++++++++++++
 requirements.txt                 |  4 ++
 3 files changed, 90 insertions(+)
 create mode 100644 platypush/plugins/ml/__init__.py
 create mode 100644 platypush/plugins/ml/cv.py

diff --git a/platypush/plugins/ml/__init__.py b/platypush/plugins/ml/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/platypush/plugins/ml/cv.py b/platypush/plugins/ml/cv.py
new file mode 100644
index 000000000..2f49c105f
--- /dev/null
+++ b/platypush/plugins/ml/cv.py
@@ -0,0 +1,86 @@
+import os
+
+from platypush.plugins import Plugin, action
+
+
+class MlModel:
+    def __init__(self, model_file, classes=None):
+        import cv2
+
+        self.model_file = os.path.abspath(os.path.expanduser(model_file))
+        self.classes = classes or []
+        self.model = cv2.dnn.readNet(model_file)
+
+    def predict(self, img, resize=None, color_convert=None):
+        import cv2
+        import numpy as np
+
+        if isinstance(img, str):
+            img = cv2.imread(os.path.abspath(os.path.expanduser(img)))
+
+        if color_convert:
+            if isinstance(color_convert, str):
+                color_convert = getattr(cv2, color_convert)
+
+            img = cv2.cvtColor(img, color_convert)
+
+        if resize:
+            img = cv2.dnn.blobFromImage(img, size=tuple(resize), mean=0.5)
+        else:
+            img = cv2.dnn.blobFromImage(img, mean=0.5)
+
+        self.model.setInput(img)
+        output = self.model.forward()
+        prediction = int(np.argmax(output))
+
+        if self.classes:
+            prediction = self.classes[prediction]
+
+        return prediction
+
+
+class MlCvPlugin(Plugin):
+    """
+    Plugin to train and make computer vision predictions using machine learning models.
+
+    Requires:
+
+        * **numpy** (``pip install numpy``)
+        * **opencv** (``pip install cv2``)
+
+    Also make sure that your OpenCV installation comes with the ``dnn`` module. To test it::
+
+        >>> import cv2.dnn
+
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.models = {}
+
+    @action
+    def predict(self, img, model_file, classes=None, resize=None, color_convert=None):
+        """
+        Make predictions for an input image using a model file. Supported model formats include all the
+        types supported by cv2.dnn (currently supported: Caffe, TensorFlow, Torch, Darknet, DLDT).
+
+        :param model_file: Path to the model file
+        :param img: Path to the image
+        :param classes: List of string labels associated with the output values (e.g. ['negative', 'positive']).
+            If not set then the index of the output neuron with highest value will be returned.
+        :param resize: Tuple or list with the resize factor to be applied to the image before being fed to
+            the model (default: None)
+        :param color_convert: Color conversion to be applied to the image before being fed to the model.
+            It points to a cv2 color conversion constant (e.g. ``cv2.COLOR_BGR2GRAY``) and it can be either
+            the constant value itself or a string (e.g. 'COLOR_BGR2GRAY').
+        """
+
+        model_file = os.path.abspath(os.path.expanduser(model_file))
+
+        if model_file not in self.models:
+            self.models[model_file] = MlModel(model_file, classes=classes)
+
+        return self.models[model_file].predict(img, resize=resize, color_convert=color_convert)
+
+
+# vim:sw=4:ts=4:et:
diff --git a/requirements.txt b/requirements.txt
index bbe26db13..8c9539807 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -168,3 +168,7 @@ pyScss
 
 # Support for MLX90640 thermal camera
 # Pillow
+
+# Support for machine learning CV plugin
+# cv2
+# numpy