You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
gesture_rec/ml_gestures/predict.py

19 lines
676 B
Python

import joblib
import numpy as np
from ml_gestures.feature_extractor import normalize_landmarks
class MLGesturePredictor:
def __init__(self, model_path, class_names):
data = joblib.load(model_path)
self.model = data['model']
self.class_names = data['class_names']
if 'none' not in self.class_names:
self.class_names.append('none') # запасной вариант
def predict(self, landmarks):
features = normalize_landmarks(landmarks).reshape(1, -1)
pred_id = self.model.predict(features)[0]
if pred_id < len(self.class_names):
return self.class_names[pred_id]
return 'none'