Optimized code. Improved import of yolo and detectron2 models. Improved work with language models.

Author: rybakov-ks

particleanalyzer/core/Detectron2Loader.py

@@ -9,123 +9,88 @@
 logger = setup_logger()
 logger.setLevel("ERROR")
 logging.disable(logging.CRITICAL)
-
-# Скрываем предупреждения PyTorch
 warnings.filterwarnings("ignore", category=UserWarning)
-"""Работаем с моделями Detectron2"""
 
+"""Работаем с моделями Detectron2"""
 
 class Detectron2Loader:
-    def __init__(self, device=None):
-
-        base_path = os.path.dirname(__file__)
-        self.model_path = lambda name: os.path.join(base_path, "..", "model", name)
-
-        if device is None:
-            self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        elif isinstance(device, torch.device):
-            self.device = device.type  # Преобразуем torch.device в строку
-        else:
-            self.device = str(device)
-        # Инициализация конфигураций
-        self.configs = {
-            "R101": self._init_r101_config(),
-            "X101": self._init_x101_config(),
-            "Cascade_R50": self._init_cascade_r50_config(),
-            "Cascade_X152": self._init_cascade_x152_config(),
-        }
-
-        self.config_paths = {
-            "R101": self.model_path("faster_rcnn_R_101_FPN_3x.yaml"),
-            "X101": self.model_path("faster_rcnn_X_101_32x8d_FPN_3x.yaml"),
-            "Cascade_R50": self.model_path("cascade_mask_rcnn_R_50_FPN_3x.yaml"),
-            "Cascade_X152": self.model_path(
-                "cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.yaml"
-            ),
-        }
-
-        self.model_paths = {
-            "R101": self.model_path("/faster_rcnn_R_101_FPN_3x.pth"),
-            "X101": self.model_path("faster_rcnn_X_101_32x8d_FPN_3x.pth"),
-            "Cascade_R50": self.model_path("cascade_mask_rcnn_R_50_FPN_3x.pth"),
-            "Cascade_X152": self.model_path(
-                "cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.pth"
-            ),
+    MODEL_MAPPING = {
+        "R101": {
+            "config_file": "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml",
+            "weights_file": "faster_rcnn_R_101_FPN_3x.pth",
+            "config_path": "faster_rcnn_R_101_FPN_3x.yaml"
+        },
+        "X101": {
+            "config_file": "COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml",
+            "weights_file": "faster_rcnn_X_101_32x8d_FPN_3x.pth",
+            "config_path": "faster_rcnn_X_101_32x8d_FPN_3x.yaml"
+        },
+        "Cascade_R50": {
+            "config_file": "Misc/cascade_mask_rcnn_R_50_FPN_3x.yaml",
+            "weights_file": "cascade_mask_rcnn_R_50_FPN_3x.pth",
+            "config_path": "cascade_mask_rcnn_R_50_FPN_3x.yaml"
+        },
+        "Cascade_X152": {
+            "config_file": "Misc/cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.yaml",
+            "weights_file": "cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.pth",
+            "config_path": "cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.yaml"
         }
+    }
 
-        # Сохраняем конфигурации в файлы
-        self._save_configs()
+    def __init__(self, device=None):
+        self._base_path = os.path.join(os.path.dirname(__file__), "..", "model")
+        self.device = self._get_device(device)
+        self.configs = {}
+        self._init_models()
 
-    def _init_r101_config(self):
-        cfg = get_cfg()
-        cfg.merge_from_file(
-            model_zoo.get_config_file(
-                "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"
-            )
-        )
-        cfg.OUTPUT_DIR = self.model_path("")
-        cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "faster_rcnn_R_101_FPN_3x.pth")
-        cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
-        cfg.MODEL.DEVICE = self.device
-        return cfg
+    def _get_device(self, device):
+        if device is None:
+            return "cuda" if torch.cuda.is_available() else "cpu"
+        return device.type if isinstance(device, torch.device) else str(device)
 
-    def _init_x101_config(self):
-        cfg = get_cfg()
-        cfg.merge_from_file(
-            model_zoo.get_config_file(
-                "COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml"
-            )
-        )
-        cfg.OUTPUT_DIR = self.model_path("")
-        cfg.MODEL.WEIGHTS = os.path.join(
-            cfg.OUTPUT_DIR, "faster_rcnn_X_101_32x8d_FPN_3x.pth"
-        )
-        cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
-        cfg.MODEL.DEVICE = self.device
-        return cfg
+    def _model_path(self, name: str) -> str:
+        return os.path.join(self._base_path, name)
 
-    def _init_cascade_r50_config(self):
+    def _init_model_config(self, model_name):
         cfg = get_cfg()
-        cfg.merge_from_file(
-            model_zoo.get_config_file("Misc/cascade_mask_rcnn_R_50_FPN_3x.yaml")
-        )
-        cfg.OUTPUT_DIR = self.model_path("")
-        cfg.MODEL.WEIGHTS = os.path.join(
-            cfg.OUTPUT_DIR, "cascade_mask_rcnn_R_50_FPN_3x.pth"
-        )
+        model_data = self.MODEL_MAPPING[model_name]
+        
+        cfg.merge_from_file(model_zoo.get_config_file(model_data["config_file"]))
+        cfg.OUTPUT_DIR = self._base_path
+        cfg.MODEL.WEIGHTS = self._model_path(model_data["weights_file"])
         cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
         cfg.MODEL.DEVICE = self.device
