test_vae_visualization.py

"""
학습된 VAE 모델로 원본 컬러 이미지와 복구 이미지를 비교 시각화
"""
import torch
import matplotlib.pyplot as plt
import numpy as np
import os
from pathlib import Path

# Add src to path
import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from models.model import VAE
from data.dataset import load_mnist


def load_trained_model(checkpoint_path, device, latent_dim=64):
    """
    학습된 VAE 모델 로드

    Args:
        checkpoint_path: 모델 체크포인트 경로
        device: torch device
        latent_dim: 잠재 벡터 차원

    Returns:
        model: 로드된 VAE 모델
    """
    model = VAE(in_channels=3, latent_dim=latent_dim).to(device)

    checkpoint = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()

    print(f"✓ Model loaded from: {checkpoint_path}")
    print(f"  Trained for {checkpoint['epoch'] + 1} epochs")

    return model


def visualize_reconstructions(model, test_loader, device, num_samples=8, save_path='./outputs/figures/vae_reconstruction.png'):
    """
    원본 이미지와 복구 이미지를 나란히 표시

    Args:
        model: VAE 모델
        test_loader: 테스트 데이터로더
        device: torch device
        num_samples: 시각화할 샘플 수
        save_path: 저장 경로
    """
    model.eval()

    # 테스트 데이터에서 첫 배치 가져오기
    images, labels = next(iter(test_loader))
    images = images[:num_samples].to(device)

    with torch.no_grad():
        recon_logits, mu, logvar = model(images)
        reconstructions = torch.sigmoid(recon_logits)

    # CPU로 이동 및 numpy 변환
    images_np = images.cpu().numpy()
    reconstructions_np = reconstructions.cpu().numpy()

    # 시각화
    fig, axes = plt.subplots(2, num_samples, figsize=(16, 5))

    for i in range(num_samples):
        # 원본 이미지 - RGB 컬러 표시
        orig_img = images_np[i].transpose(1, 2, 0)  # (C, H, W) -> (H, W, C)
        orig_img = np.clip(orig_img, 0, 1)
        axes[0, i].imshow(orig_img)
        axes[0, i].set_title(f'Original #{i+1}', fontsize=10, fontweight='bold')
        axes[0, i].axis('off')

        # 복구된 이미지 - RGB 컬러 표시
        recon_img = reconstructions_np[i].transpose(1, 2, 0)
        recon_img = np.clip(recon_img, 0, 1)
        axes[1, i].imshow(recon_img)
        axes[1, i].set_title(f'Reconstructed #{i+1}', fontsize=10, fontweight='bold')
        axes[1, i].axis('off')

    # 행 레이블 추가
    fig.text(0.02, 0.75, 'Original', va='center', fontsize=13, fontweight='bold', rotation='vertical')
    fig.text(0.02, 0.25, 'Reconstructed', va='center', fontsize=13, fontweight='bold', rotation='vertical')

    plt.suptitle('VAE Color Image Reconstruction (CIFAR10 64x64)', fontsize=14, fontweight='bold', y=0.98)
    plt.tight_layout()

    # 저장
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"✓ Reconstruction comparison saved to: {save_path}")
    plt.close()


def visualize_side_by_side(model, test_loader, device, num_samples=8, save_path='./outputs/figures/vae_side_by_side.png'):
    """
    원본과 복구 이미지를 쌍으로 나란히 표시 (한 줄에 원본-복구 쌍)

    Args:
        model: VAE 모델
        test_loader: 테스트 데이터로더
        device: torch device
        num_samples: 시각화할 샘플 수
        save_path: 저장 경로
    """
    model.eval()

    # 테스트 데이터에서 첫 배치 가져오기
    images, labels = next(iter(test_loader))
    images = images[:num_samples].to(device)

    with torch.no_grad():
        recon_logits, mu, logvar = model(images)
        reconstructions = torch.sigmoid(recon_logits)

    # CPU로 이동 및 numpy 변환
    images_np = images.cpu().numpy()
    reconstructions_np = reconstructions.cpu().numpy()

    # 시각화 (원본-복구 쌍을 한 행에 표시)
    fig, axes = plt.subplots(num_samples, 2, figsize=(6, 3*num_samples))

    for i in range(num_samples):
        # 원본 이미지
        orig_img = images_np[i].transpose(1, 2, 0)
        orig_img = np.clip(orig_img, 0, 1)
        axes[i, 0].imshow(orig_img)
        axes[i, 0].set_title(f'Original (Sample #{i+1})', fontsize=11, fontweight='bold')
        axes[i, 0].axis('off')

        # 복구된 이미지
        recon_img = reconstructions_np[i].transpose(1, 2, 0)
        recon_img = np.clip(recon_img, 0, 1)
        axes[i, 1].imshow(recon_img)
        axes[i, 1].set_title(f'Reconstructed (Sample #{i+1})', fontsize=11, fontweight='bold')
        axes[i, 1].axis('off')

        # 화살표 추가
        fig.text(0.5, 0.5 - i * (0.9 / num_samples) - 0.04, '→',
                ha='center', va='center', fontsize=20, color='gray')

    plt.suptitle('VAE Reconstruction - Original vs Reconstructed', fontsize=14, fontweight='bold', y=0.995)
    plt.tight_layout()

    # 저장
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"✓ Side-by-side comparison saved to: {save_path}")
    plt.close()


def get_device():
    """사용 가능한 최적의 디바이스 선택"""
    if torch.cuda.is_available():
        device = torch.device('cuda')
        print(f"✓ CUDA available: {torch.cuda.get_device_name(0)}")
    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        device = torch.device('mps')
        print(f"✓ Apple Metal (MPS) available")
    else:
        device = torch.device('cpu')
        print(f"Using CPU")

    return device


def main():
    print("="*70)
    print("VAE COLOR IMAGE RECONSTRUCTION TEST (CIFAR10)")
    print("="*70)

    # 설정
    BATCH_SIZE = 16
    LATENT_DIM = 64
    MODEL_PATH = './outputs/best_models/best_model.pth'

    # 디바이스 선택
    device = get_device()
    print()

    # 모델 로드
    print("Loading trained model...")
    model = load_trained_model(MODEL_PATH, device, LATENT_DIM)
    print()

    # 데이터 로드
    print("Loading test data...")
    _, test_loader = load_mnist(data_dir='./data', batch_size=BATCH_SIZE)
    print(f"✓ Test dataset loaded: {len(test_loader.dataset)} samples")
    print()

    # 시각화
    print("Generating visualizations...")
    print("-" * 70)
    visualize_reconstructions(model, test_loader, device, num_samples=8)
    visualize_side_by_side(model, test_loader, device, num_samples=6)

    print("-" * 70)
    print("="*70)
    print("TEST COMPLETED!")
    print("="*70)


if __name__ == "__main__":
    main()