datasetlib.py

import os
import random
import numpy as np
import pandas as pd
import scipy.io
import matplotlib.pyplot as plt
import tensorflow as tf
import cv2



def unpickle(file):
    """CIFAR-10 데이터를 불러오는 함수"""
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict


def load_mnist(data_dir):
    """MNIST 데이터 로드"""
    train_data = pd.read_csv(os.path.join(data_dir, 'mnist_train.csv')).values
    test_data = pd.read_csv(os.path.join(data_dir, 'mnist_test.csv')).values

    x_train, y_train = train_data[:, 1:].reshape(-1, 28, 28, 1) / 255.0, train_data[:, 0]
    x_test, y_test = test_data[:, 1:].reshape(-1, 28, 28, 1) / 255.0, test_data[:, 0]
    return (x_train, y_train), (x_test, y_test)


def load_fmnist(data_dir):
    """Fashion-MNIST 데이터 로드"""
    train_data = pd.read_csv(os.path.join(data_dir, 'fashion-mnist_train.csv')).values
    test_data = pd.read_csv(os.path.join(data_dir, 'fashion-mnist_test.csv')).values

    x_train, y_train = train_data[:, 1:].reshape(-1, 28, 28, 1) / 255.0, train_data[:, 0]
    x_test, y_test = test_data[:, 1:].reshape(-1, 28, 28, 1) / 255.0, test_data[:, 0]
    return (x_train, y_train), (x_test, y_test)


def load_cifar10(data_dir):
    """CIFAR-10 데이터 로드"""
    x_train, y_train = [], []
    for i in range(1, 6):
        batch = unpickle(os.path.join(data_dir, f'data_batch_{i}'))
        x_train.append(batch[b'data'].reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1) / 255.0)
        y_train.append(np.array(batch[b'labels']))
    x_train = np.concatenate(x_train)
    y_train = np.concatenate(y_train)

    test_batch = unpickle(os.path.join(data_dir, 'test_batch'))
    x_test = test_batch[b'data'].reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1) / 255.0
    y_test = np.array(test_batch[b'labels'])
    return (x_train, y_train), (x_test, y_test)


def load_svhn(data_dir):
    """SVHN 데이터 로드"""
    train_data = scipy.io.loadmat(os.path.join(data_dir, 'SVHN_train_32x32.mat'))
    test_data = scipy.io.loadmat(os.path.join(data_dir, 'SVHN_test_32x32.mat'))

    x_train = train_data['X'].transpose(3, 0, 1, 2) / 255.0
    y_train = train_data['y'].flatten() - 1

    x_test = test_data['X'].transpose(3, 0, 1, 2) / 255.0
    y_test = test_data['y'].flatten() - 1
    return (x_train, y_train), (x_test, y_test)


