model.py - master - madin/Archipelag

model.py
396 строк15 KB
madin
upload files
12 авг 2025, 13:36
12 авг 2025, 13:366c6518e
import os
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, Subset
from torchvision import transforms, models
from torch.utils.tensorboard import SummaryWriter
import cv2
import datetime
import matplotlib.pyplot as plt
from matplotlib import font_manager
from sklearn.metrics import mean_squared_error, mean_absolute_error



# Отключим предупреждения о шрифтах
import warnings
warnings.filterwarnings("ignore")

print("PyTorch версия:", torch.__version__)
print("CUDA доступна:", torch.cuda.is_available())
print("Количество CUDA-устройств:", torch.cuda.device_count())

if torch.cuda.is_available():
    print("Текущее CUDA устройство:", torch.cuda.current_device())
    print("Имя устройства:", torch.cuda.get_device_name(torch.cuda.current_device()))
else:
    print("CUDA недоступна. Используется CPU.")

# === Настройки ===
DATA_DIR = "D:\\drone\\dataset_v1"
IMAGE_DIR = os.path.join(DATA_DIR, "images","train")
#print('IMAGE_DIR_1',IMAGE_DIR)
CSV_PATH = 'D:\\drone\\dataset_v1\\labels\\train\\_train_annotations.csv'
MODEL_SAVE_PATH = "best_model_50_100_128_2.pth"
RESULTS_DIR = "D:\\drone\\results"
LOG_DIR = "D:\\drone\\runs"
PLOTS_DIR = "D:\\drone\\results\\plots"
REPORT_CSV = "D:\\drone\\results\\validation_report.csv"

os.makedirs(RESULTS_DIR, exist_ok=True)
os.makedirs(PLOTS_DIR, exist_ok=True)
os.makedirs(LOG_DIR, exist_ok=True)

# Гиперпараметры
BATCH_SIZE = 128
NUM_EPOCHS = 100
LEARNING_RATE = 1e-4
NUM_PARAMS = 2 # 3 params - Lat,Lon, Alt? 2 params - Lat,Lon
IMG_SIZE = 224
PATIENCE = 10 #5
SHOW_COMPARISON_EVERY_N_EPOCH = 5

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Используемое устройство: {device}")

# TensorBoard
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
writer = SummaryWriter(log_dir=os.path.join(LOG_DIR, f"resnet34_reg_{current_time}"))

writer.add_text("Hyperparameters", str({
    "batch_size": BATCH_SIZE,
    "epochs": NUM_EPOCHS,
    "learning_rate": LEARNING_RATE,
    "model": "resnet34",
    "img_size": IMG_SIZE,
    "patience": PATIENCE,
    "show_comparison": SHOW_COMPARISON_EVERY_N_EPOCH
}))

# Трансформации
train_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    #transforms.RandomHorizontalFlip(),
    #transforms.RandomRotation(10),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

val_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


# 🧩 Dataset с именем файла
class ImageParamDataset(Dataset):
    def __init__(self, csv_file, img_dir, transform=None):
        self.data_frame = pd.read_csv(csv_file)
        self.img_dir = img_dir
        self.transform = transform

    def __len__(self):
        return len(self.data_frame)

    def __getitem__(self, idx):
        img_name = self.data_frame.iloc[idx, 0]
        img_path = os.path.join(self.img_dir, img_name)
        image = cv2.imread(img_path)
        if image is None:
            raise FileNotFoundError(f"Не удалось загрузить: {img_path}")
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        params = self.data_frame.iloc[idx, 1:1+NUM_PARAMS].astype(np.float32).values

        if self.transform:
            image = self.transform(image)

        return image, torch.tensor(params, dtype=torch.float32), img_name


# 🧠 Модель
class ResNet34Regressor(nn.Module):
    def __init__(self, num_params=7):
        super(ResNet34Regressor, self).__init__()
        self.base_model = models.resnet34(pretrained=True)
        num_features = self.base_model.fc.in_features
        self.base_model.fc = nn.Identity()
        self.regressor = nn.Sequential(
            nn.Linear(num_features, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, num_params)
        )

    def forward(self, x):
        features = self.base_model(x)
        return self.regressor(features)


# 🔔 Логирование градиентов
def log_gradients(model, step):
    for name, param in model.named_parameters():
        if param.grad is not None:
            writer.add_histogram(f"Gradients/{name}", param.grad.data, step)


