11 P3 프로젝트: EuroSAT 위성영상 분류

import os
os.environ.setdefault("KERAS_BACKEND", "tensorflow")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")

from pathlib import Path
from copy import deepcopy
import hashlib
import random

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from PIL import Image
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, ConfusionMatrixDisplay

import tensorflow as tf
import keras
from keras import layers
import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

SEED = 42
random.seed(SEED)
np.random.seed(SEED)
keras.utils.set_random_seed(SEED)
torch.manual_seed(SEED)

DATA_ROOT = Path("datasets")
OUTPUT_ROOT = Path("p3_eurosat_outputs")
DATA_ROOT.mkdir(exist_ok=True, parents=True)
OUTPUT_ROOT.mkdir(exist_ok=True, parents=True)

print("Keras:", keras.__version__, "TensorFlow:", tf.__version__)
print("PyTorch:", torch.__version__)
print("TensorFlow GPU:", tf.config.list_physical_devices("GPU"))
print("PyTorch GPU:", torch.cuda.is_available())

RUN_KERAS_BASELINE = False
RUN_KERAS_AUGMENTED = False
RUN_XCEPTION_TRANSFER = False
RUN_XCEPTION_FINE_TUNING = False
RUN_TORCH_BASELINE = False
RUN_TORCH_IMPROVED = False
RUN_FINAL_TEST = False     # 검증 비교가 끝난 뒤에만 True

KERAS_BATCH_SIZE = 64
XCEPTION_BATCH_SIZE = 16
TORCH_BATCH_SIZE = 64
BASELINE_EPOCHS = 5
TRANSFER_EPOCHS = 4
FINE_TUNE_EPOCHS = 2
TORCH_EPOCHS = 5

# 이미 내려받은 파일이 있으면 다시 다운로드하지 않는다.
eurosat = datasets.EuroSAT(root=str(DATA_ROOT), download=True)

class_names = list(eurosat.classes)
image_paths = np.asarray([path for path, _ in eurosat.samples])
labels = np.asarray(eurosat.targets, dtype=np.int64)
num_classes = len(class_names)

print("Total images:", len(eurosat))
print("Number of classes:", num_classes)
print("Classes:", class_names)
print("First image:", image_paths[0])

assert len(eurosat) == 27000, "EuroSAT RGB 전체 27,000장을 읽었는지 확인하세요."
assert num_classes == 10

# 데이터 목록 자체가 바뀌면 예전에 저장한 분할을 잘못 재사용하지 않도록 확인한다.
source_signature = hashlib.sha256(
    "\n".join(str(Path(p).relative_to(DATA_ROOT)) for p in image_paths).encode("utf-8")
).hexdigest()

split_path = OUTPUT_ROOT / "eurosat_split_seed42.npz"
if split_path.exists():
    with np.load(split_path) as saved:
        assert str(saved["signature"].item()) == source_signature, "데이터 순서가 변경되었습니다. 분할 파일을 확인하세요."
        train_idx = saved["train_idx"]
        val_idx = saved["val_idx"]
        test_idx = saved["test_idx"]
else:
    indices = np.arange(len(labels))
    train_idx, remain_idx = train_test_split(
        indices, train_size=0.70, stratify=labels, random_state=SEED
    )
    val_idx, test_idx = train_test_split(
        remain_idx, test_size=0.50, stratify=labels[remain_idx], random_state=SEED
    )
    np.savez(split_path, train_idx=train_idx, val_idx=val_idx,
             test_idx=test_idx, signature=source_signature)

assert len(train_idx) == 18900 and len(val_idx) == 4050 and len(test_idx) == 4050
assert len(np.unique(np.concatenate([train_idx, val_idx, test_idx]))) == len(labels)
assert set(train_idx).isdisjoint(val_idx)
assert set(train_idx).isdisjoint(test_idx)
assert set(val_idx).isdisjoint(test_idx)

print("Train / Validation / Test:", len(train_idx), len(val_idx), len(test_idx))

counts = pd.DataFrame({
    "Train": np.bincount(labels[train_idx], minlength=num_classes),
    "Validation": np.bincount(labels[val_idx], minlength=num_classes),
    "Test": np.bincount(labels[test_idx], minlength=num_classes),
}, index=class_names)
display(counts)