# 이미지넷 데이터 로드는 검증하지 않았음. 2025.02.13
def load_imagenet(data_dir):
    """
    ImageNet 데이터셋을 로드하고, 학습(train) 및 테스트(test) 데이터로 분리한 후 그룹을 나눕니다.

    Args:
        data_dir (str): ImageNet 데이터가 저장된 경로 (예: './data/ImageNet')

    Returns:
        dict: 
            {
                "class_dict": {클래스명: 인덱스}, 
                "group1_train": [(파일경로, 클래스명), ...],
                "group1_test": [(파일경로, 클래스명), ...],
                "group2_train": [(파일경로, 클래스명), ...],
                "group2_test": [(파일경로, 클래스명), ...]
            }
    """
    train_dir = os.path.join(data_dir, "train")
    test_dir = os.path.join(data_dir, "val")

    # 클래스 리스트 및 매핑 생성
    classes = sorted(os.listdir(train_dir))
    class_dict = {c: i for i, c in enumerate(classes)}

    # 랜덤하게 그룹 1과 그룹 2로 나누기
    group1 = random.sample(classes, len(classes) // 2)
    group2 = [c for c in classes if c not in group1]

    # 데이터 분리 함수
    def split_data(class_list, data_dir):
        """
        주어진 클래스 목록에서 학습 또는 테스트 데이터를 로드하는 내부 함수.
        """
        image_label_pairs = []
        for c in class_list:
            image_path = os.path.join(data_dir, c)
            image_files = sorted(os.listdir(image_path))

            for img in image_files:
                image_label_pairs.append((os.path.join(image_path, img), c))

        return image_label_pairs

    # 그룹별 학습 및 테스트 데이터 생성
    group1_train = split_data(group1, train_dir)
    group1_test = split_data(group1, test_dir)
    group2_train = split_data(group2, train_dir)
    group2_test = split_data(group2, test_dir)

    return {
        "class_dict": class_dict,
        "group1_train": group1_train,
        "group1_test": group1_test,
        "group2_train": group2_train,
        "group2_test": group2_test,
    }


def load_dataset(dataset_name):
    """데이터셋 로드 함수"""
    data_dir = f'./data/{dataset_name}'
    if dataset_name == 'MNIST':
        return load_mnist(data_dir)
    elif dataset_name == 'FMNIST':
        return load_fmnist(data_dir)
    elif dataset_name == 'CIFAR10':
        return load_cifar10(data_dir)
    elif dataset_name == 'SVHN':
        return load_svhn(data_dir)
    elif dataset_name == 'ImageNet':
        return load_imagenet(data_dir)
    else:
        raise ValueError("지원하지 않는 데이터셋입니다.")


def create_datasets(images, labels, batch_size=32):
    return tf.data.Dataset.from_tensor_slices((images, labels)).shuffle(100000).batch(batch_size)


def show_sample_images(images, labels, dataset_name, num_samples=5):
    """데이터 샘플 시각화"""
    plt.figure(figsize=(10, 2))

    for i in range(num_samples):
        plt.subplot(1, num_samples, i + 1)
        if images.shape[-1] == 1:  # 흑백 이미지 (MNIST, FMNIST)
            plt.imshow(images[i].squeeze(), cmap='gray')
        else:  # 컬러 이미지 (CIFAR10, SVHN)
            plt.imshow(images[i])
        plt.title(f"Label: {labels[i]}")
        plt.axis('off')

    plt.suptitle(f"{dataset_name} Samples")
    plt.show()


def load_datasets_all():
    """데이터 로드 및 샘플 시각화 테스트"""
    datasets_names = ['MNIST', 'FMNIST', 'CIFAR10', 'SVHN']
    datasets = {}
    for dataset_name in datasets_names:
        a_dataset = {}
        (train_images, train_labels), (test_images, test_labels) = load_dataset(dataset_name)
        a_dataset['train_x'] = train_images
        a_dataset['train_y'] = train_labels
        a_dataset['test_x'] = test_images
        a_dataset['test_y'] = test_labels
        datasets[dataset_name] = a_dataset
    return datasets


def test_load_dataset():
    """데이터 로드 및 샘플 시각화 테스트"""
    datasets = ['MNIST', 'FMNIST', 'CIFAR10', 'SVHN']
    for dataset in datasets:
        print(f"🔍 Testing {dataset} dataset...")

        try:
            (train_images, train_labels), (test_images, test_labels) = load_dataset(dataset)

            # 데이터셋 기본 정보 출력
            print(f"✅ {dataset} 데이터 로드 성공!")
            print(f"  - 학습 데이터 크기: {train_images.shape}, 레이블 크기: {train_labels.shape}")
            print(f"  - 테스트 데이터 크기: {test_images.shape}, 레이블 크기: {test_labels.shape}")

            # 데이터 타입 및 범위 확인
            assert train_images.dtype in [np.float32, np.float64], "❌ 이미지 데이터 타입 오류"
            assert train_labels.dtype in [np.int32, np.int64, np.uint8], "❌ 라벨 데이터 타입 오류"
            assert 0 <= train_images.min() <= 1 and 0 <= train_images.max() <= 1, "❌ 이미지 데이터 범위 오류 (0~1 사이여야 함)"
            assert 0 <= test_images.min() <= 1 and 0 <= test_images.max() <= 1, "❌ 이미지 데이터 범위 오류 (0~1 사이여야 함)"

            print(f"✅ {dataset} 데이터셋 테스트 통과!\n")

            # 샘플 이미지 출력
            show_sample_images(train_images[:5], train_labels[:5], dataset)

        except Exception as e:
            print(f"❌ {dataset} 데이터 로드 실패: {str(e)}\n")



def split_train_validation(datasets, val_ratio=0.2):
    """
    Train 데이터를 80%만 유지하고, 20%를 validation 데이터로 분리하여 datasets에 추가.
    
    Args:
        datasets (dict): MNIST, FMNIST, CIFAR10, SVHN 데이터셋이 포함된 딕셔너리.
        val_ratio (float): validation 데이터 비율 (default: 0.2)
    
    Returns:
        dict: validation 데이터가 추가된 datasets
    """
    for dataset_name, dataset in datasets.items():
        train_x, train_y = dataset['train_x'], dataset['train_y']
        
        # 데이터 섞기 (Shuffle)
        indices = np.arange(len(train_x))
        np.random.shuffle(indices)
        
        # Validation 개수 계산
        val_size = int(len(train_x) * val_ratio)

        # Validation 데이터 분리
        val_x, val_y = train_x[:val_size], train_y[:val_size]

        # Train 데이터 업데이트 (80%만 유지)
        train_x, train_y = train_x[val_size:], train_y[val_size:]

        # 업데이트된 데이터 저장
        dataset['train_x'], dataset['train_y'] = train_x, train_y
        dataset['val_x'], dataset['val_y'] = val_x, val_y

        print(f"{dataset_name}: Train -> {len(train_x)}, Validation -> {len(val_x)}")

    return datasets


# 주어진 데이터셋을 resize한다
def resize_datasets(datasets, image_size_to_resize= (32,32) ):
    datasets_resized = {}
    for idx, (dataset_name, dataset) in enumerate(datasets.items()):
        if dataset['train_x'].shape[1] < image_size_to_resize[0]:
            a_dataset_resized = {}
            a_dataset_resized['train_x'] = np.array([cv2.resize(img, image_size_to_resize, interpolation=cv2.INTER_LINEAR) for img in dataset['train_x']])
            a_dataset_resized['val_x'] = np.array([cv2.resize(img, image_size_to_resize, interpolation=cv2.INTER_LINEAR) for img in dataset['val_x']])
            a_dataset_resized['test_x'] = np.array([cv2.resize(img, image_size_to_resize, interpolation=cv2.INTER_LINEAR) for img in dataset['test_x']])
            
            if len(a_dataset_resized['train_x'].shape) < 4:
                a_dataset_resized['train_x'] = np.expand_dims(a_dataset_resized['train_x'],axis=-1)
                a_dataset_resized['val_x'] = np.expand_dims(a_dataset_resized['val_x'],axis=-1)
                a_dataset_resized['test_x'] = np.expand_dims(a_dataset_resized['test_x'],axis=-1)         
                
            a_dataset_resized['train_y'] = dataset['train_y']
            a_dataset_resized['val_y'] = dataset['val_y']
            a_dataset_resized['test_y'] = dataset['test_y']
        else:
            a_dataset_resized = dataset
        datasets_resized[dataset_name] = a_dataset_resized
    return datasets_resized


# 주어진 데이터셋을 하나로 병합한다
def merge_datasets(datasets, image_size_to_resize = None):
    # image_size_to_resize = (32,32)
    if image_size_to_resize is not None:
        datasets_resized = {}
        for idx, (dataset_name, dataset) in enumerate(datasets.items()):
            if dataset['train_x'].shape[1] < image_size_to_resize[0]:
                a_dataset_resized = {}
                a_dataset_resized['train_x'] = np.array([cv2.resize(img, image_size_to_resize, interpolation=cv2.INTER_LINEAR) for img in dataset['train_x']])
                a_dataset_resized['val_x'] = np.array([cv2.resize(img, image_size_to_resize, interpolation=cv2.INTER_LINEAR) for img in dataset['val_x']])
                a_dataset_resized['test_x'] = np.array([cv2.resize(img, image_size_to_resize, interpolation=cv2.INTER_LINEAR) for img in dataset['test_x']])
                
                if len(a_dataset_resized['train_x'].shape) < 4:
                    a_dataset_resized['train_x'] = np.expand_dims(a_dataset_resized['train_x'],axis=-1)
                    a_dataset_resized['val_x'] = np.expand_dims(a_dataset_resized['val_x'],axis=-1)
                    a_dataset_resized['test_x'] = np.expand_dims(a_dataset_resized['test_x'],axis=-1)


                    a_dataset_resized['train_x'] = np.repeat(a_dataset_resized['train_x'], repeats=3, axis=-1)
                    a_dataset_resized['val_x'] = np.repeat(a_dataset_resized['val_x'], repeats=3, axis=-1)
                    a_dataset_resized['test_x'] = np.repeat(a_dataset_resized['test_x'], repeats=3, axis=-1)
                a_dataset_resized['train_y'] = dataset['train_y']
                a_dataset_resized['val_y'] = dataset['val_y']
                a_dataset_resized['test_y'] = dataset['test_y']
            else:
                a_dataset_resized = dataset
            datasets_resized[dataset_name] = a_dataset_resized

        datasets = datasets_resized


    dataset_merged = {}
    dataset_merged['merged'] = {}
    cnt_y_idx = 0
    db_name_list = []
    for idx, (dataset_name, dataset) in enumerate(datasets.items()):
        db_name_list.append((dataset_name,cnt_y_idx))
        if idx==0:
            dataset_merged['merged']['train_x'] = dataset['train_x']
            dataset_merged['merged']['train_y'] = dataset['train_y'] + cnt_y_idx
            dataset_merged['merged']['val_x'] = dataset['val_x']
            dataset_merged['merged']['val_y'] = dataset['val_y'] + cnt_y_idx
            dataset_merged['merged']['test_x'] = dataset['test_x']
            dataset_merged['merged']['test_y'] = dataset['test_y'] + cnt_y_idx
        else:
            dataset_merged['merged']['train_x'] = np.concatenate((dataset_merged['merged']['train_x'], dataset['train_x']),axis=0)
            dataset_merged['merged']['train_y'] = np.concatenate((dataset_merged['merged']['train_y'], dataset['train_y'] + cnt_y_idx),axis=0)
            dataset_merged['merged']['val_x'] = np.concatenate((dataset_merged['merged']['val_x'], dataset['val_x']),axis=0)
            dataset_merged['merged']['val_y'] = np.concatenate((dataset_merged['merged']['val_y'], dataset['val_y'] + cnt_y_idx),axis=0)
            dataset_merged['merged']['test_x'] = np.concatenate((dataset_merged['merged']['test_x'], dataset['test_x']),axis=0)
            dataset_merged['merged']['test_y'] = np.concatenate((dataset_merged['merged']['test_y'], dataset['test_y'] + cnt_y_idx),axis=0)
        
        cnt_y_idx += (np.max(dataset['train_y']) +1)
    
    return dataset_merged, db_name_list