IDKHowToCodeFr/tinyml-backend

1import joblib2import numpy as np3import os4import sys5 6sys.path.append(os.path.dirname(os.path.abspath(__file__)))7from preprocessing import resolve_model_dir8 9class EnsembleModel:10    def __init__(self):11        self.model_dir = resolve_model_dir()12        self.model_names = ['knn', 'svm', 'logreg', 'rf', 'small_nn']13        self.models = {}14        self.load_models()15        self.label_encoder = joblib.load(f'{self.model_dir}/label_encoder.pkl')16        self.weights = {'knn': 0.15, 'svm': 0.15, 'logreg': 0.20, 'rf': 0.25, 'small_nn': 0.25}17        18    def load_models(self):19        for name in self.model_names:20            path = f'{self.model_dir}/{name}.pkl'21            if os.path.exists(path):22                self.models[name] = joblib.load(path)23                24    def predict(self, X):25        individual_preds = {}26        individual_probs = {}27        28        for name, model in self.models.items():29            probs = model.predict_proba(X)30            individual_probs[name] = probs[0]31            pred_indices = np.argmax(probs, axis=1)32            individual_preds[name] = str(self.label_encoder.inverse_transform(pred_indices)[0])33            34        num_classes = list(individual_probs.values())[0].shape[0]35        weighted_probs = np.zeros(num_classes)36        37        total_weight = 038        for name in self.models.keys():39            w = self.weights.get(name, 1.0)40            weighted_probs += individual_probs[name] * w41            total_weight += w42            43        weighted_probs /= total_weight44        final_pred_idx = np.argmax(weighted_probs)45        final_pred = str(self.label_encoder.inverse_transform([final_pred_idx])[0])46        confidence = float(np.max(weighted_probs))47        48        class_probs = {49            str(self.label_encoder.inverse_transform([i])[0]): float(prob)50            for i, prob in enumerate(weighted_probs)51        }52        53        individual_conf = {54            name: float(np.max(probs)) for name, probs in individual_probs.items()55        }56        57        return final_pred, confidence, individual_preds, class_probs, self.weights, individual_conf58