+        
         return cfg
 
-    def _init_cascade_x152_config(self):
-        cfg = get_cfg()
-        cfg.merge_from_file(
-            model_zoo.get_config_file(
-                "Misc/cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.yaml"
-            )
-        )
-        cfg.OUTPUT_DIR = self.model_path("")
-        cfg.MODEL.WEIGHTS = os.path.join(
-            cfg.OUTPUT_DIR, "cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.pth"
-        )
-        cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
-        cfg.MODEL.DEVICE = self.device
-        return cfg
+    def _init_models(self):
+        self.configs = {
+            name: self._init_model_config(name)
+            for name in self.MODEL_MAPPING
+        }
+        
+        self.config_paths = {
+            name: self._model_path(self.MODEL_MAPPING[name]["config_path"])
+            for name in self.MODEL_MAPPING
+        }
+        
+        self.model_paths = {
+            name: self._model_path(self.MODEL_MAPPING[name]["weights_file"])
+            for name in self.MODEL_MAPPING
+        }
+        
+        self._save_configs()
 
     def _save_configs(self):
-        """Сохраняет конфигурации в файлы"""
         for model_name, cfg in self.configs.items():
             with open(self.config_paths[model_name], "w") as f:
                 f.write(cfg.dump())
 
     def get_config(self, model_name: str):
-        """Возвращает конфигурацию модели"""
         return self.configs.get(model_name)
 
     def get_config_path(self, model_name: str):
-        """Возвращает путь к файлу конфигурации"""
         return self.config_paths.get(model_name)
 
     def get_model_path(self, model_name: str):
-        """Возвращает путь к весам модели"""
-        return self.model_paths.get(model_name)
+        return self.model_paths.get(model_name)

particleanalyzer/core/LLMAnalysis.py

@@ -1,11 +1,10 @@
 import json
-from typing import Dict, List, Tuple, Literal, Optional
+from typing import Dict, List, Tuple, Literal
 import pandas as pd
 import numpy as np
 from openai import OpenAI
 from huggingface_hub import InferenceClient
 from particleanalyzer.core.language_context import LanguageContext
-from particleanalyzer.core.languages import translations
 
 
 class LLMAnalysis:
@@ -17,13 +16,18 @@ def __init__(
         self.provider = provider
         self.api_key = api_key
 
-        if provider == "openrouter":
+        if self.api_key.startswith("hf_"):
+            provider == "huggingface"
+            self.client = InferenceClient(provider="fireworks-ai", api_key=api_key)
+            self.model_list = ["deepseek-ai/DeepSeek-V3"]
+        elif self.api_key.startswith("sk-or-"):
             self.client = OpenAI(
                 base_url="https://openrouter.ai/api/v1",
                 api_key=api_key,
             )
-        elif provider == "huggingface":
-            self.client = InferenceClient(provider=huggingface_model, api_key=api_key)
+            self.model_list = ["deepseek/deepseek-chat:free", "deepseek/deepseek-chat-v3-0324", "google/gemini-2.0-flash-001", 
+                                "openai/gpt-4o-mini"]
+            provider == "openrouter"
         else:
             raise ValueError("Неизвестный провайдер. Доступные варианты: 'openrouter', 'huggingface'")
 

particleanalyzer/core/ModelManager.py

@@ -6,7 +6,6 @@
 try:
     from detectron2.engine import DefaultPredictor
     from .Detectron2Loader import Detectron2Loader
-
     DETECTRON2_AVAILABLE = True
 except ImportError:
     DETECTRON2_AVAILABLE = False
@@ -15,101 +14,36 @@
 class ModelManager:
     def __init__(self, device=None):
         self.device = device
-        # URL сервера
         self.SERVER_URL = "https://rybakov-k.ru/model/"
-
-        # Путь к моделям
+        
+        # Инициализация путей
         base_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
         self.MODELS_DIR = os.path.join(base_path, "model")
         os.makedirs(self.MODELS_DIR, exist_ok=True)
 
-        # Проверяем и загружаем модели
-        self._ensure_models_available()
-
-        # Инициализируем YOLO
-        self.yolo_loader = YOLOLoader()
-
-        # Инициализируем detectron_loader только если detectron2 доступен
+        yolo_files = list(YOLOLoader.MODEL_MAPPING.values())
+        detectron_files = []
         if DETECTRON2_AVAILABLE:
-            self.detectron_loader = Detectron2Loader(device=self.device)
-        else:
-            self.detectron_loader = None
-
-        self.model_types = {
-            "Yolo11 (dataset 1)": "yolo",
-            "Yolo12 (dataset 1)": "yolo",
-            "Yolo11 (dataset 2)": "yolo",
-            "Yolo12 (dataset 2)": "yolo",
-            "R101": "detectron",
-            "X101": "detectron",
-            "Cascade_R50": "detectron",
-            "Cascade_X152": "detectron",
-        }
-
-    def get_model(self, model_name: str):
-        """Возвращает модель по имени"""
-        model_type = self.model_types.get(model_name)
-
-        if model_type == "yolo":
-            return self.yolo_loader.get_model(model_name)
-        elif model_type == "detectron":
-            return self.detectron_loader.get_config(model_name)
-        else:
-            raise ValueError(f"Unknown model type: {model_name}")
+            detectron_files = [
+                *[v["weights_file"] for v in Detectron2Loader.MODEL_MAPPING.values()]
+            ]
 
-    def get_predictor(self, model_name: str):
-        """Для Detectron возвращает готовый predictor"""
-        if model_name in self.detectron_loader.configs:
-            cfg = self.detectron_loader.get_config(model_name)
-            return DefaultPredictor(cfg)
-        return None
-
-    def get_model_path(self, model_name: str) -> str:
-        """Возвращает путь к модели по её имени"""
-        if model_name in self.yolo_loader.models:
-            return self.yolo_loader.get_model_path(model_name)
-        elif model_name in self.detectron_loader.configs:
-            return self.detectron_loader.get_model_path(model_name)
-        else:
-            raise ValueError(f"Model {model_name} not found")
-
-    def get_config_path(self, model_name: str) -> str:
-        """Возвращает путь к конфигу (только для Detectron)"""
-        if model_name in self.detectron_loader.configs:
-            return self.detectron_loader.get_config_path(model_name)
-        raise ValueError(
-            f"Config for {model_name} not available (YOLO models don't use config files)"
-        )
-
-    def _ensure_models_available(self):
-        """Проверяет и загружает только необходимые модели"""
-        # Базовые файлы для YOLO (всегда нужны)
-        required_files = [
-            "Yolo11_d1.pt",
-            "Yolo11_d2.pt",
-            "Yolo12_d1.pt",
-            "Yolo12_d2.pt",
-        ]
+        # Проверяем и загружаем модели
+        self._ensure_models_available(yolo_files + detectron_files)
 
-        # Добавляем модели Detectron2 только если доступен
-        if DETECTRON2_AVAILABLE:
-            required_files.extend(
-                [
-                    "faster_rcnn_R_101_FPN_3x.pth",
-                    "faster_rcnn_X_101_32x8d_FPN_3x.pth",
-                    "cascade_mask_rcnn_R_50_FPN_3x.pth",
-                    "cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.pth",
-                ]
-            )
+        # Инициализация загрузчиков
+        self.yolo_loader = YOLOLoader()
+        self.detectron_loader = Detectron2Loader(device=self.device) if DETECTRON2_AVAILABLE else None
 
+    def _ensure_models_available(self, required_files):
+        """Проверяет и загружает необходимые файлы моделей"""
         for filename in required_files:
             file_path = os.path.join(self.MODELS_DIR, filename)
             if not os.path.exists(file_path):
                 self._download_file(filename)
 
     def _download_file(self, filename):
         """Скачивает файл с сервера"""
-        # Если это модель Detectron2 и библиотека недоступна - пропускаем
         if filename.endswith((".pth", ".yaml")) and not DETECTRON2_AVAILABLE:
             print(f"Skipping {filename} (Detectron2 not available)")
             return False
@@ -138,3 +72,32 @@ def _download_file(self, filename):
             if os.path.exists(save_path):
                 os.remove(save_path)
             return False
+
+    def get_model(self, model_name: str):
+        """Возвращает модель по имени"""
+        if model_name in self.yolo_loader.MODEL_MAPPING:
+            return self.yolo_loader.get_model(model_name)
+        elif DETECTRON2_AVAILABLE and model_name in self.detectron_loader.MODEL_CONFIGS:
+            return self.detectron_loader.get_config(model_name)
+        raise ValueError(f"Unknown model: {model_name}")
+
+    def get_predictor(self, model_name: str):
+        """Для Detectron возвращает готовый predictor"""
+        if DETECTRON2_AVAILABLE and model_name in self.detectron_loader.MODEL_CONFIGS:
+            cfg = self.detectron_loader.get_config(model_name)
+            return DefaultPredictor(cfg)
+        return None
+
+    def get_model_path(self, model_name: str) -> str:
+        """Возвращает путь к модели по её имени"""
+        if model_name in self.yolo_loader.MODEL_MAPPING:
+            return self.yolo_loader.get_model_path(model_name)
+        elif DETECTRON2_AVAILABLE and model_name in self.detectron_loader.MODEL_CONFIGS:
+            return self.detectron_loader.get_model_path(model_name)
+        raise ValueError(f"Model {model_name} not found")
+
+    def get_config_path(self, model_name: str) -> str:
+        """Возвращает путь к конфигу (только для Detectron)"""
+        if DETECTRON2_AVAILABLE and model_name in self.detectron_loader.MODEL_CONFIGS:
+            return self.detectron_loader.get_config_path(model_name)
+        raise ValueError(f"Config for {model_name} not available")