Source code for platypush.plugins.ml.cv

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


[docs]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 """
[docs] def __init__(self, **kwargs): super().__init__(**kwargs) self.models = {}
[docs] @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: