version1

.env.example

@@ -0,0 +1,18 @@
+# API Settings
+API_V1_PREFIX=/api/v1
+PROJECT_NAME=Wagon Classification API
+VERSION=1.0.0
+
+# Model Settings
+MODEL_PATH=models/best_model.pth
+CLASS_NAMES=pered,zad,none
+
+# Security
+MAX_UPLOAD_SIZE=10485760
+ALLOWED_EXTENSIONS=.jpg,.jpeg,.png,.bmp
+
+# CORS
+ALLOWED_ORIGINS=http://localhost,http://localhost:8000,http://127.0.0.1:8000
+
+# Logging
+LOG_LEVEL=INFO

.gitignore

@@ -205,3 +205,9 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+
+uploads/
+wagon_classification/
+wagon_data/
+

Dockerfile

@@ -0,0 +1,31 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+# Устанавливаем системные зависимости
+RUN apt-get update && apt-get install -y \
+    gcc \
+    g++ \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    libsm6 \
+    libxext6 \
+    libxrender-dev \
+    libgomp1 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Копируем зависимости
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Копируем приложение
+COPY . .
+
+# Создаем необходимые папки
+RUN mkdir -p models uploads
+
+# Открываем порт
+EXPOSE 8000
+
+# Запускаем приложение
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]

app/_init__.py

@@ -0,0 +1,3 @@
+"""Wagon Classification API - приложение для классификации вагонов"""
+
+__version__ = "1.0.0"

app/api/__init__.py

@@ -0,0 +1,5 @@
+"""API маршруты приложения"""
+
+from app.api.routes import router
+
+__all__ = ['router']

app/api/dependencies.py

@@ -0,0 +1,23 @@
+"""
+Зависимости для API маршрутов
+"""
+
+from fastapi import Request, HTTPException, status
+from app.models.wagon_model import get_classifier
+
+
+async def verify_model_loaded():
+    """Проверка, что модель загружена"""
+    try:
+        classifier = get_classifier()
+        if classifier.model is None:
+            raise HTTPException(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+                detail="Модель не загружена"
+            )
+        return classifier
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail=f"Ошибка загрузки модели: {str(e)}"
+        )

app/api/routes.py

@@ -0,0 +1,182 @@
+"""
+API эндпоинты для классификации вагонов
+"""
+
+import os
+import uuid
+import logging
+from typing import List
+from fastapi import APIRouter, File, UploadFile, HTTPException, status
+from fastapi.responses import JSONResponse
+from PIL import Image
+import io
+
+from app.models.schemas import PredictionResponse, ErrorResponse, HealthResponse, BatchPredictionResponse
+from app.models.wagon_model import get_classifier
+from app.config import settings
+from app.utils.image_utils import validate_image_file, process_image
+
+logger = logging.getLogger(__name__)
+router = APIRouter()
+
+
+@router.get(
+    "/health",
+    response_model=HealthResponse,
+    tags=["System"],
+    summary="Проверка здоровья сервиса"
+)
+async def health_check():
+    """
+    Проверка работоспособности API и наличия модели
+    """
+    try:
+        classifier = get_classifier()
+        return HealthResponse(
+            status="healthy",
+            model_loaded=True,
+            device=classifier.device,
+            version=settings.VERSION
+        )
+    except Exception as e:
+        logger.error(f"Health check failed: {e}")
+        return HealthResponse(
+            status="unhealthy",
+            model_loaded=False,
+            device="unknown",
+            version=settings.VERSION
+        )
+
+
+@router.post(
+    "/predict",
+    response_model=PredictionResponse,
+    tags=["Prediction"],
+    summary="Классификация одного изображения",
+    responses={
+        400: {"model": ErrorResponse, "description": "Ошибка валидации"},
+        413: {"model": ErrorResponse, "description": "Файл слишком большой"},
+        500: {"model": ErrorResponse, "description": "Внутренняя ошибка сервера"}
+    }
+)
+async def predict_image(
+    file: UploadFile = File(..., description="Изображение вагона")
+):
+    """
+    Классифицирует изображение вагона
+    
+    Определяет:
+    - **pered** - передняя часть вагона
+    - **zad** - задняя часть вагона  
+    - **none** - вагон не обнаружен
+    
+    Возвращает предсказанный класс и уверенность модели.
+    """
+    try:
+        # Валидация файла
+        validate_image_file(file, settings)
+        
+        # Загружаем изображение
+        image = process_image(file)
+        
+        # Получаем модель и делаем предсказание
+        classifier = get_classifier()
+        predicted_class, confidence, probabilities = classifier.predict(image)
+        
+        # Формируем ответ
+        response_data = {
+            "class": predicted_class,
+            "class_name": classifier.class_names_ru.get(predicted_class, predicted_class),
+            "confidence": confidence,
+            "probabilities": probabilities
+        }
+        
+        return PredictionResponse(
+            status="success",
+            data=response_data,
+            request_id=str(uuid.uuid4())
+        )
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Непредвиденная ошибка: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "code": "INTERNAL_ERROR",
+                "message": "Внутренняя ошибка сервера"
+            }
+        )
+
+
+@router.post(
+    "/predict-batch",
+    tags=["Prediction"],
+    summary="Пакетная классификация изображений"
+)
+async def predict_batch(
+    files: List[UploadFile] = File(..., description="Список изображений")
+):
+    """
+    Классифицирует несколько изображений одновременно
+    
+    Максимальное количество файлов не ограничено, но каждый файл
+    должен соответствовать требованиям по размеру и формату.
+    """
+    try:
+        classifier = get_classifier()
+        results = []
+        
+        for file in files:
+            try:
+                # Валидация
+                validate_image_file(file, settings)
+                image = process_image(file)
+                
+                # Предсказание
+                predicted_class, confidence, probabilities = classifier.predict(image)
+                
+                results.append({
+                    "filename": file.filename,
+                    "success": True,
+                    "result": {
+                        "class": predicted_class,
+                        "class_name": classifier.class_names_ru.get(predicted_class, predicted_class),
+                        "confidence": confidence,
+                        "probabilities": probabilities
+                    }
+                })
+                
+            except HTTPException as e:
+                results.append({
+                    "filename": file.filename,
+                    "success": False,
+                    "error": e.detail.get("message", str(e.detail))
+                })
+            except Exception as e:
+                results.append({
+                    "filename": file.filename,
+                    "success": False,
+                    "error": str(e)
+                })
+        
+        return JSONResponse(
+            status_code=status.HTTP_200_OK,
+            content={
+                "status": "success",
+                "results": results,
+                "total": len(results),
+                "successful": sum(1 for r in results if r["success"])
+            }
+        )
+        
+    except Exception as e:
+        logger.error(f"Ошибка пакетной обработки: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "code": "BATCH_ERROR",
+                "message": "Ошибка пакетной обработки"
+            }
+        )

app/config.py

@@ -0,0 +1,56 @@
+"""
+Конфигурация приложения
+Загружает настройки из переменных окружения
+"""
+
+import os
+from pathlib import Path
+from typing import List
+from pydantic_settings import BaseSettings
+from dotenv import load_dotenv
+
+# Загружаем переменные окружения
+load_dotenv()
+
+# Базовые пути
+BASE_DIR = Path(__file__).resolve().parent.parent
+MODEL_DIR = BASE_DIR / "models"
+UPLOAD_DIR = BASE_DIR / "uploads"
+
+# Создаем необходимые папки
+UPLOAD_DIR.mkdir(exist_ok=True)
+MODEL_DIR.mkdir(exist_ok=True)
+
+
+class Settings(BaseSettings):
+    """Настройки приложения"""
+    
+    # API настройки
+    API_V1_PREFIX: str = "/api/v1"
+    PROJECT_NAME: str = "Wagon Classification API"
+    VERSION: str = "1.0.0"
+    
+    # Модель
+    MODEL_PATH: str = str(MODEL_DIR / "best_model.pth")
+    CLASS_NAMES: List[str] = ["pered", "zad", "none"]
+    
+    # Безопасность
+    MAX_UPLOAD_SIZE: int = 10 * 1024 * 1024  # 10 MB
+    ALLOWED_EXTENSIONS: set = {".jpg", ".jpeg", ".png", ".bmp"}
+    
+    # CORS
+    ALLOWED_ORIGINS: List[str] = [
+        "http://localhost",
+        "http://localhost:8000",
+        "http://127.0.0.1:8000"
+    ]
+    
+    # Логирование
+    LOG_LEVEL: str = "INFO"
+    
+    class Config:
+        env_file = ".env"
+        case_sensitive = True
+
+
+settings = Settings()

app/main.py