# 각 클래스의 훈련 이미지를 한 장씩 표시한다.
fig, axes = plt.subplots(2, 5, figsize=(12, 5))
for class_id, ax in enumerate(axes.flat):
    image_idx = next(i for i in train_idx if labels[i] == class_id)
    with Image.open(image_paths[image_idx]) as image:
        ax.imshow(image.convert("RGB"))
    ax.set_title(class_names[class_id], fontsize=9)
    ax.axis("off")
plt.tight_layout()

AUTOTUNE = tf.data.AUTOTUNE

def decode_image(path, label):
    content = tf.io.read_file(path)
    image = tf.io.decode_jpeg(content, channels=3)
    image = tf.image.resize(image, [64, 64])
    return tf.cast(image, tf.float32), tf.cast(label, tf.int32)  # 아직 [0,255]


def make_keras_dataset(indices, batch_size, training=False):
    ds = tf.data.Dataset.from_tensor_slices((image_paths[indices], labels[indices]))
    if training:
        ds = ds.shuffle(len(indices), seed=SEED, reshuffle_each_iteration=True)
    return (ds.map(decode_image, num_parallel_calls=AUTOTUNE)
              .batch(batch_size).prefetch(AUTOTUNE))

keras_train_ds = make_keras_dataset(train_idx, KERAS_BATCH_SIZE, training=True)
keras_val_ds = make_keras_dataset(val_idx, KERAS_BATCH_SIZE)
keras_test_ds = make_keras_dataset(test_idx, KERAS_BATCH_SIZE)

example_images, example_labels = next(iter(keras_train_ds))
print("Keras batch:", example_images.shape, example_labels.shape)
print("Pixel range:", float(tf.reduce_min(example_images)), float(tf.reduce_max(example_images)))

plain_transform = transforms.ToTensor()
aug_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.RandomRotation(90),
    transforms.ToTensor(),
])

# 기준 모델과 개선 모델을 구조 중심으로 비교하기 위해 우선 동일한 전처리를 사용한다.
# 추가 실험에서 aug_transform을 적용할 수 있다.
torch_data_plain = datasets.EuroSAT(root=str(DATA_ROOT), download=False,
                                    transform=plain_transform)

train_torch_ds = Subset(torch_data_plain, train_idx.tolist())
val_torch_ds = Subset(torch_data_plain, val_idx.tolist())
test_torch_ds = Subset(torch_data_plain, test_idx.tolist())

pin_memory = torch.cuda.is_available()
train_loader = DataLoader(train_torch_ds, batch_size=TORCH_BATCH_SIZE,
                          shuffle=True, num_workers=0, pin_memory=pin_memory)
val_loader = DataLoader(val_torch_ds, batch_size=TORCH_BATCH_SIZE,
                        shuffle=False, num_workers=0, pin_memory=pin_memory)
test_loader = DataLoader(test_torch_ds, batch_size=TORCH_BATCH_SIZE,
                         shuffle=False, num_workers=0, pin_memory=pin_memory)

torch_x, torch_y = next(iter(train_loader))
print("PyTorch batch:", tuple(torch_x.shape), tuple(torch_y.shape))
print("Pixel range:", torch_x.min().item(), torch_x.max().item())

def build_keras_cnn(use_augmentation=False):
    inputs = keras.Input(shape=(64, 64, 3))
    x = inputs
    if use_augmentation:
        x = layers.RandomFlip("horizontal_and_vertical")(x)
        x = layers.RandomRotation(0.25, fill_mode="reflect")(x)
    x = layers.Rescaling(1.0 / 255.0)(x)
    x = layers.Conv2D(32, 3, padding="same", activation="relu")(x)
    x = layers.MaxPooling2D(2)(x)
    x = layers.Conv2D(64, 3, padding="same", activation="relu")(x)
    x = layers.MaxPooling2D(2)(x)
    x = layers.Conv2D(128, 3, padding="same", activation="relu")(x)
    x = layers.GlobalAveragePooling2D()(x)
    outputs = layers.Dense(num_classes, activation="softmax")(x)
    model = keras.Model(inputs, outputs)
    model.compile(optimizer=keras.optimizers.Adam(1e-3),
                  loss="sparse_categorical_crossentropy", metrics=["accuracy"])
    return model

keras_base = build_keras_cnn(use_augmentation=False)
keras_base.summary()

