feat: add metrics

2026-06-16 04:59:51 +00:00
parent 14968dd4d4
commit 934a4cbff4
9 changed files with 550 additions and 624 deletions
@@ -0,0 +1,283 @@
+import os
+import random
+import warnings
+from pathlib import Path
+from PIL import Image
+import pandas as pd
+import numpy as np
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import torchvision.transforms as T
+import timm
+
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.metrics import confusion_matrix
+
+# Подавление предупреждений цветовых профилей
+warnings.filterwarnings("ignore", message=".*Unknown Adobe color transform code.*")
+
+# Настройки окружения
+DATA_ROOT = Path("./NFS/Thesis/Emoset/EmoSet-118K")
+# ВАЖНО: Добавили путь для медиа файлов
+MEDIA_DIR = Path("./src/scripts/media")
+MEDIA_DIR.mkdir(parents=True, exist_ok=True)
+
+BATCH_SIZE = 64
+EPOCHS = 30
+LR = 5e-5
+NUM_WORKERS = 32
+PATIENCE = 7
+
+# Маппинг классов
+CLASS_MAPPING = {
+    "amusement": 0, "anger": 1, "awe": 2, "contentment": 3,
+    "disgust": 4, "excitement": 5, "fear": 6, "sadness": 7
+}
+# Инвертированный маппинг для графиков
+INV_CLASS_MAPPING = {v: k for k, v in CLASS_MAPPING.items()}
+CLASS_NAMES = [INV_CLASS_MAPPING[i] for i in range(len(CLASS_MAPPING))]
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Устройство: {DEVICE}")
+
+# Фиксация генераторов псевдослучайных чисел
+def set_seed(seed=42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+set_seed()
+
+# Инициализация структур данных
+class EmoSetDataset(Dataset):
+    def __init__(self, root: Path | str, split: str, transform=None):
+        self.root = Path(root) / split
+        self.df = pd.read_csv(self.root / "labels.csv")
+        self.transform = transform
+        
+        # Фильтрация датафрейма
+        self.df = self.df[self.df["label"].isin(CLASS_MAPPING.keys())].reset_index(drop=True)
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        row = self.df.iloc[idx]
+        img_path = self.root / "images" / row["filename"]
+
+        try:
+            img = Image.open(img_path).convert("RGB")
+        except Exception:
+            img = Image.new("RGB", (256, 256), (0, 0, 0))
+
+        if self.transform:
+            img_tensor = self.transform(img)
+        else:
+            img_tensor = T.ToTensor()(img)
+            
+        label_idx = CLASS_MAPPING[row["label"]]
+        return img_tensor, label_idx
+
+# Трансформации
+base_tf = [
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+]
+
+train_transform = T.Compose([
+    T.Resize(256, antialias=True),
+    T.RandomCrop(224),
+    T.RandomHorizontalFlip(),
+    *base_tf
+])
+
+val_transform = T.Compose([
+    T.Resize(256, antialias=True),
+    T.CenterCrop(224),
+    *base_tf
+])
+
+train_ds = EmoSetDataset(DATA_ROOT, "train", transform=train_transform)
+val_ds   = EmoSetDataset(DATA_ROOT, "val", transform=val_transform)
+
+train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=True)
+val_loader   = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)
+
+# Инициализация модели и оптимизатора
+model = timm.create_model("resnet50", pretrained=True, num_classes=8, drop_rate=0.3)
+model.to(DEVICE)
+
+criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
+
+optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-3)
+scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)
+
+# Функции для отрисовки графиков
+def plot_learning_curves(history):
+    """Отрисовка графиков функции потерь и точности"""
+    epochs = range(1, len(history['train_loss']) + 1)
+    
+    plt.figure(figsize=(14, 5))
+    
+    # График Loss
+    plt.subplot(1, 2, 1)
+    plt.plot(epochs, history['train_loss'], 'b-', label='Train Loss')
+    plt.plot(epochs, history['val_loss'], 'r--', label='Validation Loss')
+    plt.title('График функции потерь (Loss)', fontsize=14)
+    plt.xlabel('Эпохи', fontsize=12)
+    plt.ylabel('Loss', fontsize=12)
+    plt.legend()
+    plt.grid(True, linestyle=':', alpha=0.7)
+    
+    # График Accuracy
+    plt.subplot(1, 2, 2)
+    plt.plot(epochs, history['train_acc'], 'b-', label='Train Accuracy')
+    plt.plot(epochs, history['val_acc'], 'r--', label='Validation Accuracy')
+    plt.title('График точности (Accuracy)', fontsize=14)
+    plt.xlabel('Эпохи', fontsize=12)
+    plt.ylabel('Accuracy', fontsize=12)
+    plt.legend()
+    plt.grid(True, linestyle=':', alpha=0.7)
+    
+    plt.tight_layout()
+    plot_path = MEDIA_DIR / "training_history.png"
+    plt.savefig(plot_path, dpi=300, bbox_inches='tight')
+    plt.close()
+    print(f"[INFO] График обучения сохранен в: {plot_path}")
+
+def plot_confusion_matrix(y_true, y_pred):
+    """Отрисовка тепловой матрицы ошибок"""
+    cm = confusion_matrix(y_true, y_pred)
+    
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
+                xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES,
+                cbar_kws={'label': 'Количество сэмплов'})
+    
+    plt.title('Матрица ошибок (Confusion Matrix) - ResNet50', fontsize=16, pad=20)
+    plt.ylabel('Истинные классы (Ground Truth)', fontsize=12)
+    plt.xlabel('Предсказанные классы (Predicted)', fontsize=12)
+    
+    plt.xticks(rotation=45, ha='right')
+    plt.yticks(rotation=0)
+    
+    plt.tight_layout()
+    cm_path = MEDIA_DIR / "confusion_matrix_emoset.png"
+    plt.savefig(cm_path, dpi=300, bbox_inches='tight')
+    plt.close()
+    print(f"[INFO] Матрица ошибок сохранена в: {cm_path}")
+
+# Логика эпохи обучения
+def train_epoch(current_model, loader):
+    current_model.train()
+    total_loss, correct_preds, total_samples = 0.0, 0, 0
+
+    for imgs, labels in tqdm(loader, desc="Тренировка", leave=False, smoothing=0):
+        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
+
+        optimizer.zero_grad(set_to_none=True)
+        logits = current_model(imgs)
+        loss = criterion(logits, labels)
+
+        loss.backward()
+        optimizer.step()
+
+        total_loss += loss.item() * imgs.size(0)
+        preds = logits.argmax(dim=1)
+        correct_preds += (preds == labels).sum().item()
+        total_samples += labels.size(0)
+
+    return total_loss / total_samples, correct_preds / total_samples
+
+# Логика эпохи валидации с сохранением предсказаний для матрицы ошибок
+@torch.no_grad()
+def val_epoch(current_model, loader, return_preds=False):
+    current_model.eval()
+    total_loss, correct_preds, total_samples = 0.0, 0, 0
+    all_preds, all_labels = [], []
+
+    for imgs, labels in tqdm(loader, desc="Валидация", leave=False, smoothing=0):
+        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
+
+        logits = current_model(imgs)
+        loss = criterion(logits, labels)
+
+        total_loss += loss.item() * imgs.size(0)
+        preds = logits.argmax(dim=1)
+        
+        correct_preds += (preds == labels).sum().item()
+        total_samples += labels.size(0)
+        
+        if return_preds:
+            all_preds.extend(preds.cpu().numpy())
+            all_labels.extend(labels.cpu().numpy())
+
+    avg_loss = total_loss / total_samples
+    avg_acc = correct_preds / total_samples
+    
+    if return_preds:
+        return avg_loss, avg_acc, all_labels, all_preds
+    return avg_loss, avg_acc
+
+if __name__ == "__main__":
+    best_val_acc = 0.0
+    best_val_loss = float('inf')
+    epochs_no_improve = 0
+    checkpoint_path = "./emosetV2_resnet50_best.pth"
+    
+    # Словарь для хранения истории обучения
+    history = {
+        'train_loss': [], 'train_acc': [],
+        'val_loss': [], 'val_acc': []
+    }
+    
+    # Переменные для хранения лучших предсказаний для матрицы
+    best_labels, best_preds = [], []
+
+    print("Старт обучения.")
+
+    for epoch in range(1, EPOCHS + 1):
+        train_loss, train_acc = train_epoch(model, train_loader)
+        
+        # Получаем предсказания только если это может быть лучшая эпоха
+        val_loss, val_acc, val_labels, val_preds = val_epoch(model, val_loader, return_preds=True)
+
+        scheduler.step()
+
+        # Запись в историю
+        history['train_loss'].append(train_loss)
+        history['train_acc'].append(train_acc)
+        history['val_loss'].append(val_loss)
+        history['val_acc'].append(val_acc)
+
+        print(f"[{epoch}/{EPOCHS}] Train Loss: {train_loss:.4f}, Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f}, Acc: {val_acc:.4f}")
+
+        # Сохранение лучших весов по Accuracy
+        if val_acc > best_val_acc:
+            best_val_acc = val_acc
+            best_labels = val_labels # Сохраняем предсказания лучшей модели
+            best_preds = val_preds
+            torch.save(model.state_dict(), checkpoint_path)
+            print(f"Сохранен чекпоинт (Acc: {best_val_acc:.4f})")
+
+        # Оценка переобучения по Loss (Early Stopping)
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            epochs_no_improve = 0
+        else:
+            epochs_no_improve += 1
+            if epochs_no_improve >= PATIENCE:
+                print(f"Ранняя остановка: метрика валидации не улучшается {PATIENCE} эпох.")
+                break
+
+    print("Процесс обучения завершен. Генерирую графики для диссертации...")
+    plot_learning_curves(history)
+    plot_confusion_matrix(best_labels, best_preds)
+    print("Все медиафайлы успешно созданы!")