lab4.py - master - supremeur_93/Lab4

lab4.py
860 строк33 KB
Diarra Bourama
upload files
16 янв 2026, 08:14
16 янв 2026, 08:145f64c08
# lab4_vae.py
import os
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import seaborn as sns
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA

# ============================================================
# Конфигурация
# ============================================================
OUTPUT_DIR = "lab4_results"
os.makedirs(OUTPUT_DIR, exist_ok=True)

np.random.seed(42)
tf.random.set_seed(42)

print("🔧 Лабораторная работа 4: Вариационный автоэнкодер (VAE)")

# ============================================================
# Часть 1: Загрузка и подготовка данных
# ============================================================
print("\n" + "=" * 60)
print("📥 ЧАСТЬ 1: Загрузка датасета Fashion-MNIST")
print("=" * 60)

# Загрузка Fashion-MNIST
(x_train, y_train), (x_test, y_test) = keras.datasets.fashion_mnist.load_data()

# Нормализация и преобразование формы данных
x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

# Добавление размерности канала
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

print(f"✅ Форма тренировочных данных: {x_train.shape}")
print(f"✅ Форма тестовых данных: {x_test.shape}")

# Названия классов Fashion-MNIST
class_names = [
    "T-shirt/top",
    "Trouser",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]

# Визуализация примеров исходных данных
print("\n🖼️ Визуализация примеров данных...")
plt.figure(figsize=(12, 12))
for i in range(25):
    plt.subplot(5, 5, i + 1)
    plt.imshow(x_train[i].squeeze(), cmap="gray")
    plt.title(class_names[y_train[i]], fontsize=8)
    plt.axis("off")
plt.suptitle("Примеры из датасета Fashion-MNIST", fontsize=16)
plt.tight_layout()
plt.savefig(
    os.path.join(OUTPUT_DIR, "dataset_samples.png"), dpi=150, bbox_inches="tight"
)
plt.close()

# ============================================================
# Часть 2: Реализация архитектуры VAE
# ============================================================
print("\n" + "=" * 60)
print("🏗️ ЧАСТЬ 2: Создание архитектуры VAE")
print("=" * 60)

# Параметры модели
LATENT_DIM = 2  # Размерность латентного пространства (2D для визуализации)


# Слой Sampling с reparameterization trick
class Sampling(layers.Layer):
    """Слой для сэмплирования из нормального распределения с reparameterization trick"""

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon


# ============================================================
# Энкодер
# ============================================================
print("\n🔨 Создание энкодера...")
encoder_inputs = keras.Input(shape=(28, 28, 1), name="encoder_input")
x = layers.Conv2D(32, 3, activation="relu", strides=2, padding="same")(encoder_inputs)
x = layers.Conv2D(64, 3, activation="relu", strides=2, padding="same")(x)
x = layers.Flatten()(x)
x = layers.Dense(16, activation="relu")(x)

# Предсказание параметров распределения
z_mean = layers.Dense(LATENT_DIM, name="z_mean")(x)
z_log_var = layers.Dense(LATENT_DIM, name="z_log_var")(x)
z = Sampling()([z_mean, z_log_var])

encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")
encoder.summary()

# ============================================================
# Декодер
# ============================================================
print("\n🔨 Создание декодера...")
latent_inputs = keras.Input(shape=(LATENT_DIM,), name="decoder_input")
x = layers.Dense(7 * 7 * 64, activation="relu")(latent_inputs)
x = layers.Reshape((7, 7, 64))(x)
x = layers.Conv2DTranspose(64, 3, activation="relu", strides=2, padding="same")(x)
x = layers.Conv2DTranspose(32, 3, activation="relu", strides=2, padding="same")(x)
decoder_outputs = layers.Conv2DTranspose(1, 3, activation="sigmoid", padding="same")(x)

decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")
decoder.summary()

# ============================================================
# Модель VAE
# ============================================================
print("\n🔨 Создание полной модели VAE...")


class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
        self.reconstruction_loss_tracker = keras.metrics.Mean(
            name="reconstruction_loss"
        )
        self.kl_loss_tracker = keras.metrics.Mean(name="kl_loss")

    @property
    def metrics(self):
        return [
            self.total_loss_tracker,
            self.reconstruction_loss_tracker,
            self.kl_loss_tracker,
        ]

    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(data)
            reconstruction = self.decoder(z)

            # Reconstruction loss (бинарная кросс-энтропия)
            reconstruction_loss = tf.reduce_mean(
                tf.reduce_sum(
                    keras.losses.binary_crossentropy(data, reconstruction), axis=(1, 2)
                )
            )

            # KL divergence loss
            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))

            # Total loss
            total_loss = reconstruction_loss + kl_loss

        gradients = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_weights))

        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)

        return {
            "loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
        }

    def test_step(self, data):
        z_mean, z_log_var, z = self.encoder(data)
        reconstruction = self.decoder(z)

        reconstruction_loss = tf.reduce_mean(
            tf.reduce_sum(
                keras.losses.binary_crossentropy(data, reconstruction), axis=(1, 2)
            )
        )

        kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
        kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))

        total_loss = reconstruction_loss + kl_loss

        return {
            "loss": total_loss,
            "reconstruction_loss": reconstruction_loss,
            "kl_loss": kl_loss,
        }


# Создание модели VAE
vae = VAE(encoder, decoder)
vae.compile(optimizer=keras.optimizers.Adam(learning_rate=1e-3))

print(f"\n📊 Параметры модели:")
print(f"   • Размерность латентного пространства: {LATENT_DIM}")
print(f"   • Параметры энкодера: {encoder.count_params():,}")
print(f"   • Параметры декодера: {decoder.count_params():,}")
print(f"   • Всего параметров: {encoder.count_params() + decoder.count_params():,}")

# ============================================================
# Часть 3: Обучение модели
# ============================================================
print("\n" + "=" * 60)
print("🔄 ЧАСТЬ 3: Обучение модели VAE")
print("=" * 60)

EPOCHS = 30
BATCH_SIZE = 128

print(f"🎯 Гиперпараметры обучения:")
print(f"   • Эпохи: {EPOCHS}")
print(f"   • Размер батча: {BATCH_SIZE}")
print(f"   • Оптимизатор: Adam (lr=1e-3)")

# Обучение модели
print("\n⏳ Начало обучения...")
history = vae.fit(
    x_train,
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    validation_data=(x_test, x_test),
    verbose=1,
)

print("\n✅ Обучение завершено!")

# Сохранение модели
encoder.save(os.path.join(OUTPUT_DIR, "vae_encoder.h5"))
decoder.save(os.path.join(OUTPUT_DIR, "vae_decoder.h5"))
print(f"💾 Модели сохранены в {OUTPUT_DIR}/")

# ============================================================
# Часть 4: Визуализация процесса обучения
# ============================================================
print("\n" + "=" * 60)
print("📈 ЧАСТЬ 4: Анализ процесса обучения")
print("=" * 60)

plt.figure(figsize=(15, 5))

# Total Loss
plt.subplot(1, 3, 1)
plt.plot(history.history["loss"], label="Training Loss", linewidth=2)
plt.plot(history.history["val_loss"], label="Validation Loss", linewidth=2)
plt.title("Total Loss", fontsize=14)
plt.xlabel("Эпоха")
plt.ylabel("Loss")
plt.legend()
plt.grid(True, alpha=0.3)

# Reconstruction Loss
plt.subplot(1, 3, 2)
plt.plot(
    history.history["reconstruction_loss"],
    label="Training Recon Loss",
    linewidth=2,
)
plt.plot(
    history.history["val_reconstruction_loss"],
    label="Validation Recon Loss",
    linewidth=2,
)
plt.title("Reconstruction Loss", fontsize=14)
plt.xlabel("Эпоха")
plt.ylabel("Loss")
plt.legend()
plt.grid(True, alpha=0.3)

# KL Loss
plt.subplot(1, 3, 3)
plt.plot(history.history["kl_loss"], label="Training KL Loss", linewidth=2)
plt.plot(history.history["val_kl_loss"], label="Validation KL Loss", linewidth=2)
plt.title("KL Divergence Loss", fontsize=14)
plt.xlabel("Эпоха")
plt.ylabel("Loss")
plt.legend()
plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig(
    os.path.join(OUTPUT_DIR, "training_history.png"), dpi=150, bbox_inches="tight"
)
plt.close()

print("✅ Графики обучения сохранены")

# ============================================================
# Часть 5: Визуализация реконструкций
# ============================================================
print("\n" + "=" * 60)
print("🖼️ ЧАСТЬ 5: Визуализация реконструкций")
print("=" * 60)


def plot_reconstructions(model, data, n_samples=10):
    """Визуализация оригинальных изображений и их реконструкций"""
    samples = data[:n_samples]
    z_mean, z_log_var, z = model.encoder.predict(samples, verbose=0)
    reconstructions = model.decoder.predict(z, verbose=0)

    plt.figure(figsize=(20, 4))
    for i in range(n_samples):
        # Оригинал
        plt.subplot(2, n_samples, i + 1)
        plt.imshow(samples[i].squeeze(), cmap="gray")
        plt.title("Оригинал", fontsize=8)
        plt.axis("off")

        # Реконструкция
        plt.subplot(2, n_samples, n_samples + i + 1)
        plt.imshow(reconstructions[i].squeeze(), cmap="gray")
        plt.title("Реконструкция", fontsize=8)
        plt.axis("off")

    plt.tight_layout()
    plt.savefig(
        os.path.join(OUTPUT_DIR, "reconstructions.png"), dpi=150, bbox_inches="tight"
    )
    plt.close()


plot_reconstructions(vae, x_test)
print("✅ Реконструкции сохранены")

# ============================================================
# Часть 6: Анализ латентного пространства
# ============================================================
print("\n" + "=" * 60)
print("🔍 ЧАСТЬ 6: Анализ латентного пространства")
print("=" * 60)

# Получение латентных представлений
print("📊 Кодирование тестовых данных в латентное пространство...")
z_mean, z_log_var, z = vae.encoder.predict(x_test, verbose=0)

print(f"   • Размерность латентных векторов: {z_mean.shape}")
print(f"   • Среднее z_mean: {np.mean(z_mean):.4f}")
print(f"   • Стд. отклонение z_mean: {np.std(z_mean):.4f}")

# Визуализация латентного пространства (2D)
if LATENT_DIM == 2:
    print("\n🎨 Визуализация 2D латентного пространства...")
    plt.figure(figsize=(12, 10))
    scatter = plt.scatter(
        z_mean[:, 0], z_mean[:, 1], c=y_test, cmap="tab10", alpha=0.6, s=10
    )
    plt.colorbar(scatter, label="Класс")
    plt.xlabel("z[0]", fontsize=12)
    plt.ylabel("z[1]", fontsize=12)
    plt.title("Латентное пространство VAE (цвета = классы)", fontsize=14)
    plt.grid(True, alpha=0.3)

    # Добавление легенды классов
    for i, name in enumerate(class_names):
        plt.scatter([], [], c=[plt.cm.tab10(i)], label=name, s=50)
    plt.legend(loc="center left", bbox_to_anchor=(1.15, 0.5), fontsize=9)

    plt.tight_layout()
    plt.savefig(
        os.path.join(OUTPUT_DIR, "latent_space_2d.png"), dpi=150, bbox_inches="tight"
    )
    plt.close()
    print("✅ 2D латентное пространство визуализировано")

# Визуализация с t-SNE (для любой размерности)
print("\n🔬 Визуализация латентного пространства с t-SNE...")
z_tsne = TSNE(n_components=2, random_state=42, perplexity=30).fit_transform(z_mean)

plt.figure(figsize=(12, 10))
scatter = plt.scatter(
    z_tsne[:, 0], z_tsne[:, 1], c=y_test, cmap="tab10", alpha=0.6, s=10
)
plt.colorbar(scatter, label="Класс")
plt.xlabel("t-SNE component 1", fontsize=12)
plt.ylabel("t-SNE component 2", fontsize=12)
plt.title("Латентное пространство VAE (t-SNE проекция)", fontsize=14)
plt.grid(True, alpha=0.3)

# Легенда
for i, name in enumerate(class_names):
    plt.scatter([], [], c=[plt.cm.tab10(i)], label=name, s=50)
plt.legend(loc="center left", bbox_to_anchor=(1.15, 0.5), fontsize=9)

plt.tight_layout()
plt.savefig(
    os.path.join(OUTPUT_DIR, "latent_space_tsne.png"), dpi=150, bbox_inches="tight"
)
plt.close()
print("✅ t-SNE визуализация сохранена")

# Визуализация с PCA
print("\n📐 Визуализация латентного пространства с PCA...")
z_pca = PCA(n_components=2).fit_transform(z_mean)

plt.figure(figsize=(12, 10))
scatter = plt.scatter(z_pca[:, 0], z_pca[:, 1], c=y_test, cmap="tab10", alpha=0.6, s=10)
plt.colorbar(scatter, label="Класс")
plt.xlabel("PCA component 1", fontsize=12)
plt.ylabel("PCA component 2", fontsize=12)
plt.title("Латентное пространство VAE (PCA проекция)", fontsize=14)
plt.grid(True, alpha=0.3)

# Легенда
for i, name in enumerate(class_names):
    plt.scatter([], [], c=[plt.cm.tab10(i)], label=name, s=50)
plt.legend(loc="center left", bbox_to_anchor=(1.15, 0.5), fontsize=9)

plt.tight_layout()
plt.savefig(
    os.path.join(OUTPUT_DIR, "latent_space_pca.png"), dpi=150, bbox_inches="tight"
)
plt.close()
print("✅ PCA визуализация сохранена")

# ============================================================
# Часть 7: Генерация новых изображений
# ============================================================
print("\n" + "=" * 60)
print("✨ ЧАСТЬ 7: Генерация новых изображений")
print("=" * 60)


def plot_latent_space_manifold(decoder, n=20, figsize=15):
    """Визуализация 2D манифолда латентного пространства"""
    digit_size = 28
    scale = 2.0
    figure = np.zeros((digit_size * n, digit_size * n))

    # Создание сетки точек в латентном пространстве
    grid_x = np.linspace(-scale, scale, n)
    grid_y = np.linspace(-scale, scale, n)[::-1]

    for i, yi in enumerate(grid_y):
        for j, xi in enumerate(grid_x):
            z_sample = np.array([[xi, yi]])
            x_decoded = decoder.predict(z_sample, verbose=0)
            digit = x_decoded[0].reshape(digit_size, digit_size)
            figure[
                i * digit_size : (i + 1) * digit_size,
                j * digit_size : (j + 1) * digit_size,
            ] = digit

    plt.figure(figsize=(figsize, figsize))
    start_range = digit_size // 2
    end_range = n * digit_size + start_range
    pixel_range = np.arange(start_range, end_range, digit_size)
    sample_range_x = np.round(grid_x, 1)
    sample_range_y = np.round(grid_y, 1)
    plt.xticks(pixel_range, sample_range_x)
    plt.yticks(pixel_range, sample_range_y)
    plt.xlabel("z[0]", fontsize=14)
    plt.ylabel("z[1]", fontsize=14)
    plt.imshow(figure, cmap="Greys_r")
    plt.title("2D манифолд сгенерированных изображений", fontsize=16)
    plt.tight_layout()
    plt.savefig(
        os.path.join(OUTPUT_DIR, "latent_manifold.png"), dpi=150, bbox_inches="tight"
    )
    plt.close()


if LATENT_DIM == 2:
    print("🎨 Генерация 2D манифолда...")
    plot_latent_space_manifold(vae.decoder, n=20)
    print("✅ 2D манифолд сохранен")

# Генерация случайных изображений
print("\n🎲 Генерация случайных изображений из N(0, I)...")
n_samples = 25
random_latent_vectors = np.random.normal(size=(n_samples, LATENT_DIM))
generated_images = vae.decoder.predict(random_latent_vectors, verbose=0)

plt.figure(figsize=(12, 12))
for i in range(n_samples):
    plt.subplot(5, 5, i + 1)
    plt.imshow(generated_images[i].squeeze(), cmap="gray")
    plt.axis("off")
plt.suptitle("Случайно сгенерированные изображения", fontsize=16)
plt.tight_layout()
plt.savefig(
    os.path.join(OUTPUT_DIR, "random_generated.png"), dpi=150, bbox_inches="tight"
)
plt.close()
print("✅ Случайные изображения сохранены")

# ============================================================
# Часть 8: Интерполяция в латентном пространстве
# ============================================================
print("\n" + "=" * 60)
print("🔀 ЧАСТЬ 8: Интерполяция в латентном пространстве")
print("=" * 60)


def interpolate_images(encoder, decoder, img1, img2, n_steps=10):
    """Интерполяция между двумя изображениями в латентном пространстве"""
    # Получение латентных представлений
    z1_mean, _, _ = encoder.predict(img1[np.newaxis, ...], verbose=0)
    z2_mean, _, _ = encoder.predict(img2[np.newaxis, ...], verbose=0)

    # Линейная интерполяция
    interpolated_z = []
    for alpha in np.linspace(0, 1, n_steps):
        interpolated_z.append((1 - alpha) * z1_mean + alpha * z2_mean)

    interpolated_z = np.vstack(interpolated_z)

    # Декодирование
    interpolated_images = decoder.predict(interpolated_z, verbose=0)

    return interpolated_images


# Примеры интерполяции
print("🎨 Создание интерполяций между разными классами...")

# Выбор примеров из разных классов
interpolation_pairs = [
    (0, 1),  # T-shirt → Trouser
    (2, 3),  # Pullover → Dress
    (5, 7),  # Sandal → Sneaker
]

for idx, (class1, class2) in enumerate(interpolation_pairs):
    # Найти примеры каждого класса
    idx1 = np.where(y_test == class1)[0][0]
    idx2 = np.where(y_test == class2)[0][0]

    img1 = x_test[idx1]
    img2 = x_test[idx2]

    # Интерполяция
    interpolated = interpolate_images(vae.encoder, vae.decoder, img1, img2, n_steps=10)

    # Визуализация
    plt.figure(figsize=(20, 4))
    for i in range(10):
        plt.subplot(1, 10, i + 1)
        plt.imshow(interpolated[i].squeeze(), cmap="gray")
        plt.axis("off")
        if i == 0:
            plt.title(class_names[class1], fontsize=10)
        elif i == 9:
            plt.title(class_names[class2], fontsize=10)

    plt.suptitle(
        f'Интерполяция: "{class_names[class1]}" → "{class_names[class2]}"', fontsize=14
    )
    plt.tight_layout()
    plt.savefig(
        os.path.join(OUTPUT_DIR, f"interpolation_{idx+1}.png"),
        dpi=150,
        bbox_inches="tight",
    )
    plt.close()

print("✅ Интерполяции сохранены")

# ============================================================
# Часть 9: Анализ размерностей латентного пространства
# ============================================================
print("\n" + "=" * 60)
print("🔬 ЧАСТЬ 9: Анализ семантики латентных размерностей")
print("=" * 60)


def analyze_latent_dimensions(decoder, n_samples=10, n_dims=LATENT_DIM):
    """Анализ влияния каждой размерности латентного пространства"""
    if n_dims > 2:
        n_dims = min(n_dims, 10)  # Ограничение для визуализации

    fig, axes = plt.subplots(n_dims, n_samples, figsize=(n_samples * 2, n_dims * 2))
    if n_dims == 1:
        axes = axes[np.newaxis, :]

    for dim in range(n_dims):
        # Варьирование одной размерности
        values = np.linspace(-3, 3, n_samples)

        for i, value in enumerate(values):
            z_sample = np.zeros((1, LATENT_DIM))
            z_sample[0, dim] = value
            generated = decoder.predict(z_sample, verbose=0)

            axes[dim, i].imshow(generated[0].squeeze(), cmap="gray")
            axes[dim, i].axis("off")
            if i == 0:
                axes[dim, i].set_ylabel(f"Dim {dim}", fontsize=10)
            if dim == 0:
                axes[dim, i].set_title(f"{value:.1f}", fontsize=8)

    plt.suptitle("Влияние латентных размерностей на генерацию", fontsize=14)
    plt.tight_layout()
    plt.savefig(
        os.path.join(OUTPUT_DIR, "latent_dimensions_analysis.png"),
        dpi=150,
        bbox_inches="tight",
    )
    plt.close()


print("🎨 Анализ влияния латентных размерностей...")
analyze_latent_dimensions(vae.decoder, n_samples=10)
print("✅ Анализ размерностей сохранен")

# ============================================================
# Часть 10: Сравнение с обычным автоэнкодером
# ============================================================
print("\n" + "=" * 60)
print("⚖️ ЧАСТЬ 10: Сравнение VAE с обычным автоэнкодером")
print("=" * 60)

print("🔨 Создание обычного автоэнкодера для сравнения...")

# Обычный автоэнкодер (без вариационной части)
ae_encoder_inputs = keras.Input(shape=(28, 28, 1), name="ae_encoder_input")
x = layers.Conv2D(32, 3, activation="relu", strides=2, padding="same")(
    ae_encoder_inputs
)
x = layers.Conv2D(64, 3, activation="relu", strides=2, padding="same")(x)
x = layers.Flatten()(x)
x = layers.Dense(16, activation="relu")(x)
ae_latent = layers.Dense(LATENT_DIM, name="ae_latent")(x)
ae_encoder = keras.Model(ae_encoder_inputs, ae_latent, name="ae_encoder")

ae_decoder_inputs = keras.Input(shape=(LATENT_DIM,), name="ae_decoder_input")
x = layers.Dense(7 * 7 * 64, activation="relu")(ae_decoder_inputs)
x = layers.Reshape((7, 7, 64))(x)
x = layers.Conv2DTranspose(64, 3, activation="relu", strides=2, padding="same")(x)
x = layers.Conv2DTranspose(32, 3, activation="relu", strides=2, padding="same")(x)
ae_decoder_outputs = layers.Conv2DTranspose(1, 3, activation="sigmoid", padding="same")(
    x
)
ae_decoder = keras.Model(ae_decoder_inputs, ae_decoder_outputs, name="ae_decoder")

# Полная модель AE
ae_inputs = keras.Input(shape=(28, 28, 1))
ae_latent_output = ae_encoder(ae_inputs)
ae_outputs = ae_decoder(ae_latent_output)
autoencoder = keras.Model(ae_inputs, ae_outputs, name="autoencoder")

autoencoder.compile(optimizer="adam", loss="binary_crossentropy")

print("⏳ Обучение обычного автоэнкодера...")
ae_history = autoencoder.fit(
    x_train,
    x_train,
    epochs=30,
    batch_size=128,
    validation_data=(x_test, x_test),
    verbose=0,
)

print("✅ Обучение автоэнкодера завершено")

# Сравнение реконструкций
print("\n📊 Сравнение реконструкций VAE и AE...")
n_compare = 10
test_samples = x_test[:n_compare]

# VAE реконструкции
z_mean_vae, _, z_vae = vae.encoder.predict(test_samples, verbose=0)
vae_reconstructions = vae.decoder.predict(z_vae, verbose=0)

# AE реконструкции
ae_latent_output = ae_encoder.predict(test_samples, verbose=0)
ae_reconstructions = ae_decoder.predict(ae_latent_output, verbose=0)

plt.figure(figsize=(20, 6))
for i in range(n_compare):
    # Оригинал
    plt.subplot(3, n_compare, i + 1)
    plt.imshow(test_samples[i].squeeze(), cmap="gray")
    if i == 0:
        plt.ylabel("Оригинал", fontsize=10)
    plt.axis("off")

    # VAE реконструкция
    plt.subplot(3, n_compare, n_compare + i + 1)
    plt.imshow(vae_reconstructions[i].squeeze(), cmap="gray")
    if i == 0:
        plt.ylabel("VAE", fontsize=10)
    plt.axis("off")

    # AE реконструкция
    plt.subplot(3, n_compare, 2 * n_compare + i + 1)
    plt.imshow(ae_reconstructions[i].squeeze(), cmap="gray")
    if i == 0:
        plt.ylabel("AE", fontsize=10)
    plt.axis("off")

plt.suptitle("Сравнение: VAE vs обычный Автоэнкодер", fontsize=14)
plt.tight_layout()
plt.savefig(
    os.path.join(OUTPUT_DIR, "vae_vs_ae_comparison.png"), dpi=150, bbox_inches="tight"
)
plt.close()

print("✅ Сравнение сохранено")

# Сравнение латентного пространства
print("\n🔍 Сравнение латентного пространства VAE и AE...")

if LATENT_DIM == 2:
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 8))

    # VAE латентное пространство
    scatter1 = ax1.scatter(
        z_mean[:, 0], z_mean[:, 1], c=y_test, cmap="tab10", alpha=0.6, s=10
    )
    ax1.set_xlabel("z[0]", fontsize=12)
    ax1.set_ylabel("z[1]", fontsize=12)
    ax1.set_title("VAE Латентное пространство", fontsize=14)
    ax1.grid(True, alpha=0.3)

    # AE латентное пространство
    ae_latent_test = ae_encoder.predict(x_test, verbose=0)
    scatter2 = ax2.scatter(
        ae_latent_test[:, 0],
        ae_latent_test[:, 1],
        c=y_test,
        cmap="tab10",
        alpha=0.6,
        s=10,
    )
    ax2.set_xlabel("z[0]", fontsize=12)
    ax2.set_ylabel("z[1]", fontsize=12)
    ax2.set_title("AE Латентное пространство", fontsize=14)
    ax2.grid(True, alpha=0.3)

    plt.colorbar(scatter1, ax=ax1, label="Класс")
    plt.colorbar(scatter2, ax=ax2, label="Класс")

    plt.tight_layout()
    plt.savefig(
        os.path.join(OUTPUT_DIR, "latent_space_comparison.png"),
        dpi=150,
        bbox_inches="tight",
    )
    plt.close()

    print("✅ Сравнение латентного пространства сохранено")

# ============================================================
# Финальный отчет
# ============================================================
print("\n" + "=" * 60)
print("📋 ФИНАЛЬНЫЙ ОТЧЕТ")
print("=" * 60)

# Вычисление финальных метрик
final_vae_loss = history.history["val_loss"][-1]
final_vae_recon = history.history["val_reconstruction_loss"][-1]
final_vae_kl = history.history["val_kl_loss"][-1]
final_ae_loss = ae_history.history["val_loss"][-1]

report_path = os.path.join(OUTPUT_DIR, "lab4_detailed_report.txt")
with open(report_path, "w", encoding="utf-8") as f:
    f.write("=" * 60 + "\n")
    f.write("ЛАБОРАТОРНАЯ РАБОТА 4: ВАРИАЦИОННЫЙ АВТОЭНКОДЕР (VAE)\n")
    f.write("=" * 60 + "\n\n")

    f.write("КОНФИГУРАЦИЯ ЭКСПЕРИМЕНТА:\n")
    f.write("-" * 40 + "\n")
    f.write(f"Датасет: Fashion-MNIST\n")
    f.write(f"Размер изображений: 28x28x1\n")
    f.write(f"Тренировочных образцов: {len(x_train)}\n")
    f.write(f"Тестовых образцов: {len(x_test)}\n")
    f.write(f"Количество классов: 10\n")
    f.write(f"Размерность латентного пространства: {LATENT_DIM}\n")
    f.write(f"Эпох обучения: {EPOCHS}\n")
    f.write(f"Размер батча: {BATCH_SIZE}\n\n")

    f.write("АРХИТЕКТУРА VAE:\n")
    f.write("-" * 40 + "\n")
    f.write("Энкодер:\n")
    f.write("  Conv2D(32) → Conv2D(64) → Flatten → Dense(16)\n")
    f.write(f"  → Dense({LATENT_DIM}) для μ\n")
    f.write(f"  → Dense({LATENT_DIM}) для log(σ²)\n")
    f.write("  → Sampling (reparameterization trick)\n\n")
    f.write("Декодер:\n")
    f.write("  Dense(7×7×64) → Reshape(7,7,64)\n")
    f.write("  → Conv2DTranspose(64) → Conv2DTranspose(32)\n")
    f.write("  → Conv2DTranspose(1, sigmoid)\n\n")

    f.write("РЕЗУЛЬТАТЫ ОБУЧЕНИЯ:\n")
    f.write("-" * 40 + "\n")
    f.write(f"VAE Final Validation Loss: {final_vae_loss:.4f}\n")
    f.write(f"VAE Reconstruction Loss: {final_vae_recon:.4f}\n")
    f.write(f"VAE KL Divergence Loss: {final_vae_kl:.4f}\n")
    f.write(f"AE Final Validation Loss: {final_ae_loss:.4f}\n\n")

    f.write("КЛЮЧЕВЫЕ ВЫВОДЫ:\n")
    f.write("-" * 40 + "\n")
    f.write("1. VAE успешно обучается генерировать изображения Fashion-MNIST\n")
    f.write("2. Латентное пространство организовано семантически\n")
    f.write("3. Reparameterization trick позволяет обучать вероятностную модель\n")
    f.write("4. KL divergence регуляризует латентное пространство\n")
    f.write("5. Интерполяция в латентном пространстве дает плавные переходы\n")
    f.write("6. VAE генерирует более разнообразные образцы чем обычный AE\n\n")

    f.write("СРАВНЕНИЕ VAE И AE:\n")
    f.write("-" * 40 + "\n")
    f.write("VAE:\n")
    f.write("  + Генерация новых образцов из N(0,I)\n")
    f.write("  + Организованное латентное пространство\n")
    f.write("  + Плавная интерполяция\n")
    f.write("  - Немного размытые реконструкции\n\n")
    f.write("AE:\n")
    f.write("  + Более четкие реконструкции\n")
    f.write("  + Проще обучать\n")
    f.write("  - Не гарантирует качественную генерацию\n")
    f.write('  - Латентное пространство может иметь "дыры"\n\n')

    f.write("РЕКОМЕНДАЦИИ ДЛЯ УЛУЧШЕНИЯ:\n")
    f.write("-" * 40 + "\n")
    f.write("1. Увеличить размерность латентного пространства (8-128)\n")
    f.write("2. Использовать β-VAE для лучшего disentanglement\n")
    f.write("3. Добавить больше слоев в энкодер/декодер\n")
    f.write("4. Экспериментировать с другими функциями потерь (perceptual loss)\n")
    f.write("5. Использовать conditional VAE для контролируемой генерации\n")

print(f"\n✅ Детальный отчет сохранен: {report_path}")

# Финальная статистика
print(f"\n📊 СВОДКА РЕЗУЛЬТАТОВ:")
print(f"   • VAE Validation Loss: {final_vae_loss:.4f}")
print(f"   • VAE Reconstruction Loss: {final_vae_recon:.4f}")
print(f"   • VAE KL Divergence: {final_vae_kl:.4f}")
print(f"   • AE Validation Loss: {final_ae_loss:.4f}")

print(f"\n📁 ВСЕ РЕЗУЛЬТАТЫ СОХРАНЕНЫ В: {OUTPUT_DIR}/")
print(f"🖼️ ГРАФИКИ:")
print(f"   • training_history.png")
print(f"   • latent_space_2d.png")
print(f"   • latent_manifold.png")
print(f"   • interpolation_*.png")
print(f"   • vae_vs_ae_comparison.png")
print(f"💾 МОДЕЛИ: vae_encoder.h5, vae_decoder.h5")
print(f"📋 ОТЧЕТ: lab4_detailed_report.txt")

print("\n" + "=" * 60)
print("🎉 ЛАБОРАТОРНАЯ РАБОТА 4 УСПЕШНО ЗАВЕРШЕНА!")
print("=" * 60)