# 훈련된 Keras 모델의 검증 성능과 체크포인트 경로를 보관한다.
keras_records = {}

def train_keras_model(model, name, train_ds, val_ds, epochs, batch_size_note):
    model_path = OUTPUT_ROOT / (name + ".keras")
    callbacks = [
        keras.callbacks.EarlyStopping(monitor="val_loss", patience=2,
                                      restore_best_weights=True),
        keras.callbacks.ModelCheckpoint(str(model_path), monitor="val_loss",
                                        save_best_only=True),
    ]
    history = model.fit(train_ds, validation_data=val_ds, epochs=epochs,
                        callbacks=callbacks, verbose=1)
    best_epoch = int(np.argmin(history.history["val_loss"]))
    keras_records[name] = {
        "val_loss": float(history.history["val_loss"][best_epoch]),
        "val_accuracy": float(history.history["val_accuracy"][best_epoch]),
        "best_epoch": best_epoch + 1,
        "path": str(model_path),
        "history": history.history,
        "batch_size": batch_size_note,
    }
    return history

if RUN_KERAS_BASELINE:
    keras.utils.set_random_seed(SEED)
    keras_base = build_keras_cnn(False)
    keras_base_history = train_keras_model(
        keras_base, "keras_baseline", keras_train_ds, keras_val_ds,
        BASELINE_EPOCHS, KERAS_BATCH_SIZE
    )
else:
    print("Keras 기준 모델을 훈련하려면 RUN_KERAS_BASELINE=True")

# 실제 증식 이미지와 원본 이미지를 비교한다.
example_path = image_paths[train_idx[0]]
with Image.open(example_path) as img:
    original = np.asarray(img.convert("RGB"), dtype=np.float32)

visual_augmentation = keras.Sequential([
    layers.RandomFlip("horizontal_and_vertical"),
    layers.RandomRotation(0.25, fill_mode="reflect"),
])
fig, axes = plt.subplots(2, 5, figsize=(12, 5))
axes = axes.flat
axes[0].imshow(original.astype("uint8"))
axes[0].set_title("Original")
for i, ax in enumerate(list(axes)[1:], 1):
    augmented = visual_augmentation(original[None, ...], training=True).numpy()[0]
    ax.imshow(np.clip(augmented, 0, 255).astype("uint8"))
    ax.set_title(f"Augmented {i}")
for ax in fig.axes:
    ax.axis("off")
plt.tight_layout()

if RUN_KERAS_AUGMENTED:
    keras.utils.set_random_seed(SEED)
    keras_augmented = build_keras_cnn(use_augmentation=True)
    keras_aug_history = train_keras_model(
        keras_augmented, "keras_augmented", keras_train_ds, keras_val_ds,
        BASELINE_EPOCHS, KERAS_BATCH_SIZE
    )
else:
    print("Keras 증식 모델을 훈련하려면 RUN_KERAS_AUGMENTED=True")

def build_xception_transfer():
    base = keras.applications.Xception(
        include_top=False, weights="imagenet", input_shape=(96, 96, 3)
    )
    base.trainable = False

    inputs = keras.Input(shape=(64, 64, 3))
    x = layers.RandomFlip("horizontal_and_vertical")(inputs)
    x = layers.RandomRotation(0.25, fill_mode="reflect")(x)
    x = layers.Resizing(96, 96)(x)
    x = layers.Rescaling(1 / 127.5, offset=-1)(x)
    x = base(x, training=False)
    x = layers.GlobalAveragePooling2D()(x)
    x = layers.Dropout(0.2)(x)
    outputs = layers.Dense(num_classes, activation="softmax")(x)
    model = keras.Model(inputs, outputs, name="eurosat_xception")
    model.compile(optimizer=keras.optimizers.Adam(1e-3),
                  loss="sparse_categorical_crossentropy", metrics=["accuracy"])
    return model, base

xception_model, xception_base = None, None
if RUN_XCEPTION_TRANSFER:
    keras.utils.set_random_seed(SEED)
    xception_model, xception_base = build_xception_transfer()
    xception_model.summary()
    xception_train_ds = make_keras_dataset(train_idx, XCEPTION_BATCH_SIZE, training=True)
    xception_val_ds = make_keras_dataset(val_idx, XCEPTION_BATCH_SIZE)
    xception_history = train_keras_model(
        xception_model, "xception_transfer", xception_train_ds, xception_val_ds,
        TRANSFER_EPOCHS, XCEPTION_BATCH_SIZE
    )
else:
    print("전이 학습을 시작하려면 RUN_XCEPTION_TRANSFER=True")

if RUN_XCEPTION_FINE_TUNING:
    assert RUN_XCEPTION_TRANSFER and xception_model is not None, (
        "전이 학습 셀을 먼저 실행하세요."
    )
    xception_base.trainable = True
    for layer in xception_base.layers[:-20]:
        layer.trainable = False
    for layer in xception_base.layers:
        if isinstance(layer, layers.BatchNormalization):
            layer.trainable = False

    print("Trainable weights in base:", len(xception_base.trainable_weights))
    xception_model.compile(optimizer=keras.optimizers.Adam(1e-5),
                           loss="sparse_categorical_crossentropy",
                           metrics=["accuracy"])
    finetune_history = train_keras_model(
        xception_model, "xception_finetuned", xception_train_ds, xception_val_ds,
        FINE_TUNE_EPOCHS, XCEPTION_BATCH_SIZE
    )
else:
    print("미세 조정은 RUN_XCEPTION_FINE_TUNING=True로 변경한 뒤 실행하세요.")

def plot_training_history(history, title):
    n = len(history["loss"])
    epochs = np.arange(1, n + 1)
    fig, axes = plt.subplots(1, 2, figsize=(10, 3.5))
    for key, ax, label in [("loss", axes[0], "Loss"),
                           ("accuracy", axes[1], "Accuracy")]:
        ax.plot(epochs, history[key], marker="o", label="Train")
        ax.plot(epochs, history["val_" + key], marker="o", label="Validation")
        ax.set(xlabel="Epoch", ylabel=label, title=f"{title}: {label}")
        ax.set_xticks(np.arange(1, n + 1, max(1, int(np.ceil(n / 20)))))
        ax.legend()
    plt.tight_layout()

if keras_records:
    rows = [{"Model": name, "Best val loss": rec["val_loss"],
             "Val accuracy": rec["val_accuracy"], "Best epoch": rec["best_epoch"]}
            for name, rec in keras_records.items()]
    display(pd.DataFrame(rows).sort_values("Best val loss"))
    for name, rec in keras_records.items():
        plot_training_history(rec["history"], name)
