app/models/model_loader.py - master - sofnuk2004/server1

app/models/model_loader.py
150 строк8 KB
Sofya Nikolaeva
Save
19 мар 2026, 20:54
19 мар 2026, 20:548648371
    import os
    import json
    from app.models.dl_models import MobileNetV2Classifier, ResNet50Classifier

    class ModelLoader:
        """Класс для загрузки и управления моделями"""
        
        _instance = None
        _models = {}
        
        def __new__(cls):
            if cls._instance is None:
                cls._instance = super(ModelLoader, cls).__new__(cls)
            return cls._instance
        
        def get_models_base_path(self):
            """Получение пути к папке с моделями"""
            current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            models_path = os.path.join(current_dir, 'static', 'models')
            return models_path
        
        def load_model(self, model_name):
            """Загрузка модели по имени"""
            if model_name not in self._models:
                models_base = self.get_models_base_path()
                
                # Проверяем существование папки модели
                model_dir = os.path.join(models_base, model_name)
                if not os.path.exists(model_dir):
                    print(f"❌ Папка модели не найдена: {model_dir}")
                    # Пробуем найти модель в подпапках
                    for root, dirs, files in os.walk(models_base):
                        if model_name.lower() in root.lower():
                            model_dir = root
                            break
                
                if model_name.lower() == 'mobilenetv2':
                    model_path = os.path.join(model_dir, 'final_mobilenetv2_model.h5')
                    class_path = os.path.join(model_dir, 'data_generators_info_binary.json')
                    
                    if os.path.exists(model_path):
                        model = MobileNetV2Classifier(model_path, class_path)
                        self._models[model_name] = model
                        print(f"✅ Модель {model_name} загружена из {model_path}")
                    else:
                        print(f"❌ Модель не найдена по пути: {model_path}")
                        
                elif model_name.lower() == 'resnet50':
                    model_path = os.path.join(model_dir, 'final_resnet50_fixed.h5')
                    class_path = os.path.join(model_dir, 'data_generators_info.json')
                    
                    if os.path.exists(model_path):
                        model = ResNet50Classifier(model_path, class_path)
                        self._models[model_name] = model
                        print(f"✅ Модель {model_name} загружена из {model_path}")
                    else:
                        print(f"❌ Модель не найдена по пути: {model_path}")
                
                else:
                    # Поиск модели по имени в папках
                    for root, dirs, files in os.walk(models_base):
                        if model_name in root or model_name in dirs:
                            model_dir = root if os.path.isfile(os.path.join(root, 'model.keras')) else os.path.join(root, model_name)
                            keras_files = [f for f in os.listdir(model_dir) if f.endswith('.keras')]
                            if keras_files:
                                model_path = os.path.join(model_dir, keras_files[0])
                                json_files = [f for f in os.listdir(model_dir) if f.endswith('.json')]
                                class_path = os.path.join(model_dir, json_files[0]) if json_files else None
                                
                                model = MobileNetV2Classifier(model_path, class_path)
                                self._models[model_name] = model
                                print(f"✅ Модель {model_name} загружена из {model_path}")
                                break
            
            return self._models.get(model_name)
        
        def get_model(self, model_name):
            """Получение загруженной модели"""
            return self._models.get(model_name)
        
        def unload_model(self, model_name):
            """Выгрузка модели для освобождения памяти"""
            if model_name in self._models:
                del self._models[model_name]
                print(f"📤 Модель {model_name} выгружена")
        
        def list_models(self):
            """Список загруженных моделей"""
            return list(self._models.keys())
        
        def get_available_models(self):
            """Получение списка доступных моделей в файловой системе"""
            models_base = self.get_models_base_path()
            available = []
            
            if os.path.exists(models_base):
                for model_dir in os.listdir(models_base):
                    model_path = os.path.join(models_base, model_dir)
                    if os.path.isdir(model_path):
                        # Ищем .h5 файлы (предпочтительно) и .keras
                        h5_files = [f for f in os.listdir(model_path) if f.endswith('.h5')]
                        keras_files = [f for f in os.listdir(model_path) if f.endswith('.keras')]
                        json_files = [f for f in os.listdir(model_path) if f.endswith('.json')]
                        
                        # Берем первый .h5 файл, если есть, иначе .keras
                        model_file = None
                        if h5_files:
                            model_file = h5_files[0]
                        elif keras_files:
                            model_file = keras_files[0]
                        
                        if model_file:
                            available.append({
                                'name': model_dir,
                                'path': model_path,
                                'model_file': model_file,
                                'class_file': json_files[0] if json_files else None,
                                'full_path': model_path
                            })
            
            return available
        
        def get_model_info(self, model_name):
            """Получение информации о конкретной модели"""
            models_base = self.get_models_base_path()
            model_path = os.path.join(models_base, model_name)
            
            if os.path.exists(model_path) and os.path.isdir(model_path):
                keras_files = [f for f in os.listdir(model_path) if f.endswith('.keras')]
                json_files = [f for f in os.listdir(model_path) if f.endswith('.json')]
                
                # Попытка прочитать информацию из JSON
                class_info = {}
                if json_files:
                    try:
                        with open(os.path.join(model_path, json_files[0]), 'r') as f:
                            class_info = json.load(f)
                    except:
                        pass
                
                return {
                    'name': model_name,
                    'path': model_path,
                    'model_file': keras_files[0] if keras_files else None,
                    'class_file': json_files[0] if json_files else None,
                    'class_info': class_info,
                    'is_loaded': model_name in self._models
                }
            
            return None