# 📊 Валидация: сравнение + визуализация + отчёт
def evaluate_and_compare(model, val_loader, epoch):
    model.eval()
    preds_list = []
    truths_list = []
    names_list = []

    with torch.no_grad():
        for images, targets, filenames in val_loader:
            images = images.to(device)
            outputs = model(images)
            preds_list.append(outputs.cpu().numpy())
            truths_list.append(targets.numpy())
            names_list.extend(filenames)

    preds = np.vstack(preds_list)
    truths = np.vstack(truths_list)
    errors = np.abs(preds - truths)

    # === Метрики ===
    mae_per_param = np.mean(errors, axis=0)
    mse_per_param = np.mean((preds - truths) ** 2, axis=0)
    total_mae = np.mean(mae_per_param)
    total_mse = np.mean(mse_per_param)

    # === Логирование в TensorBoard ===
    writer.add_scalar("Eval/Total_MAE", total_mae, epoch)
    writer.add_scalar("Eval/Total_MSE", total_mse, epoch)
    for i in range(NUM_PARAMS):
        writer.add_scalar(f"Eval/MAE_param_{i+1}", mae_per_param[i], epoch)
        writer.add_scalar(f"Eval/MSE_param_{i+1}", mse_per_param[i], epoch)
        # Гистограмма ошибок MAE по параметрам
        writer.add_histogram(f"Error_MAE/param_{i+1}", errors[:, i], epoch)

    # === Вывод в консоль ===
    print(f"\n📊 Сравнение предсказаний и истины (Валидация, эпоха {epoch+1}):")
    print(f"Средний MAE: {total_mae:.4f}, Средний MSE: {total_mse:.4f}")
    print(f"MAE по параметрам: {np.round(mae_per_param, 4)}")
    print(f"MSE по параметрам: {np.round(mse_per_param, 4)}")

    # === Сохранение CSV-отчёта ===
    param_names = [f"param_{i+1}" for i in range(NUM_PARAMS)]
    df_truth = pd.DataFrame(truths, columns=[f"true_{p}" for p in param_names])
    df_pred = pd.DataFrame(preds, columns=[f"pred_{p}" for p in param_names])
    df_err = pd.DataFrame(errors, columns=[f"error_{p}" for p in param_names])
    df_full = pd.concat([
        pd.DataFrame(names_list, columns=["filename"]),
        df_truth, df_pred, df_err
    ], axis=1)

    # Сохраняем при последнем вызове или можно при каждом
    if epoch == NUM_EPOCHS - 1:
        df_full.to_csv(REPORT_CSV, index=False)
        print(f"\n📎 Полный отчёт о валидации сохранён в '{REPORT_CSV}'")

    # === График: предсказание vs. истина ===
    plt.figure(figsize=(14, 8))
    for i in range(NUM_PARAMS):
        plt.subplot(2, 4, i+1)
        plt.scatter(truths[:, i], preds[:, i], alpha=0.6, s=10)
        min_val = min(truths[:, i].min(), preds[:, i].min())
        max_val = max(truths[:, i].max(), preds[:, i].max())
        plt.plot([min_val, max_val], [min_val, max_val], 'r--', lw=1)
        plt.xlabel("Истинное значение")
        plt.ylabel("Предсказанное")
        plt.title(f"param_{i+1} | MAE={mae_per_param[i]:.3f}")
        plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plot_path = os.path.join(PLOTS_DIR, f"preds_vs_true_epoch{epoch+1}.png")
    plt.savefig(plot_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"📈 График сохранён: {plot_path}")

    # === Доп. график: гистограммы ошибок ===
    plt.figure(figsize=(14, 8))
    for i in range(NUM_PARAMS):
        plt.subplot(2, 4, i+1)
        plt.hist(errors[:, i], bins=20, alpha=0.7, color='skyblue', edgecolor='black')
        plt.xlabel("MAE")
        plt.ylabel("Частота")
        plt.title(f"Ошибка param_{i+1} | Среднее: {mae_per_param[i]:.3f}")
        plt.grid(True, alpha=0.3)
    plt.tight_layout()
    hist_path = os.path.join(PLOTS_DIR, f"error_hist_epoch{epoch+1}.png")
    plt.savefig(hist_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"📊 Гистограммы ошибок сохранены: {hist_path}")

    return total_mse