@@ -0,0 +1,108 @@
+"""
+Основной файл приложения FastAPI
+"""
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse
+import logging
+import os
+from pathlib import Path
+
+from app.api.routes import router
+from app.config import settings
+from app.utils.logger import setup_logging
+
+# Настройка логирования
+setup_logging(settings.LOG_LEVEL)
+logger = logging.getLogger(__name__)
+
+# Создаем приложение
+app = FastAPI(
+    title=settings.PROJECT_NAME,
+    version=settings.VERSION,
+    description="API для классификации вагонов по изображениям\n\n"
+                "Определяет переднюю и заднюю часть вагона на фотографии.",
+    docs_url="/docs",
+    redoc_url="/redoc",
+    openapi_tags=[
+        {
+            "name": "System",
+            "description": "Системные эндпоинты (health check)"
+        },
+        {
+            "name": "Prediction",
+            "description": "Эндпоинты для классификации изображений"
+        }
+    ]
+)
+
+# Настройка CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=settings.ALLOWED_ORIGINS,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Подключаем API роуты
+app.include_router(router, prefix=settings.API_V1_PREFIX)
+
+# Статические файлы (веб-интерфейс)
+static_dir = Path(__file__).parent / "static"
+if static_dir.exists():
+    app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")
+
+
+@app.get("/")
+async def root():
+    """Главная страница"""
+    static_index = static_dir / "index.html"
+    if static_index.exists():
+        return FileResponse(str(static_index))
+    return {
+        "message": settings.PROJECT_NAME,
+        "version": settings.VERSION,
+        "docs": "/docs",
+        "health": f"{settings.API_V1_PREFIX}/health"
+    }
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Загрузка модели при старте"""
+    logger.info("=" * 50)
+    logger.info(f"Запуск {settings.PROJECT_NAME} v{settings.VERSION}")
+    logger.info("=" * 50)
+    
+    # Проверяем существование модели
+    if not os.path.exists(settings.MODEL_PATH):
+        logger.warning(f"⚠️ Модель не найдена: {settings.MODEL_PATH}")
+        logger.info("Пожалуйста, обучите модель командой: python train_model.py")
+    else:
+        try:
+            # Предварительная загрузка модели
+            from app.models.wagon_model import get_classifier
+            classifier = get_classifier()
+            logger.info(f"✅ Модель загружена на устройство: {classifier.device}")
+            logger.info(f"📋 Доступные классы: {classifier.class_names}")
+        except Exception as e:
+            logger.error(f"❌ Ошибка при загрузке модели: {e}")
+
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    """Очистка при завершении"""
+    logger.info("Остановка API сервиса")
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(
+        "app.main:app",
+        host="0.0.0.0",
+        port=8000,
+        reload=True
+    )

app/models/__init__.py

@@ -0,0 +1,12 @@
+"""Модели данных и ML модель"""
+
+from app.models.wagon_model import WagonClassifier, get_classifier
+from app.models.schemas import PredictionResponse, ErrorResponse, HealthResponse
+
+__all__ = [
+    'WagonClassifier',
+    'get_classifier',
+    'PredictionResponse',
+    'ErrorResponse',
+    'HealthResponse'
+]

app/models/schemas.py

@@ -0,0 +1,70 @@
+"""
+Pydantic схемы для валидации данных API
+"""
+
+from pydantic import BaseModel, Field
+from typing import Dict, Optional, List
+from datetime import datetime
+
+
+class PredictionResponse(BaseModel):
+    """Ответ API с предсказанием"""
+    status: str = Field(..., example="success")
+    data: Dict = Field(..., description="Результат классификации")
+    timestamp: datetime = Field(default_factory=datetime.now)
+    request_id: Optional[str] = Field(None, description="Уникальный идентификатор запроса")
+    
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "status": "success",
+                "data": {
+                    "class": "pered",
+                    "class_name": "передняя часть вагона",
+                    "confidence": 0.95,
+                    "probabilities": {
+                        "pered": 0.95,
+                        "zad": 0.03,
+                        "none": 0.02
+                    }
+                },
+                "timestamp": "2024-01-15T10:30:00",
+                "request_id": "123e4567-e89b-12d3-a456-426614174000"
+            }
+        }
+
+
+class ErrorResponse(BaseModel):
+    """Ответ при ошибке"""
+    status: str = Field(..., example="error")
+    error: Dict = Field(..., description="Детали ошибки")
+    timestamp: datetime = Field(default_factory=datetime.now)
+    
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "status": "error",
+                "error": {
+                    "code": "INVALID_IMAGE",
+                    "message": "Файл не является корректным изображением",
+                    "details": "Поддерживаются форматы: jpg, jpeg, png"
+                },
+                "timestamp": "2024-01-15T10:30:00"
+            }
+        }
+
+
+class HealthResponse(BaseModel):
+    """Проверка здоровья сервиса"""
+    status: str = Field(..., description="Статус сервиса")
+    model_loaded: bool = Field(..., description="Загружена ли модель")
+    device: str = Field(..., description="Устройство выполнения")
+    version: str = Field(..., description="Версия API")
+
+
+class BatchPredictionResponse(BaseModel):
+    """Ответ для пакетной классификации"""
+    status: str = Field(..., example="success")
+    results: List[Dict] = Field(..., description="Результаты для каждого файла")
+    total: int = Field(..., description="Всего файлов")
+    successful: int = Field(..., description="Успешно обработано")

app/models/wagon_model.py

@@ -0,0 +1,199 @@
+"""
+Обертка для модели машинного обучения
+Загружает обученную модель и выполняет предсказания
+"""
+
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
+import os
+import logging
+from typing import Dict, Tuple, List
+
+logger = logging.getLogger(__name__)
+
+
+class WagonClassifier:
+    """
+    Классификатор вагонов
+    Загружает обученную модель и выполняет инференс
+    """
+    
+    def __init__(self, model_path: str, class_names: List[str], device: str = None):
+        """
+        Инициализация классификатора
+        
+        Args:
+            model_path: Путь к файлу с весами модели (.pth)
+            class_names: Список названий классов
+            device: Устройство для выполнения (cuda/cpu)
+        """
+        self.model_path = model_path
+        self.class_names = class_names
+        self.num_classes = len(class_names)
+        
+        # Определяем устройство
+        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
+        logger.info(f"Используется устройство: {self.device}")
+        
+        # Загружаем модель
+        self.model = self._load_model()
+        
+        # Трансформации для изображений
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225]
+            )
+        ])
+        
+        # Русские названия классов для вывода
+        self.class_names_ru = {
+            'pered': 'передняя часть вагона',
+            'zad': 'задняя часть вагона',
+            'none': 'вагон не обнаружен'
+        }
+        
+        logger.info(f"Модель загружена. Классы: {self.class_names}")
+    
+    def _load_model(self) -> nn.Module:
+        """
+        Загрузка модели из файла
+        
+        Returns:
+            Загруженная модель в режиме evaluation
+        """
+        # Создаем архитектуру модели (должна совпадать с train_model.py)
+        model = models.efficientnet_b2(weights=None)
+        in_features = model.classifier[1].in_features
+        model.classifier = nn.Sequential(
+            nn.Dropout(p=0.3),
+            nn.Linear(in_features, self.num_classes)
+        )
+        
+        # Проверяем существование файла
+        if not os.path.exists(self.model_path):
+            raise FileNotFoundError(f"Модель не найдена: {self.model_path}")
+        
+        # Загружаем веса
+        checkpoint = torch.load(self.model_path, map_location=self.device)
+        
+        # Поддерживаем разные форматы сохранения
+        if 'model_state_dict' in checkpoint:
+            model.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            model.load_state_dict(checkpoint)
+        
+        # Перемещаем на устройство и переводим в режим оценки
+        model = model.to(self.device)
+        model.eval()
+        
+        return model
+    
+    def _preprocess_image(self, image: Image.Image) -> torch.Tensor:
+        """
+        Предобработка изображения перед подачей в модель
+        
+        Args:
+            image: PIL Image
+            
+        Returns:
+            Тензор, готовый для инференса
+        """
+        # Конвертируем в RGB если нужно
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        # Применяем трансформации
+        input_tensor = self.transform(image)
+        input_tensor = input_tensor.unsqueeze(0)  # Добавляем batch dimension
+        input_tensor = input_tensor.to(self.device)
+        
+        return input_tensor
+    
+    def predict(self, image: Image.Image) -> Tuple[str, float, Dict[str, float]]:
+        """
+        Предсказание для одного изображения
+        
+        Args:
+            image: PIL Image
+            
+        Returns:
+            Tuple: (предсказанный_класс, уверенность, словарь_вероятностей)
+        """
+        try:
+            # Предобработка
+            input_tensor = self._preprocess_image(image)
+            
+            # Инференс
+            with torch.no_grad():
+                outputs = self.model(input_tensor)
+                probabilities = torch.nn.functional.softmax(outputs, dim=1)
+                
+                # Получаем предсказание
+                predicted_idx = torch.argmax(probabilities, dim=1).item()
+                confidence = probabilities[0][predicted_idx].item()
+                predicted_class = self.class_names[predicted_idx]
+                
+                # Все вероятности
+                all_probs = {
+                    class_name: probabilities[0][i].item()
+                    for i, class_name in enumerate(self.class_names)
+                }
+            
+            logger.info(f"Предсказание: {predicted_class} с уверенностью {confidence:.2%}")
+            
+            return predicted_class, confidence, all_probs
+            
+        except Exception as e:
+            logger.error(f"Ошибка при предсказании: {e}")
+            raise
+    
+    def predict_batch(self, images: List[Image.Image]) -> List[Dict]:
+        """
+        Пакетное предсказание для нескольких изображений
+        
+        Args:
+            images: Список PIL Image
+            
+        Returns:
+            Список результатов для каждого изображения
+        """
+        results = []
+        for image in images:
+            pred_class, confidence, probs = self.predict(image)
+            results.append({
+                'class': pred_class,
+                'class_name': self.class_names_ru.get(pred_class, pred_class),
+                'confidence': confidence,
+                'probabilities': probs
+            })
+        return results
+
+
+# Глобальный экземпляр модели (синглтон)
+_classifier_instance = None
+
+
+def get_classifier() -> WagonClassifier:
+    """
+    Получить экземпляр классификатора (синглтон)
+    Модель загружается только один раз при первом вызове
+    
+    Returns:
+        Экземпляр WagonClassifier
+    """
+    global _classifier_instance
+    
+    if _classifier_instance is None:
+        from app.config import settings
+        
+        _classifier_instance = WagonClassifier(
+            model_path=settings.MODEL_PATH,
+            class_names=settings.CLASS_NAMES
+        )
+    
+    return _classifier_instance

app/services/prediction_service.py

@@ -0,0 +1,67 @@
+"""
+Сервис для работы с предсказаниями
+Содержит бизнес-логику обработки изображений
+"""
+
+import logging
+from typing import Dict, Any, List
+from PIL import Image
+
+from app.models.wagon_model import WagonClassifier
+
+logger = logging.getLogger(__name__)
+
+
+class PredictionService:
+    """Сервис для выполнения предсказаний"""
+    
+    def __init__(self, classifier: WagonClassifier):
+        self.classifier = classifier
+    
+    def predict_single(self, image: Image.Image) -> Dict[str, Any]:
+        """
+        Предсказание для одного изображения
+        
+        Args:
+            image: PIL Image
+            
+        Returns:
+            Словарь с результатами предсказания
+        """
+        predicted_class, confidence, probabilities = self.classifier.predict(image)
+        
+        return {
+            "class": predicted_class,
+            "class_name": self.classifier.class_names_ru.get(predicted_class, predicted_class),
+            "confidence": confidence,
+            "probabilities": probabilities
+        }
+    
+    def predict_batch(self, images: List[Image.Image]) -> List[Dict[str, Any]]:
+        """
+        Предсказание для нескольких изображений
+        
+        Args:
+            images: Список PIL Image
+            
+        Returns:
+            Список результатов
+        """
+        results = []
+        for image in images:
+            try:
+                result = self.predict_single(image)
+                results.append(result)
+            except Exception as e:
+                logger.error(f"Ошибка при предсказании: {e}")
+                results.append({"error": str(e)})
+        
+        return results
+    
+    def get_model_info(self) -> Dict[str, Any]:
+        """Получить информацию о модели"""
+        return {
+            "device": self.classifier.device,
+            "classes": self.classifier.class_names,
+            "num_classes": self.classifier.num_classes
+        }

app/static/index.html

@@ -0,0 +1,54 @@
+<!DOCTYPE html>
+<html lang="ru">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0, user-scalable=yes">
+    <title>Классификатор вагонов - WagonDetector</title>
+    <link rel="stylesheet" href="/static/style.css">
+</head>
+<body>
+    <div class="container">
+        <header>
+            <h1>🚂 Классификатор вагонов</h1>
+            <p>Определение передней и задней части вагона по фотографии</p>
+        </header>
+
+        <div class="upload-area" id="uploadArea">
+            <div class="upload-content">
+                <div class="upload-icon">📸</div>
+                <h3>Загрузите изображение вагона</h3>
+                <p>Перетащите файл сюда или нажмите для выбора</p>
+                <p class="file-info">Поддерживаются: JPG, JPEG, PNG (до 10 MB)</p>
+                <input type="file" id="fileInput" accept="image/jpeg,image/jpg,image/png" hidden>
+                <button class="btn btn-primary" id="selectFileBtn">Выбрать файл</button>
+            </div>
+        </div>
+
+        <div class="preview-area" id="previewArea" style="display: none;">
+            <div class="preview-image">
+                <img id="previewImg" alt="Предпросмотр">
+                <button class="btn-remove" id="removeImageBtn">✕</button>
+            </div>
+            <button class="btn btn-success" id="predictBtn">🔍 Распознать вагон</button>
+        </div>
+
+        <div class="results-area" id="resultsArea" style="display: none;">
+            <h3>Результат классификации</h3>
+            <div class="result-card">
+                <div class="result-class" id="resultClass"></div>
+                <div class="result-confidence" id="resultConfidence"></div>
+                <div class="probabilities" id="probabilities"></div>
+            </div>
+        </div>
+
+        <div class="loading" id="loading" style="display: none;">
+            <div class="spinner"></div>
+            <p>Обработка изображения...</p>
+        </div>
+
+        <div class="error-message" id="errorMessage" style="display: none;"></div>
+    </div>
+
+    <script src="/static/script.js"></script>
+</body>
+</html>

app/static/script.js

@@ -0,0 +1,201 @@
+// DOM элементы
+const uploadArea = document.getElementById('uploadArea');
+const fileInput = document.getElementById('fileInput');
+const selectFileBtn = document.getElementById('selectFileBtn');
+const previewArea = document.getElementById('previewArea');
+const previewImg = document.getElementById('previewImg');
+const removeImageBtn = document.getElementById('removeImageBtn');
+const predictBtn = document.getElementById('predictBtn');
+const resultsArea = document.getElementById('resultsArea');
+const loading = document.getElementById('loading');
+const errorMessage = document.getElementById('errorMessage');
+
+let currentFile = null;
+
+// API базовый URL
+const API_URL = window.location.origin + '/api/v1';
+
+// Обработчики событий
+selectFileBtn.addEventListener('click', () => fileInput.click());
+removeImageBtn.addEventListener('click', clearImage);
+predictBtn.addEventListener('click', predictImage);
+fileInput.addEventListener('change', handleFileSelect);
+
+// Drag & Drop
+uploadArea.addEventListener('dragover', (e) => {
+    e.preventDefault();
+    uploadArea.classList.add('drag-over');
+});
+
+uploadArea.addEventListener('dragleave', () => {
+    uploadArea.classList.remove('drag-over');
+});
+
+uploadArea.addEventListener('drop', (e) => {
+    e.preventDefault();
+    uploadArea.classList.remove('drag-over');
+    const file = e.dataTransfer.files[0];
+    if (file && file.type.startsWith('image/')) {
+        handleFile(file);
+    } else {
+        showError('Пожалуйста, загрузите изображение');
+    }
+});
+
+uploadArea.addEventListener('click', () => fileInput.click());
+
+function handleFileSelect(e) {
+    const file = e.target.files[0];
+    if (file) {
+        handleFile(file);
+    }
+}
+
+function handleFile(file) {
+    // Проверка размера (10 MB)
+    if (file.size > 10 * 1024 * 1024) {
+        showError('Файл слишком большой. Максимум 10 MB');
+        return;
+    }
+    
+    currentFile = file;
+    
+    // Предпросмотр
+    const reader = new FileReader();
+    reader.onload = (e) => {
+        previewImg.src = e.target.result;
+        uploadArea.style.display = 'none';
+        previewArea.style.display = 'block';
+        resultsArea.style.display = 'none';
+        hideError();
+    };
+    reader.readAsDataURL(file);
+}
+
+function clearImage() {
+    currentFile = null;
+    fileInput.value = '';
+    previewArea.style.display = 'none';
+    uploadArea.style.display = 'block';
+    resultsArea.style.display = 'none';
+    hideError();
+}
+
+async function predictImage() {
+    if (!currentFile) {
+        showError('Сначала выберите изображение');
+        return;
+    }
+    
+    // Показываем загрузку
+    loading.style.display = 'block';
+    resultsArea.style.display = 'none';
+    hideError();
+    
+    // Создаем FormData
+    const formData = new FormData();
+    formData.append('file', currentFile);
+    
+    try {
+        // Отправляем запрос
+        const response = await fetch(`${API_URL}/predict`, {
+            method: 'POST',
+            body: formData
+        });
+        
+        const data = await response.json();
+        
+        if (!response.ok) {
+            throw new Error(data.error?.message || 'Ошибка при обработке');
+        }
+        
+        // Отображаем результаты
+        displayResults(data.data);
+        
+    } catch (error) {
+        console.error('Error:', error);
+        showError(error.message || 'Ошибка при отправке запроса');
+    } finally {
+        loading.style.display = 'none';
+    }
+}
+
+function displayResults(data) {
+    const resultClass = document.getElementById('resultClass');
+    const resultConfidence = document.getElementById('resultConfidence');
+    const probabilitiesDiv = document.getElementById('probabilities');
+    
+    // Определяем эмодзи для класса
+    let emoji = '';
+    if (data.class === 'pered') emoji = '🚂 Передняя часть';
+    else if (data.class === 'zad') emoji = '🚂 Задняя часть';
+    else emoji = '⭕ Вагон не обнаружен';
+    
+    // Отображаем основной результат
+    resultClass.innerHTML = `${emoji}<br>${data.class_name}`;
+    resultConfidence.innerHTML = `Уверенность: <strong>${(data.confidence * 100).toFixed(1)}%</strong>`;
+    
+    // Отображаем все вероятности
+    probabilitiesDiv.innerHTML = '<h4>Распределение вероятностей:</h4>';
+    
+    const classNames = {
+        'pered': 'Передняя часть',
+        'zad': 'Задняя часть',
+        'none': 'Вагон не обнаружен'
+    };
+    
+    for (const [cls, prob] of Object.entries(data.probabilities)) {
+        const percent = (prob * 100).toFixed(1);
+        const isPredicted = cls === data.class;
+        
+        const probItem = document.createElement('div');
+        probItem.className = 'prob-item';
+        probItem.innerHTML = `
+            <div class="prob-label">${classNames[cls] || cls}</div>
+            <div class="prob-bar">
+                <div class="prob-fill" style="width: ${percent}%; background: ${isPredicted ? 'linear-gradient(90deg, #667eea, #764ba2)' : '#cbd5e0'}">
+                    ${percent}%
+                </div>
+            </div>
+        `;
+        probabilitiesDiv.appendChild(probItem);
+    }
+    
+    // Показываем результаты
+    resultsArea.style.display = 'block';
+    
+    // Прокручиваем к результатам
+    resultsArea.scrollIntoView({ behavior: 'smooth' });
+}
+
+function showError(message) {
+    errorMessage.textContent = message;
+    errorMessage.style.display = 'block';
+    setTimeout(() => {
+        errorMessage.style.display = 'none';
+    }, 5000);
+}
+
+function hideError() {
+    errorMessage.style.display = 'none';
+}
+
+// Проверка здоровья API при загрузке
+async function checkHealth() {
+    try {
+        const response = await fetch(`${API_URL}/health`);
+        const data = await response.json();
+        if (data.status === 'healthy') {
+            console.log('✅ API готов к работе');
+        } else {
+            console.warn('⚠️ API не здоров:', data);
+            showError('Сервис временно недоступен');
+        }
+    } catch (error) {
+        console.error('❌ API недоступен:', error);
+        showError('Не удалось подключиться к серверу');
+    }
+}
+
+// Запускаем проверку при загрузке
+checkHealth();

app/static/style.css

@@ -0,0 +1,262 @@
+* {
+    margin: 0;
+    padding: 0;
+    box-sizing: border-box;
+}
+
+body {
+    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, sans-serif;
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    min-height: 100vh;
+    padding: 20px;
+}
+
+.container {
+    max-width: 800px;
+    margin: 0 auto;
+    background: white;
+    border-radius: 20px;
+    box-shadow: 0 20px 60px rgba(0,0,0,0.3);
+    overflow: hidden;
+    padding: 30px;
+}
+
+header {
+    text-align: center;
+    margin-bottom: 30px;
+}
+
+header h1 {
+    color: #333;
+    font-size: 28px;
+    margin-bottom: 10px;
+}
+
+header p {
+    color: #666;
+    font-size: 14px;
+}
+
+.upload-area {
+    border: 2px dashed #ddd;
+    border-radius: 10px;
+    padding: 40px;
+    text-align: center;
+    transition: all 0.3s ease;
+    cursor: pointer;
+}
+
+.upload-area:hover {
+    border-color: #667eea;
+    background: #f8f9ff;
+}
+
+.upload-area.drag-over {
+    border-color: #667eea;
+    background: #f0f2ff;
+}
+
+.upload-icon {
+    font-size: 48px;
+    margin-bottom: 15px;
+}
+
+.upload-content h3 {
+    color: #333;
+    margin-bottom: 10px;
+}
+
+.upload-content p {
+    color: #666;
+    margin-bottom: 5px;
+}
+
+.file-info {
+    font-size: 12px;
+    color: #999;
+}
+
+.btn {
+    padding: 10px 20px;
+    border: none;
+    border-radius: 5px;
+    font-size: 14px;
+    cursor: pointer;
+    transition: all 0.3s ease;
+    margin-top: 15px;
+}
+
+.btn-primary {
+    background: #667eea;
+    color: white;
+}
+
+.btn-primary:hover {
+    background: #5a67d8;
+    transform: translateY(-2px);
+}
+
+.btn-success {
+    background: #48bb78;
+    color: white;
+    font-size: 16px;
+    padding: 12px 30px;
+}
+
+.btn-success:hover {
+    background: #38a169;
+    transform: translateY(-2px);
+}
+
+.preview-area {
+    text-align: center;
+    margin-top: 20px;
+}
+
+.preview-image {
+    position: relative;
+    display: inline-block;
+    margin-bottom: 20px;
+}
+
+.preview-image img {
+    max-width: 100%;
+    max-height: 400px;
+    border-radius: 10px;
+    box-shadow: 0 5px 15px rgba(0,0,0,0.2);
+}
+
+.btn-remove {
+    position: absolute;
+    top: -10px;
+    right: -10px;
+    width: 30px;
+    height: 30px;
+    border-radius: 50%;
+    background: #f56565;
+    color: white;
+    border: none;
+    cursor: pointer;
+    font-size: 18px;
+    transition: all 0.3s ease;
+}
+
+.btn-remove:hover {
+    background: #e53e3e;
+    transform: scale(1.1);
+}
+
+.results-area {
+    margin-top: 30px;
+}
+
+.results-area h3 {
+    text-align: center;
+    color: #333;
+    margin-bottom: 20px;
+}
+
+.result-card {
+    background: #f7fafc;
+    border-radius: 10px;
+    padding: 20px;
+    text-align: center;
+}
+
+.result-class {
+    font-size: 24px;
+    font-weight: bold;
+    color: #667eea;
+    margin-bottom: 10px;
+}
+
+.result-confidence {
+    font-size: 18px;
+    color: #48bb78;
+    margin-bottom: 20px;
+}
+
+.probabilities {
+    text-align: left;
+    margin-top: 20px;
+}
+
+.probabilities h4 {
+    margin-bottom: 10px;
+    color: #4a5568;
+    font-size: 14px;
+}
+
+.prob-item {
+    margin-bottom: 10px;
+}
+
+.prob-label {
+    font-size: 14px;
+    color: #666;
+    margin-bottom: 5px;
+}
+
+.prob-bar {
+    background: #e2e8f0;
+    height: 30px;
+    border-radius: 5px;
+    overflow: hidden;
+    position: relative;
+}
+
+.prob-fill {
+    background: linear-gradient(90deg, #667eea, #764ba2);
+    height: 100%;
+    display: flex;
+    align-items: center;
+    justify-content: flex-end;
+    padding-right: 10px;
+    color: white;
+    font-size: 12px;
+    font-weight: bold;
+    transition: width 0.5s ease;
+}
+
+.loading {
+    text-align: center;
+    padding: 40px;
+}
+
+.spinner {
+    border: 4px solid #f3f3f3;
+    border-top: 4px solid #667eea;
+    border-radius: 50%;
+    width: 50px;
+    height: 50px;
+    animation: spin 1s linear infinite;
+    margin: 0 auto 20px;
+}
+
+@keyframes spin {
+    0% { transform: rotate(0deg); }
+    100% { transform: rotate(360deg); }
+}
+
+.error-message {
+    background: #fed7d7;
+    color: #c53030;
+    padding: 15px;
+    border-radius: 10px;
+    margin-top: 20px;
+    text-align: center;
+}
+
+@media (max-width: 768px) {
+    .container {
+        padding: 20px;
+    }
+    
+    .upload-area {
+        padding: 20px;
+    }
+    
+    .result-class {
+        font-size: 20px;
+    }
+}

app/utils/__init__.py

@@ -0,0 +1,6 @@
+"""Вспомогательные утилиты"""
+
+from app.utils.image_utils import validate_image_file, process_image
+from app.utils.logger import setup_logging
+
+__all__ = ['validate_image_file', 'process_image', 'setup_logging']

app/utils/image_utils.py

@@ -0,0 +1,95 @@
+"""
+Утилиты для работы с изображениями
+"""
+
+import os
+import io
+import logging
+from fastapi import UploadFile, HTTPException, status
+from PIL import Image, ImageFile
+
+# Разрешаем загрузку усеченных изображений
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+logger = logging.getLogger(__name__)
+
+
+def validate_image_file(file: UploadFile, settings) -> bool:
+    """
+    Проверка валидности файла изображения
+    
+    Args:
+        file: Загруженный файл
+        settings: Настройки приложения
+        
+    Returns:
+        True если файл валиден
+        
+    Raises:
+        HTTPException при ошибках валидации
+    """
+    # Проверяем расширение
+    ext = os.path.splitext(file.filename)[1].lower()
+    if ext not in settings.ALLOWED_EXTENSIONS:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail={
+                "code": "INVALID_EXTENSION",
+                "message": f"Неподдерживаемый формат. Разрешенные: {', '.join(settings.ALLOWED_EXTENSIONS)}"
+            }
+        )
+    
+    # Проверяем размер
+    file.file.seek(0, 2)
+    size = file.file.tell()
+    file.file.seek(0)
+    
+    if size > settings.MAX_UPLOAD_SIZE:
+        raise HTTPException(
+            status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
+            detail={
+                "code": "FILE_TOO_LARGE",
+                "message": f"Файл слишком большой. Максимум: {settings.MAX_UPLOAD_SIZE // (1024*1024)} MB"
+            }
+        )
+    
+    return True
+
+
+def process_image(file: UploadFile) -> Image.Image:
+    """
+    Загрузка и предобработка изображения
+    
+    Args:
+        file: Загруженный файл
+        
+    Returns:
+        PIL Image объект
+        
+    Raises:
+        HTTPException при ошибках загрузки
+    """
+    try:
+        # Читаем файл
+        contents = file.file.read()
+        
+        # Пытаемся открыть как изображение
+        image = Image.open(io.BytesIO(contents))
+        
+        # Проверяем, что изображение можно прочитать
+        image.verify()
+        
+        # Переоткрываем (после verify нужно заново)
+        image = Image.open(io.BytesIO(contents))
+        
+        return image
+        
+    except Exception as e:
+        logger.error(f"Ошибка загрузки изображения: {e}")
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail={
+                "code": "INVALID_IMAGE",
+                "message": "Файл не является корректным изображением"
+            }
+        )

app/utils/logger.py

@@ -0,0 +1,38 @@
+"""
+Настройка логирования
+"""
+
+import logging
+import sys
+from typing import Optional
+
+
+def setup_logging(level: Optional[str] = None):
+    """
+    Настройка логирования для приложения
+    
+    Args:
+        level: Уровень логирования (DEBUG, INFO, WARNING, ERROR)
+    """
+    log_level = level or logging.INFO
+    
+    # Настройка формата
+    formatter = logging.Formatter(
+        '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S'
+    )
+    
+    # Обработчик для stdout
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setFormatter(formatter)
+    
+    # Настройка корневого логгера
+    root_logger = logging.getLogger()
+    root_logger.setLevel(log_level)
+    root_logger.addHandler(console_handler)
+    
+    # Уменьшаем логи от библиотек
+    logging.getLogger("urllib3").setLevel(logging.WARNING)
+    logging.getLogger("torch").setLevel(logging.WARNING)
+    
+    return root_logger

docker-compose.yml

@@ -0,0 +1,22 @@
+version: '3.8'
+
+services:
+  api:
+    build: .
+    ports:
+      - "8000:8000"
+    volumes:
+      - ./models:/app/models
+      - ./uploads:/app/uploads
+      - ./app:/app/app
+      - ./logs:/app/logs
+    environment:
+      - MODEL_PATH=/app/models/best_model.pth
+      - LOG_LEVEL=INFO
+    restart: unless-stopped
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:8000/api/v1/health"]
+      interval: 30s
+      timeout: 10s
+      retries: 3
+      start_period: 40s

requirements.txt

@@ -0,0 +1,27 @@
+tensorflow
+opencv-python
+matplotlib
+numpy
+scikit-learn
+torch
+torchvision
+seaborn
+tqdm
+Pillow
+pandas
+patool
+unrar
+winrar
+fastapi
+uvicorn[standard]
+python-multipart
+python-dotenv
+pydantic
+pydantic-settings
+aiofiles
+python-jose[cryptography]
+python-magic
+pytest
+pytest-asyncio
+httpx
+

train_model.py

@@ -0,0 +1,1010 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
+from torchvision import models, transforms
+import torch.cuda.amp as amp
+import os
+import shutil
+import numpy as np
+from PIL import Image, ImageFile
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import confusion_matrix, classification_report
+import seaborn as sns
+import warnings
+warnings.filterwarnings('ignore')
+
+# ================================================
+# ВКЛЮЧАЕМ ОБРАБОТКУ УСЕЧЕННЫХ ИЗОБРАЖЕНИЙ
+# ================================================
+ImageFile.LOAD_TRUNCATED_IMAGES = True  # Разрешаем загрузку усеченных изображений
+
+# ================================================
+# КОНФИГУРАЦИЯ (С ПРАВИЛЬНЫМИ НАЗВАНИЯМИ КЛАССОВ)
+# ================================================
+class Config:
+    # Пути (изменены для Windows)
+    BASE_DIR = os.path.join(os.getcwd(), 'wagon_classification')
+    DATA_DIR = os.path.join(os.getcwd(), 'wagon_classification', 'data', 'processed')
+    EXTRACTED_DIR = os.path.join(os.getcwd(), 'wagon_data', 'extracted')
+    MODEL_SAVE_PATH = os.path.join(os.getcwd(), 'wagon_classification', 'best_model.pth')
+
+    # Параметры - ИСПРАВЛЕНО: pered вместо prered
+    CLASS_NAMES = ['pered', 'zad', 'none']  # Изменено prered -> pered
+    NUM_CLASSES = 3
+
+    # Гиперпараметры
+    BATCH_SIZE = 32  # Оптимальный размер для T4
+    NUM_EPOCHS = 15
+
+    # Устройство
+    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    @staticmethod
+    def print_info():
+        print(f"\n📊 КОНФИГУРАЦИЯ:")
+        print(f"  • Устройство: {Config.DEVICE}")
+        if Config.DEVICE.type == 'cuda':
+            print(f"  • GPU: {torch.cuda.get_device_name(0)}")
+            print(f"  • Память: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
+        print(f"  • Классы: {Config.CLASS_NAMES}")
+        print(f"  • Batch size: {Config.BATCH_SIZE}")
+        print(f"  • Эпох: {Config.NUM_EPOCHS}")
+
+# ================================================
+# УТИЛИТЫ ДЛЯ РАБОТЫ С ИЗОБРАЖЕНИЯМИ
+# ================================================
+def load_image_safe(image_path, target_size=(224, 224)):
+    """Безопасная загрузка изображения с обработкой ошибок"""
+    try:
+        # Пытаемся открыть изображение
+        image = Image.open(image_path)
+
+        # Проверяем, что изображение валидно
+        image.verify()  # Проверка целостности файла
+
+        # Закрываем и открываем снова (после verify нужно переоткрыть)
+        image = Image.open(image_path)
+
+        # Конвертируем в RGB если нужно
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+
+        # Проверяем размеры
+        if image.size[0] == 0 or image.size[1] == 0:
+            print(f"⚠ Изображение {image_path} имеет нулевые размеры")
+            # Создаем черное изображение
+            image = Image.new('RGB', target_size, color='black')
+
+        return image
+
+    except (IOError, OSError, Image.DecompressionBombError) as e:
+        print(f"⚠ Ошибка загрузки {image_path}: {e}")
+        # Создаем черное изображение в случае ошибки
+        return Image.new('RGB', target_size, color='black')
+
+    except Exception as e:
+        print(f"⚠ Неизвестная ошибка при загрузке {image_path}: {e}")
+        return Image.new('RGB', target_size, color='black')
+
+def repair_image_file(image_path):
+    """Пытается восстановить поврежденный файл изображения"""
+    try:
+        with open(image_path, 'rb') as f:
+            data = f.read()
+
+        # Проверяем, что файл не пустой
+        if len(data) == 0:
+            print(f"❌ Файл {image_path} пустой")
+            return False
+
+        # Пытаемся восстановить как JPEG
+        if image_path.lower().endswith('.jpg') or image_path.lower().endswith('.jpeg'):
+            # Добавляем маркер конца JPEG если нужно
+            if not data.endswith(b'\xff\xd9'):
+                print(f"⚠ Восстанавливаю JPEG файл {image_path}")
+                data += b'\xff\xd9'
+                with open(image_path, 'wb') as f:
+                    f.write(data)
+                return True
+
+        return False
+
+    except Exception as e:
+        print(f"❌ Ошибка при восстановлении {image_path}: {e}")
+        return False
+
+# ================================================
+# ТРАНСФОРМАЦИИ (ПРОСТЫЕ И РАБОЧИЕ)
+# ================================================
+def get_transforms():
+    """Создание простых и рабочих трансформаций"""
+    # Обучающие трансформации
+    train_transform = transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.RandomCrop(224),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.ColorJitter(brightness=0.2, contrast=0.2),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                           std=[0.229, 0.224, 0.225])
+    ])
+
+    # Валидационные трансформации
+    val_transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                           std=[0.229, 0.224, 0.225])
+    ])
+
+    return train_transform, val_transform
+
+# ================================================
+# ПОДГОТОВКА ДАННЫХ (С ПРАВИЛЬНЫМИ ИМЕНАМИ ПАПОК)
+# ================================================
+def prepare_data_simple():
+    """Упрощенная подготовка данных с правильными именами папок"""
+    print("=" * 60)
+    print("📊 ПОДГОТОВКА ДАННЫХ")
+    print("=" * 60)
+
+    # Создаем папки
+    os.makedirs(Config.BASE_DIR, exist_ok=True)
+    os.makedirs(Config.EXTRACTED_DIR, exist_ok=True)
+    os.makedirs(Config.DATA_DIR, exist_ok=True)
+
+    # Для Windows - ручной ввод пути к архиву
+    print("\n📤 ШАГ 1: Укажите путь к архиву vagon1.rar")
+    print("Пример: C:/Users/Username/Downloads/vagon1.rar")
+    archive_path = input("Введите полный путь к архиву: ").strip()
+    
+    # Заменяем прямые слеши на обратные для Windows
+    archive_path = archive_path.replace('/', '\\')
+    
+    if not os.path.exists(archive_path):
+        print(f"\n❌ Файл не найден: {archive_path}")
+        return False
+
+    print(f"\n✅ Найден архив: {os.path.basename(archive_path)}")
+
+    # Распаковываем с использованием patool (кроссплатформенный)
+    try:
+        import patoolib
+        print("📦 Распаковка архива...")
+        patoolib.extract_archive(archive_path, outdir=Config.EXTRACTED_DIR)
+        print("✅ Архив распакован")
+    except ImportError:
+        print("⚠ Установите библиотеку patool: pip install patool")
+        print("Или распакуйте архив вручную в папку:", Config.EXTRACTED_DIR)
+        return False
+    except Exception as e:
+        print(f"⚠ Ошибка при распаковке: {e}")
+        print("Попробуйте распаковать вручную в папку:", Config.EXTRACTED_DIR)
+        return False
+
+    # Проверяем структуру - ИЩЕМ ПРАВИЛЬНЫЕ ИМЕНА ПАПОК
+    print("\n🔍 Проверка данных...")
+
+    # Список возможных имен папок (учитываем опечатки)
+    possible_folders = {
+        'pered': ['pered', 'prered', 'peredn', 'peredniy', 'front', 'перед'],
+        'zad': ['zad', 'zadn', 'zadniy', 'back', 'rear', 'зад'],
+        'none': ['none', 'non', 'empty', 'нет', 'пусто']
+    }
+
+    actual_folders = os.listdir(Config.EXTRACTED_DIR)
+    print(f"Найдены папки в extracted: {actual_folders}")
+
+    # Сопоставляем фактические папки с нашими классами
+    folder_mapping = {}
+    for target_class, possible_names in possible_folders.items():
+        for folder in actual_folders:
+            folder_lower = folder.lower()
+            if folder_lower in possible_names:
+                folder_mapping[target_class] = folder
+                print(f"  ✓ {target_class} → {folder}")
+                break
+
+    # Если не нашли все классы, пытаемся найти по содержимому
+    if len(folder_mapping) < len(Config.CLASS_NAMES):
+        print("\n⚠ Не все классы найдены. Ищем изображения...")
+        for folder in actual_folders:
+            folder_path = os.path.join(Config.EXTRACTED_DIR, folder)
+            if os.path.isdir(folder_path):
+                images = [f for f in os.listdir(folder_path)
+                         if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
+                if images:
+                    print(f"  Папка '{folder}': {len(images)} изображений")
+                    # Пробуем угадать класс
+                    if 'pered' in folder.lower() or 'перед' in folder.lower():
+                        folder_mapping['pered'] = folder
+                    elif 'zad' in folder.lower() or 'зад' in folder.lower():
+                        folder_mapping['zad'] = folder
+                    elif 'none' in folder.lower() or 'нет' in folder.lower():
+                        folder_mapping['none'] = folder
+
+    # Проверяем, что нашли все необходимые классы
+    missing_classes = []
+    for cls in Config.CLASS_NAMES:
+        if cls not in folder_mapping:
+            missing_classes.append(cls)
+
+    if missing_classes:
+        print(f"\n❌ Отсутствуют классы: {missing_classes}")
+        print("Пожалуйста, убедитесь что в архиве есть папки с именами:")
+        print("  - 'pered' (или похожее) - для передней части вагона")
+        print("  - 'zad' (или похожее) - для задней части вагона")
+        print("  - 'none' (или похожее) - для отсутствия вагона")
+        return False
+
+    print("\n✅ Все классы найдены!")
+
+    # Создаем структуру
+    print("\n📁 Создание структуры train/val...")
+    for split in ['train', 'val']:
+        for cls in Config.CLASS_NAMES:
+            os.makedirs(os.path.join(Config.DATA_DIR, split, cls), exist_ok=True)
+
+    # Распределяем данные
+    print("📊 Разделение на train/val (80/20)...")
+    total_images = 0
+
+    for target_class, source_folder in folder_mapping.items():
+        source_dir = os.path.join(Config.EXTRACTED_DIR, source_folder)
+
+        if not os.path.exists(source_dir):
+            print(f"⚠ Папка {source_folder} не найдена, пропускаем")
+            continue
+
+        images = [f for f in os.listdir(source_dir)
+                 if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
+
+        if not images:
+            print(f"⚠ В папке {source_folder} нет изображений")
+            continue
+
+        print(f"\n📂 Обрабатываем {source_folder} → {target_class}:")
+        print(f"  Найдено {len(images)} изображений")
+
+        # Разделяем
+        train_imgs, val_imgs = train_test_split(
+            images, test_size=0.2, random_state=42
+        )
+
+        # Копируем train
+        print(f"  Копируем {len(train_imgs)} в train...")
+        for img in tqdm(train_imgs, desc=f"  {target_class} train"):
+            src = os.path.join(source_dir, img)
+            dst = os.path.join(Config.DATA_DIR, 'train', target_class, img)
+            shutil.copy2(src, dst)
+
+        # Копируем val
+        print(f"  Копируем {len(val_imgs)} в val...")
+        for img in tqdm(val_imgs, desc=f"  {target_class} val"):
+            src = os.path.join(source_dir, img)
+            dst = os.path.join(Config.DATA_DIR, 'val', target_class, img)
+            shutil.copy2(src, dst)
+
+        total_images += len(images)
+        print(f"  ✓ {target_class}: {len(train_imgs)} train, {len(val_imgs)} val")
+
+    # Проверяем финальную структуру
+    print(f"\n✅ Готово! Всего {total_images} изображений")
+    print("\n📂 Финальная структура данных:")
+
+    for split in ['train', 'val']:
+        split_total = 0
+        print(f"\n  {split.upper()}:")
+        for cls in Config.CLASS_NAMES:
+            cls_dir = os.path.join(Config.DATA_DIR, split, cls)
+            if os.path.exists(cls_dir):
+                count = len([f for f in os.listdir(cls_dir)
+                           if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
+                print(f"    {cls}: {count} изображений")
+                split_total += count
+        print(f"    Всего: {split_total}")
+
+    return True
+
+# ================================================
+# ДАТАСЕТ С ОБРАБОТКОЙ ПОВРЕЖДЕННЫХ ИЗОБРАЖЕНИЙ
+# ================================================
+class RobustWagonDataset(Dataset):
+    """Надежный датасет с обработкой поврежденных изображений"""
+    def __init__(self, data_dir, transform=None, mode='train'):
+        self.image_paths = []
+        self.labels = []
+        self.transform = transform
+
+        data_path = os.path.join(data_dir, mode)
+
+        for class_idx, class_name in enumerate(Config.CLASS_NAMES):
+            class_dir = os.path.join(data_path, class_name)
+            if not os.path.exists(class_dir):
+                print(f"⚠ Папка {class_dir} не найдена!")
+                continue
+
+            images = [f for f in os.listdir(class_dir)
+                     if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
+
+            for img in images:
+                self.image_paths.append(os.path.join(class_dir, img))
+                self.labels.append(class_idx)
+
+        print(f"✅ {mode.upper()}: загружено {len(self.image_paths)} изображений")
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        # Загружаем изображение с обработкой ошибок
+        img_path = self.image_paths[idx]
+
+        # Пытаемся загрузить изображение
+        image = load_image_safe(img_path)
+
+        # Применяем трансформации
+        if self.transform:
+            image = self.transform(image)
+
+        return image, self.labels[idx]
+
+# ================================================
+# МОДЕЛЬ
+# ================================================
+def create_simple_model():
+    """Создание простой модели"""
+    # Используем EfficientNet-B2 как компромисс между скоростью и точностью
+    model = models.efficientnet_b2(weights='DEFAULT')
+
+    # Заменяем классификатор
+    in_features = model.classifier[1].in_features
+    model.classifier = nn.Sequential(
+        nn.Dropout(p=0.3),
+        nn.Linear(in_features, Config.NUM_CLASSES)
+    )
+
+    return model.to(Config.DEVICE)
+
+# ================================================
+# ОБУЧЕНИЕ (РАБОЧАЯ ВЕРСИЯ)
+# ================================================
+def train_simple_model():
+    """Простая и рабочая функция обучения"""
+    print("\n" + "="*60)
+    print("🏋️‍♂️ НАЧИНАЕМ ОБУЧЕНИЕ")
+    print("=" * 60)
+
+    # Очищаем кэш GPU
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    Config.print_info()
+
+    # Получаем трансформации
+    train_transform, val_transform = get_transforms()
+
+    # Создаем датасеты
+    print("\n📥 Загрузка данных...")
+    train_dataset = RobustWagonDataset(
+        Config.DATA_DIR,
+        transform=train_transform,
+        mode='train'
+    )
+
+    val_dataset = RobustWagonDataset(
+        Config.DATA_DIR,
+        transform=val_transform,
+        mode='val'
+    )
+
+    if len(train_dataset) == 0:
+        print("❌ Обучающие данные не найдены!")
+        return None, None
+
+    # Создаем DataLoader
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=Config.BATCH_SIZE,
+        shuffle=True,
+        num_workers=0,  # 0 для Windows чтобы избежать проблем
+        pin_memory=True if torch.cuda.is_available() else False
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=Config.BATCH_SIZE,
+        shuffle=False,
+        num_workers=0,  # 0 для Windows чтобы избежать проблем
+        pin_memory=True if torch.cuda.is_available() else False
+    )
+
+    # Создаем модель
+    print("\n🧠 Создание модели...")
+    model = create_simple_model()
+
+    # Функция потерь и оптимизатор
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=1e-4)
+    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
+
+    # История обучения
+    history = {
+        'train_loss': [], 'train_acc': [],
+        'val_loss': [], 'val_acc': []
+    }
+
+    best_val_acc = 0.0
+
+    print("\n" + "="*50)
+    print("🏁 НАЧАЛО ОБУЧЕНИЯ")
+    print("="*50)
+
+    for epoch in range(Config.NUM_EPOCHS):
+        print(f"\n📅 ЭПОХА {epoch + 1}/{Config.NUM_EPOCHS}")
+
+        # ===== ОБУЧЕНИЕ =====
+        model.train()
+        train_loss = 0.0
+        train_correct = 0
+        train_total = 0
+
+        train_bar = tqdm(train_loader, desc='Training')
+        for images, labels in train_bar:
+            # Перемещаем данные на GPU
+            images = images.to(Config.DEVICE)
+            labels = labels.to(Config.DEVICE)
+
+            # Forward pass
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+
+            # Backward pass
+            loss.backward()
+            optimizer.step()
+
+            # Статистика
+            train_loss += loss.item()
+            _, predicted = outputs.max(1)
+            train_total += labels.size(0)
+            train_correct += predicted.eq(labels).sum().item()
+
+            # Обновляем прогресс-бар
+            train_bar.set_postfix({
+                'Loss': f'{loss.item():.4f}',
+                'Acc': f'{100.*train_correct/train_total:.1f}%'
+            })
+
+        # Средние значения за эпоху
+        avg_train_loss = train_loss / len(train_loader)
+        train_accuracy = train_correct / train_total
+
+        # ===== ВАЛИДАЦИЯ =====
+        model.eval()
+        val_loss = 0.0
+        val_correct = 0
+        val_total = 0
+        all_preds = []
+        all_labels = []
+
+        with torch.no_grad():
+            val_bar = tqdm(val_loader, desc='Validation')
+            for images, labels in val_bar:
+                images = images.to(Config.DEVICE)
+                labels = labels.to(Config.DEVICE)
+
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+
+                val_loss += loss.item()
+                _, predicted = outputs.max(1)
+                val_total += labels.size(0)
+                val_correct += predicted.eq(labels).sum().item()
+
+                all_preds.extend(predicted.cpu().numpy())
+                all_labels.extend(labels.cpu().numpy())
+
+        avg_val_loss = val_loss / len(val_loader)
+        val_accuracy = val_correct / val_total
+
+        # Обновляем scheduler
+        scheduler.step()
+
+        # Сохраняем историю
+        history['train_loss'].append(avg_train_loss)
+        history['train_acc'].append(train_accuracy)
+        history['val_loss'].append(avg_val_loss)
+        history['val_acc'].append(val_accuracy)
+
+        # Сохраняем лучшую модель
+        if val_accuracy > best_val_acc:
+            best_val_acc = val_accuracy
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'val_acc': val_accuracy,
+                'train_acc': train_accuracy,
+                'class_names': Config.CLASS_NAMES
+            }, Config.MODEL_SAVE_PATH)
+            print(f"💾 Сохранена лучшая модель! Точность: {val_accuracy:.4f}")
+
+        # Выводим статистику
+        print(f"📊 Результаты эпохи {epoch + 1}:")
+        print(f"  Train Loss: {avg_train_loss:.4f}, Acc: {train_accuracy:.4f}")
+        print(f"  Val Loss: {avg_val_loss:.4f}, Acc: {val_accuracy:.4f}")
+        print(f"  LR: {scheduler.get_last_lr()[0]:.2e}")
+
+    # ===== ВИЗУАЛИЗАЦИЯ РЕЗУЛЬТАТОВ =====
+    print("\n📈 Визуализация результатов...")
+
+    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
+
+    # График потерь
+    axes[0, 0].plot(history['train_loss'], label='Train', marker='o', linewidth=2)
+    axes[0, 0].plot(history['val_loss'], label='Val', marker='s', linewidth=2)
+    axes[0, 0].set_title('Loss History', fontsize=14, fontweight='bold')
+    axes[0, 0].set_xlabel('Epoch')
+    axes[0, 0].set_ylabel('Loss')
+    axes[0, 0].legend()
+    axes[0, 0].grid(True, alpha=0.3)
+
+    # График точности
+    axes[0, 1].plot(history['train_acc'], label='Train', marker='o', linewidth=2)
+    axes[0, 1].plot(history['val_acc'], label='Val', marker='s', linewidth=2)
+    axes[0, 1].set_title('Accuracy History', fontsize=14, fontweight='bold')
+    axes[0, 1].set_xlabel('Epoch')
+    axes[0, 1].set_ylabel('Accuracy')
+    axes[0, 1].legend()
+    axes[0, 1].grid(True, alpha=0.3)
+
+    # Confusion Matrix
+    try:
+        cm = confusion_matrix(all_labels, all_preds)
+        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
+                   xticklabels=Config.CLASS_NAMES,
+                   yticklabels=Config.CLASS_NAMES,
+                   ax=axes[1, 0])
+        axes[1, 0].set_title('Confusion Matrix', fontsize=14, fontweight='bold')
+        axes[1, 0].set_xlabel('Predicted')
+        axes[1, 0].set_ylabel('True')
+    except:
+        axes[1, 0].text(0.5, 0.5, 'Confusion Matrix\nне доступна',
+                       ha='center', va='center', fontsize=12)
+        axes[1, 0].set_title('Confusion Matrix', fontsize=14, fontweight='bold')
+        axes[1, 0].axis('off')
+
+    # Classification Report
+    try:
+        report = classification_report(all_labels, all_preds,
+                                      target_names=Config.CLASS_NAMES)
+        axes[1, 1].text(0, 1, report, fontsize=10, fontfamily='monospace',
+                       verticalalignment='top', transform=axes[1, 1].transAxes)
+    except:
+        axes[1, 1].text(0.5, 0.5, 'Classification Report\nне доступен',
+                       ha='center', va='center', fontsize=12)
+
+    axes[1, 1].set_title('Classification Report', fontsize=14, fontweight='bold')
+    axes[1, 1].axis('off')
+
+    plt.tight_layout()
+    results_path = os.path.join(os.getcwd(), 'training_results.png')
+    plt.savefig(results_path, dpi=100, bbox_inches='tight')
+    plt.show()
+
+    # Выводим отчет по классификации
+    print("\n" + "="*60)
+    print("📋 ОТЧЕТ ПО КЛАССИФИКАЦИИ")
+    print("="*60)
+    try:
+        print(classification_report(all_labels, all_preds,
+                                  target_names=Config.CLASS_NAMES))
+    except:
+        print("Отчет не доступен")
+
+    print("\n" + "="*60)
+    print("🎉 ОБУЧЕНИЕ ЗАВЕРШЕНО!")
+    print("="*60)
+    print(f"🏆 Лучшая точность на валидации: {best_val_acc:.4f}")
+    print(f"💾 Модель сохранена: {Config.MODEL_SAVE_PATH}")
+    print("\n📋 Классы модели:")
+    for i, cls in enumerate(Config.CLASS_NAMES):
+        print(f"  {i}: {cls}")
+
+    return model, history
+
+# ================================================
+# ПРЕДСКАЗАНИЕ С ОБРАБОТКОЙ ПОВРЕЖДЕННЫХ ИЗОБРАЖЕНИЙ
+# ================================================
+def predict_single_image():
+    """Предсказание для одного изображения с обработкой поврежденных файлов"""
+    if not os.path.exists(Config.MODEL_SAVE_PATH):
+        print("❌ Модель не обучена!")
+        return
+
+    print("\n📤 Введите путь к изображению для классификации...")
+    image_path = input("Введите полный путь к изображению: ").strip()
+    image_path = image_path.replace('/', '\\')
+    
+    if not os.path.exists(image_path):
+        print(f"❌ Файл не найден: {image_path}")
+        return
+
+    print(f"✅ Изображение найдено: {os.path.basename(image_path)}")
+
+    # Пытаемся восстановить поврежденное изображение
+    print("🔧 Проверка целостности изображения...")
+    repair_success = repair_image_file(image_path)
+    if repair_success:
+        print("✅ Изображение восстановлено")
+
+    # Загружаем модель
+    model = create_simple_model()
+    checkpoint = torch.load(Config.MODEL_SAVE_PATH, map_location=Config.DEVICE)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+
+    # Получаем трансформации
+    _, val_transform = get_transforms()
+
+    # Загружаем и обрабатываем изображение
+    try:
+        # Используем безопасную загрузку
+        print("🖼️ Загрузка изображения...")
+        image = load_image_safe(image_path)
+
+        # Проверяем, что изображение загрузилось
+        if image is None:
+            print("❌ Не удалось загрузить изображение")
+            return None
+
+        print(f"✅ Изображение загружено. Размер: {image.size}")
+
+        # Применяем трансформации
+        print("🔄 Применение трансформаций...")
+        input_tensor = val_transform(image).unsqueeze(0).to(Config.DEVICE)
+
+        # Предсказание
+        print("🧠 Выполнение предсказания...")
+        with torch.no_grad():
+            outputs = model(input_tensor)
+            probabilities = torch.nn.functional.softmax(outputs, dim=1)
+            predicted_idx = torch.argmax(probabilities, dim=1).item()
+            confidence = probabilities[0][predicted_idx].item()
+
+        predicted_class = Config.CLASS_NAMES[predicted_idx]
+
+        # Выводим результат
+        print("\n" + "="*60)
+        print("🎯 РЕЗУЛЬТАТ КЛАССИФИКАЦИИ")
+        print("="*60)
+        print(f"📋 Класс: {predicted_class}")
+        print(f"📊 Уверенность: {confidence:.2%}")
+        print(f"\n📈 Распределение вероятностей:")
+        for i, cls in enumerate(Config.CLASS_NAMES):
+            prob = probabilities[0][i].item()
+            prob_percent = prob * 100
+            # Создаем прогресс-бар
+            bar_length = 20
+            filled_length = int(bar_length * prob)
+            bar = '█' * filled_length + '░' * (bar_length - filled_length)
+
+            # Подсвечиваем предсказанный класс
+            if i == predicted_idx:
+                print(f"  ⭐ {cls}: {bar} {prob_percent:5.1f}% ({prob:.4f})")
+            else:
+                print(f"     {cls}: {bar} {prob_percent:5.1f}% ({prob:.4f})")
+
+        # Визуализация
+        plt.figure(figsize=(14, 6))
+
+        # Изображение
+        plt.subplot(1, 3, 1)
+        plt.imshow(image)
+        plt.title(f"Входное изображение\n{os.path.basename(image_path)}", fontsize=12)
+        plt.axis('off')
+
+        # График вероятностей
+        plt.subplot(1, 3, 2)
+        colors = ['#FF6B6B', '#4ECDC4', '#45B7D1']
+        probs = probabilities[0].cpu().numpy()
+
+        bars = plt.bar(Config.CLASS_NAMES, probs, color=colors, alpha=0.7, edgecolor='black', linewidth=2)
+        bars[predicted_idx].set_alpha(1.0)
+        bars[predicted_idx].set_linewidth(3)
+        bars[predicted_idx].set_edgecolor('red')
+
+        plt.title(f"Предсказание: {predicted_class}\nУ��еренность: {confidence:.2%}",
+                 fontsize=14, fontweight='bold')
+        plt.ylim([0, 1.1])
+        plt.ylabel('Вероятность', fontsize=12)
+        plt.grid(True, alpha=0.3, axis='y')
+
+        # Добавляем значения на столбцы
+        for i, (bar, prob) in enumerate(zip(bars, probs)):
+            plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.02,
+                    f'{prob:.2%}', ha='center', va='bottom', fontsize=11,
+                    fontweight='bold' if i == predicted_idx else 'normal',
+                    color='red' if i == predicted_idx else 'black')
+
+        # Тепловая карта вероятностей
+        plt.subplot(1, 3, 3)
+        prob_matrix = probabilities.cpu().numpy().reshape(-1, 1)
+        plt.imshow(prob_matrix, cmap='RdYlGn', aspect='auto', vmin=0, vmax=1)
+        plt.colorbar(label='Вероятность')
+        plt.yticks(range(len(Config.CLASS_NAMES)), Config.CLASS_NAMES)
+        plt.xticks([])
+        plt.title('Тепловая карта вероятностей', fontsize=14, fontweight='bold')
+
+        # Добавляем значения в ячейки
+        for i, prob in enumerate(prob_matrix):
+            plt.text(0, i, f'{prob[0]:.3f}', ha='center', va='center',
+                    color='white' if prob[0] > 0.5 else 'black',
+                    fontweight='bold' if i == predicted_idx else 'normal')
+
+        plt.tight_layout()
+        plt.show()
+
+        # Дополнительная информация
+        print("\n📝 ИНТЕРПРЕТАЦИЯ РЕЗУЛЬТАТА:")
+        if predicted_class == 'pered':
+            print("  🚂 Передняя часть вагона обнаружена")
+        elif predicted_class == 'zad':
+            print("  🚂 Задняя часть вагона обнаружена")
+        elif predicted_class == 'none':
+            print("  ⭕ Вагон не обнаружен")
+
+        if confidence > 0.9:
+            print("  ✅ Высокая уверенность предсказания")
+        elif confidence > 0.7:
+            print("  ⚠ Средняя уверенность предсказания")
+        else:
+            print("  ❓ Низкая уверенность, возможно неоднозначное изображение")
+
+        return predicted_class, confidence
+
+    except Exception as e:
+        print(f"\n❌ Критическая ошибка при обработке изображения: {e}")
+        print("\n🔧 ВОЗМОЖНЫЕ РЕШЕНИЯ:")
+        print("  1. Попробуйте загрузить другое изображение")
+        print("  2. Убедитесь, что файл не поврежден")
+        print("  3. Проверьте формат файла (должен быть JPG, PNG)")
+
+        import traceback
+        traceback.print_exc()
+        return None
+
+# ================================================
+# ПАКЕТНОЕ ТЕСТИРОВАНИЕ
+# ================================================
+def batch_test_images():
+    """Тестирование модели на нескольких изображениях"""
+    if not os.path.exists(Config.MODEL_SAVE_PATH):
+        print("❌ Модель не обучена!")
+        return
+
+    print("\n📤 Введите путь к папке с изображениями для тестирования...")
+    folder_path = input("Введите полный путь к папке: ").strip()
+    folder_path = folder_path.replace('/', '\\')
+    
+    if not os.path.exists(folder_path):
+        print(f"❌ Папка не найдена: {folder_path}")
+        return
+
+    # Получаем список изображений
+    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff')
+    image_files = [f for f in os.listdir(folder_path) 
+                   if f.lower().endswith(image_extensions)]
+    
+    if not image_files:
+        print(f"❌ В папке нет изображений: {folder_path}")
+        return
+
+    print(f"✅ Найдено {len(image_files)} изображений")
+
+    # Загружаем модель
+    model = create_simple_model()
+    checkpoint = torch.load(Config.MODEL_SAVE_PATH, map_location=Config.DEVICE)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+
+    # Получаем трансформации
+    _, val_transform = get_transforms()
+
+    results = []
+
+    for image_name in image_files:
+        print(f"\n🔍 Обработка: {image_name}")
+
+        # Путь к изображению
+        image_path = os.path.join(folder_path, image_name)
+
+        try:
+            # Безопасная загрузка
+            image = load_image_safe(image_path)
+            if image is None:
+                print(f"  ❌ Не удалось загрузить {image_name}")
+                continue
+
+            # Пред��казание
+            input_tensor = val_transform(image).unsqueeze(0).to(Config.DEVICE)
+
+            with torch.no_grad():
+                outputs = model(input_tensor)
+                probabilities = torch.nn.functional.softmax(outputs, dim=1)
+                predicted_idx = torch.argmax(probabilities, dim=1).item()
+                confidence = probabilities[0][predicted_idx].item()
+
+            predicted_class = Config.CLASS_NAMES[predicted_idx]
+            results.append((image_name, predicted_class, confidence))
+
+            print(f"  ✅ {predicted_class} ({confidence:.2%})")
+
+        except Exception as e:
+            print(f"  ❌ Ошибка: {e}")
+            results.append((image_name, "ERROR", 0.0))
+
+    # Выводим сводку
+    print("\n" + "="*60)
+    print("📊 СВОДКА ПО ТЕСТИРОВАНИЮ")
+    print("="*60)
+
+    if not results:
+        print("❌ Нет результатов")
+        return
+
+    # Группируем по классам
+    class_summary = {}
+    for _, cls, conf in results:
+        if cls not in class_summary:
+            class_summary[cls] = []
+        class_summary[cls].append(conf)
+
+    print("\n📈 Статистика по классам:")
+    for cls, confidences in class_summary.items():
+        if cls == "ERROR":
+            print(f"  ❌ Ошибки: {len(confidences)} изображений")
+        else:
+            avg_conf = np.mean(confidences) if confidences else 0
+            print(f"  {cls}: {len(confidences)} изображений, средняя уверенность: {avg_conf:.2%}")
+
+    # Подробные результаты
+    print("\n📋 Подробные результаты:")
+    for i, (img_name, cls, conf) in enumerate(results, 1):
+        if cls == "ERROR":
+            print(f"  {i:2d}. ❌ {img_name}")
+        else:
+            print(f"  {i:2d}. ✅ {img_name}: {cls} ({conf:.2%})")
+
+    return results
+
+# ================================================
+# ГЛАВНОЕ МЕНЮ
+# ================================================
+def main_menu():
+    """Главное меню"""
+    print("\n" + "="*60)
+    print("🚂 КЛАССИФИКАТОР ВАГОНОВ")
+    print("="*60)
+    print(f"📱 Устройство: {Config.DEVICE}")
+    if Config.DEVICE.type == 'cuda':
+        print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
+        print(f"💾 Память: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
+
+    while True:
+        print("\n" + "="*60)
+        print("ГЛАВНОЕ МЕНЮ:")
+        print("1. 📊 Подготовить данные")
+        print("2. 🏋️‍♂️ Обучить модель")
+        print("3. 🔍 Протестировать одно изображение")
+        print("4. 📦 Протестировать несколько изображений")
+        print("5. 📈 Показать графики")
+        print("6. 🧹 Очистить кэш")
+        print("0. ❌ Выход")
+        print("="*60)
+
+        choice = input("\nВыберите действие (0-6): ").strip()
+
+        if choice == '1':
+            # Подготовка данных
+            print("\n" + "="*60)
+            print("ПОДГОТОВКА ДАННЫХ")
+            print("="*60)
+
+            success = prepare_data_simple()
+            if success:
+                print("\n✅ Данные готовы к обучению!")
+            else:
+                print("\n❌ Ошибка при подготовке данных")
+
+        elif choice == '2':
+            # Обучение модели
+            if not os.path.exists(Config.DATA_DIR):
+                print("\n❌ Данные не подготовлены! Сначала выполните шаг 1.")
+                continue
+
+            try:
+                model, history = train_simple_model()
+                if model is not None:
+                    print("\n✅ Обучение успешно завершено!")
+            except Exception as e:
+                print(f"\n❌ Ошибка при обучении: {e}")
+                import traceback
+                traceback.print_exc()
+
+        elif choice == '3':
+            # Тестирование одного изображения
+            if not os.path.exists(Config.MODEL_SAVE_PATH):
+                print("\n❌ Модель не обучена! Сначала выполните шаг 2.")
+                continue
+
+            try:
+                result = predict_single_image()
+                if result:
+                    print("\n✅ Предсказание выполнено успешно!")
+            except Exception as e:
+                print(f"\n❌ Ошибка при предсказании: {e}")
+
+        elif choice == '4':
+            # Пакетное тестирование
+            if not os.path.exists(Config.MODEL_SAVE_PATH):
+                print("\n❌ Модель н�� обучена! Сначала выполните шаг 2.")
+                continue
+
+            try:
+                results = batch_test_images()
+                if results:
+                    print("\n✅ Пакетное тестирование завершено!")
+            except Exception as e:
+                print(f"\n❌ Ошибка при пакетном тестировании: {e}")
+
+        elif choice == '5':
+            # Показать графики
+            results_path = os.path.join(os.getcwd(), 'training_results.png')
+            if os.path.exists(results_path):
+                print("\n📊 Графики обучения:")
+                try:
+                    img = plt.imread(results_path)
+                    plt.figure(figsize=(12, 8))
+                    plt.imshow(img)
+                    plt.axis('off')
+                    plt.show()
+                except Exception as e:
+                    print(f"❌ Ошибка при загрузке графиков: {e}")
+            else:
+                print("\n❌ Графики не найдены. Сначала обучите модель.")
+
+        elif choice == '6':
+            # Очистка кэша
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                print("✅ Кэш GPU очищен!")
+            else:
+                print("⚠ GPU не доступна")
+
+        elif choice == '0':
+            print("\n👋 До свидания!")
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            break
+
+        else:
+            print("\n❌ Неверный выбор. Пожалуйста, выберите от 0 до 6.")
+
+# ================================================
+# ЗАПУСК
+# ================================================
+if __name__ == "__main__":
+    print("🚂 КЛАССИФИКАТОР ВАГОНОВ ДЛЯ WINDOWS 10")
+    print("=" * 60)
+    print("📦 Зависимости:")
+    print("Установите библиотеки из файла requirements.txt:")
+    print("pip install -r requirements.txt")
+    print("=" * 60)
+    
+    # Запускаем меню
+    main_menu()