else:
    print("Keras 모델을 훈련한 뒤 검증 결과를 비교하세요.")

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class TorchBaselineCNN(nn.Module):
    def __init__(self, n_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Linear(128, n_classes)

    def forward(self, x):
        x = self.features(x).flatten(1)
        return self.classifier(x)  # logits: (batch, 10)

torch_base = TorchBaselineCNN(num_classes).to(DEVICE)
with torch.no_grad():
    print("Output shape:", tuple(torch_base(torch.zeros(2, 3, 64, 64, device=DEVICE)).shape))
print("Parameters:", sum(p.numel() for p in torch_base.parameters()))

def run_torch_epoch(model, loader, optimizer=None):
    training = optimizer is not None
    model.train(training)
    loss_fn = nn.CrossEntropyLoss()
    total_loss, correct, total = 0.0, 0, 0

    with torch.set_grad_enabled(training):
        for xb, yb in loader:
            xb, yb = xb.to(DEVICE), yb.to(DEVICE)
            if training:
                optimizer.zero_grad()
            logits = model(xb)
            loss = loss_fn(logits, yb)
            if training:
                loss.backward()
                optimizer.step()
            total_loss += loss.item() * len(yb)
            correct += (logits.argmax(dim=1) == yb).sum().item()
            total += len(yb)
    return total_loss / total, correct / total


torch_records = {}

def train_torch_model(model, name, epochs=5, learning_rate=1e-3):
    model = model.to(DEVICE)
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    history = {"loss": [], "val_loss": [], "accuracy": [], "val_accuracy": []}
    best_val_loss = float("inf")
    best_state = None
    best_epoch = 0
    for epoch in range(epochs):
        train_loss, train_acc = run_torch_epoch(model, train_loader, optimizer)
        val_loss, val_acc = run_torch_epoch(model, val_loader)
        for key, value in [("loss", train_loss), ("accuracy", train_acc),
                           ("val_loss", val_loss), ("val_accuracy", val_acc)]:
            history[key].append(float(value))
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_state = deepcopy(model.state_dict())
            best_epoch = epoch + 1
        print(f"Epoch {epoch + 1}/{epochs} | train={train_loss:.3f}, {train_acc:.3f} | "
              f"val={val_loss:.3f}, {val_acc:.3f}")
    model.load_state_dict(best_state)
    save_path = OUTPUT_ROOT / f"{name}.pt"
    torch.save(best_state, save_path)
    torch_records[name] = {"model": model, "path": str(save_path),
                           "val_loss": best_val_loss,
                           "val_accuracy": history["val_accuracy"][best_epoch - 1],
                           "best_epoch": best_epoch, "history": history}
    return model, history

if RUN_TORCH_BASELINE:
    torch.manual_seed(SEED)
    torch_base = TorchBaselineCNN(num_classes)
    torch_base, torch_base_history = train_torch_model(
        torch_base, "torch_baseline", TORCH_EPOCHS
    )
else:
    print("PyTorch 기준 모델을 훈련하려면 RUN_TORCH_BASELINE=True")

class DepthwiseSeparableConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.depthwise = nn.Conv2d(in_channels, in_channels, 3, padding=1,
                                   groups=in_channels, bias=False)
        self.pointwise = nn.Conv2d(in_channels, out_channels, 1, bias=False)

    def forward(self, x):
        return self.pointwise(self.depthwise(x))


class MiniXceptionBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        self.main = nn.Sequential(
            nn.ReLU(),
            DepthwiseSeparableConv(in_channels, out_channels),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            DepthwiseSeparableConv(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
            nn.MaxPool2d(3, stride=2, padding=1) if stride == 2 else nn.Identity(),
        )
        self.skip = (nn.Identity() if stride == 1 and in_channels == out_channels
                     else nn.Sequential(nn.Conv2d(in_channels, out_channels, 1,
                                                  stride=stride, bias=False),
                                        nn.BatchNorm2d(out_channels)))

    def forward(self, x):
        return self.main(x) + self.skip(x)


class MiniXception(nn.Module):
    def __init__(self, n_classes=10):
        super().__init__()
        self.stem = nn.Sequential(nn.Conv2d(3, 24, 3, padding=1, bias=False),
                                  nn.BatchNorm2d(24), nn.ReLU())
        self.entry = MiniXceptionBlock(24, 32, stride=2)
        self.middle = MiniXceptionBlock(32, 32)
        self.exit = MiniXceptionBlock(32, 64, stride=2)
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.head = nn.Linear(64, n_classes)

    def forward(self, x):
        x = self.stem(x)
        x = self.entry(x)
        x = self.middle(x)
        x = self.exit(x)
        return self.head(self.pool(x).flatten(1))

mini_xception = MiniXception(num_classes).to(DEVICE)
with torch.no_grad():
    print("Output shape:", tuple(mini_xception.eval()(torch.zeros(2, 3, 64, 64, device=DEVICE)).shape))
print("Parameters:", sum(p.numel() for p in mini_xception.parameters()))

if RUN_TORCH_IMPROVED:
    torch.manual_seed(SEED)
    mini_xception = MiniXception(num_classes)
    mini_xception, mini_history = train_torch_model(
        mini_xception, "torch_mini_xception", TORCH_EPOCHS
    )
else:
    print("PyTorch 개선 모델을 훈련하려면 RUN_TORCH_IMPROVED=True")

if torch_records:
    rows = [{"Model": name, "Best val loss": rec["val_loss"],
             "Val accuracy": rec["val_accuracy"], "Best epoch": rec["best_epoch"]}
            for name, rec in torch_records.items()]
    display(pd.DataFrame(rows).sort_values("Best val loss"))
    for name, rec in torch_records.items():
        plot_training_history(rec["history"], name)
else:
    print("PyTorch 모델을 훈련한 뒤 검증 결과를 비교하세요.")

# TODO: 직접 선택한 변수를 하나만 변경하여 실험한다.
#
# Keras 예: 기준 모델의 필터 수만 변경해 새 모델 정의·훈련
#   train_keras_model(my_model, "keras_my_experiment", ...)
#
# PyTorch 예: 같은 구조를 새로 초기화하고 학습률만 변경
#   train_torch_model(my_model, "torch_my_experiment", learning_rate=...)
#
# 테스트 데이터는 이 실험에 사용하지 않는다.

selected_keras_name = min(keras_records, key=lambda x: keras_records[x]["val_loss"]) if keras_records else None
selected_torch_name = min(torch_records, key=lambda x: torch_records[x]["val_loss"]) if torch_records else None

print("Selected Keras model:", selected_keras_name)
print("Selected PyTorch model:", selected_torch_name)

# 두 프레임워크 모두 프로젝트의 필수 실험을 마쳤는지 먼저 확인한다.
if not keras_records or not torch_records:
    print("일부 모델이 아직 훈련되지 않았습니다. 최종 테스트 전에 필수 비교를 완료하세요.")

keras_test_predictions = None
torch_test_predictions = None
keras_test_metrics = None
torch_test_metrics = None

def predict_torch(model, loader):
    model.eval()
    all_predictions = []
    with torch.no_grad():
        for xb, _ in loader:
            logits = model(xb.to(DEVICE))
            all_predictions.append(logits.argmax(dim=1).cpu().numpy())
    return np.concatenate(all_predictions)

if RUN_FINAL_TEST:
    assert selected_keras_name is not None and selected_torch_name is not None, (
        "먼저 두 프레임워크의 모델을 훈련하고 검증 모델을 선택하세요."
    )
    keras_best = keras.models.load_model(keras_records[selected_keras_name]["path"])
    keras_test_metrics = keras_best.evaluate(keras_test_ds, verbose=0, return_dict=True)
    keras_test_probabilities = keras_best.predict(keras_test_ds, verbose=0)
    keras_test_predictions = keras_test_probabilities.argmax(axis=1)

    torch_best = torch_records[selected_torch_name]["model"].to(DEVICE)
    torch_test_predictions = predict_torch(torch_best, test_loader)
    torch_test_loss, torch_test_accuracy = run_torch_epoch(torch_best, test_loader)
    torch_test_metrics = {"loss": torch_test_loss, "accuracy": torch_test_accuracy}

    print("Keras test:", keras_test_metrics)
    print("PyTorch test:", torch_test_metrics)
else:
    print("모델 선택을 확정했으면 RUN_FINAL_TEST=True로 바꾼 뒤 실행하세요.")

def analyze_predictions(y_true, y_pred, title):
    print(title)
    print("Accuracy:", round(accuracy_score(y_true, y_pred), 4))
    print(classification_report(y_true, y_pred, target_names=class_names,
                                zero_division=0, digits=3))
    cm = confusion_matrix(y_true, y_pred, labels=np.arange(num_classes))
    fig, ax = plt.subplots(figsize=(9, 8))
    ConfusionMatrixDisplay(cm, display_labels=class_names).plot(
        ax=ax, xticks_rotation=55, colorbar=False, values_format="d"
    )
    ax.set_title(title)
    plt.tight_layout()

if keras_test_predictions is not None:
    analyze_predictions(labels[test_idx], keras_test_predictions, "Keras test confusion matrix")
if torch_test_predictions is not None:
    analyze_predictions(labels[test_idx], torch_test_predictions, "PyTorch test confusion matrix")

# TODO: 제출할 모델의 잘못 분류한 이미지 몇 장을 확인하고 해석한다.
# Keras와 PyTorch 중 하나 또는 둘 다 사용해도 된다.

if keras_test_predictions is not None:
    final_predictions = keras_test_predictions
elif torch_test_predictions is not None:
    final_predictions = torch_test_predictions
else:
    final_predictions = None

if final_predictions is not None:
    true_labels = labels[test_idx]
    wrong_pos = np.flatnonzero(final_predictions != true_labels)
    print("Misclassified images:", len(wrong_pos))
    fig, axes = plt.subplots(2, 5, figsize=(12, 5))
    for ax, pos in zip(axes.flat, wrong_pos[:10]):
        with Image.open(image_paths[test_idx[pos]]) as img:
            ax.imshow(img.convert("RGB"))
        ax.set_title(f"True: {class_names[true_labels[pos]]}\nPred: {class_names[final_predictions[pos]]}",
                     fontsize=8)
    for ax in axes.flat:
        ax.axis("off")
    plt.tight_layout()
else:
    print("최종 테스트를 수행한 뒤 오분류 이미지를 확인하세요.")