# 🚀 Обучение
def train_model():
    dataset = ImageParamDataset(CSV_PATH, IMAGE_DIR, transform=train_transform)
    train_size = int(0.8 * len(dataset))
    val_size = len(dataset) - train_size
    train_indices = list(range(train_size))
    val_indices = list(range(train_size, train_size + val_size))
    train_dataset = Subset(dataset, train_indices)
    val_dataset = Subset(dataset, val_indices)

    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0,
                              collate_fn=lambda x: (torch.stack([i[0] for i in x]),
                                                    torch.stack([i[1] for i in x])))

    def collate_with_filename(batch):
        images = torch.stack([item[0] for item in batch])
        targets = torch.stack([item[1] for item in batch])
        filenames = [item[2] for item in batch]
        return images, targets, filenames

    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0,
                            collate_fn=collate_with_filename)

    model = ResNet34Regressor(num_params=NUM_PARAMS).to(device)
    criterion_mse = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=3, factor=0.5)

    best_val_mse = float('inf')
    patience_counter = 0
    global_step = 0

    for epoch in range(NUM_EPOCHS):
        # Обучение
        model.train()
        train_loss = 0.0
        for images, targets in train_loader:
            images, targets = images.to(device), targets.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion_mse(outputs, targets)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

            #loss = criterion_mse(outputs, targets) мне  это нужно

            if global_step % 10 == 0:
                log_gradients(model, global_step)
            global_step += 1

        avg_train_loss = train_loss / len(train_loader)




        # Валидация
        model.eval()
        val_mse_total = 0.0
        num_samples = 0
        with torch.no_grad():
            for images, targets, _ in val_loader:
                images, targets = images.to(device), targets.to(device)
                outputs = model(images)
                val_mse_total += criterion_mse(outputs, targets).item() * images.size(0)
                num_samples += images.size(0)
        avg_val_mse = val_mse_total / num_samples

        scheduler.step(avg_val_mse)

        # Логирование
        writer.add_scalar("Loss/Train_Epoch", avg_train_loss, epoch)
        writer.add_scalar("Loss/Val_MSE", avg_val_mse, epoch)
        writer.add_scalar("LR", optimizer.param_groups[0]['lr'], epoch)

        # Сохранение лучшей модели
        if avg_val_mse < best_val_mse:
            best_val_mse = avg_val_mse
            patience_counter = 0
            torch.save(model.state_dict(), MODEL_SAVE_PATH)
            print(f"✅ Эпоха {epoch+1}: Сохранена лучшая модель (Val MSE = {best_val_mse:.4f})")
        else:
            patience_counter += 1

        # Прогресс
        print(f"Epoch {epoch+1}/{NUM_EPOCHS} | Train MSE: {avg_train_loss:.4f} | Val MSE: {avg_val_mse:.4f}")

        # Сравнение с истиной
        if (epoch + 1) % SHOW_COMPARISON_EVERY_N_EPOCH == 0:
            evaluate_and_compare(model, val_loader, epoch)

        # Early stopping
        if patience_counter >= PATIENCE:
            print(f"🛑 Ранняя остановка на эпохе {epoch+1}")
            break

    writer.close()
    print("📌 Логирование в TensorBoard завершено.")

    # Финальное сравнение
    print("\n" + "="*60)
    print("✅ Финальная оценка и визуализация...")
    evaluate_and_compare(model, val_loader, epoch=NUM_EPOCHS-1)

    return model


# 🔎 Инференс по папке
def infer_folder(model, folder_path, output_csv="D:\\drone\\results\\predictions.csv"):

    results = []
    print("путь_inference",folder_path)
    img_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    model.eval()

    with torch.no_grad():
        for img_file in img_files:
            img_path = os.path.join(folder_path, img_file)
            try:
                image = cv2.imread(img_path)
                if image is None:
                    print(f"⚠️ Пропущено: {img_path}")
                    continue
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                image = val_transform(image).unsqueeze(0).to(device)
                output = model(image)

                pred = output.cpu().numpy().flatten()
                results.append([img_file] + pred.tolist())
            except Exception as e:
                print(f"❌ Ошибка: {img_file} — {e}")



    columns = ["filename"] + [f"param_{i+1}" for i in range(NUM_PARAMS)]
    df = pd.DataFrame(results, columns=columns)
    output_path = output_csv
    df.to_csv(output_path, index=False)
    print(output_path)
    print(f"\n📌 Предсказания сохранены в '{output_path}'")

    return df


# 🏁 Главный блок
if __name__ == "__main__":
    print("🚀 Запуск обучения с визуализацией и отчетами...")
    model = train_model()



    print(f"\n✅ Загружаем лучшую модель из '{MODEL_SAVE_PATH}'...")
    model.load_state_dict(torch.load(MODEL_SAVE_PATH))
    model.to(device)

    print("\n🔍 Инференс по папке...")
    infer_folder(model, folder_path = os.path.join(DATA_DIR, "images","test"))

    print(f"\n🎉 Все результаты сохранены!")
    print("📊 Графики:                results/plots/")
    print("📋 Полный отчёт:            results/validation_report.csv")
    print("📉 Предсказания:            results/predictions.csv")
    print("🚀 TensorBoard:             tensorboard --logdir=runs")