forked from platypush/platypush
86 lines
2.8 KiB
Python
86 lines
2.8 KiB
Python
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:
|