어텐션이 없는 비전 트랜스포머

!pip install -qq -U tensorflow-addons

import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow_addons as tfa

import pathlib
import glob

# 재현성을 위한 시드 설정
SEED = 42
keras.utils.set_random_seed(SEED)

class Config(object):
    # 데이터
    batch_size = 256
    buffer_size = batch_size * 2
    input_shape = (32, 32, 3)
    num_classes = 10

    # 보강
    image_size = 48

    # 아키텍쳐
    patch_size = 4
    projected_dim = 96
    num_shift_blocks_per_stages = [2, 4, 8, 2]
    epsilon = 1e-5
    stochastic_depth_rate = 0.2
    mlp_dropout_rate = 0.2
    num_div = 12
    shift_pixel = 1
    mlp_expand_ratio = 2

    # 옵티마이저
    lr_start = 1e-5
    lr_max = 1e-3
    weight_decay = 1e-4

    # 트레이닝
    epochs = 100

    # 추론
    label_map = {
        0: "airplane",
        1: "automobile",
        2: "bird",
        3: "cat",
        4: "deer",
        5: "dog",
        6: "frog",
        7: "horse",
        8: "ship",
        9: "truck",
    }
    tf_ds_batch_size = 20


config = Config()

(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
(x_train, y_train), (x_val, y_val) = (
    (x_train[:40000], y_train[:40000]),
    (x_train[40000:], y_train[40000:]),
)
print(f"Training samples: {len(x_train)}")
print(f"Validation samples: {len(x_val)}")
print(f"Testing samples: {len(x_test)}")

AUTO = tf.data.AUTOTUNE
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_ds = train_ds.shuffle(config.buffer_size).batch(config.batch_size).prefetch(AUTO)

val_ds = tf.data.Dataset.from_tensor_slices((x_val, y_val))
val_ds = val_ds.batch(config.batch_size).prefetch(AUTO)

test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_ds = test_ds.batch(config.batch_size).prefetch(AUTO)

def get_augmentation_model():
    """데이터 보강 모델 빌드."""
    data_augmentation = keras.Sequential(
        [
            layers.Resizing(config.input_shape[0] + 20, config.input_shape[0] + 20),
            layers.RandomCrop(config.image_size, config.image_size),
            layers.RandomFlip("horizontal"),
            layers.Rescaling(1 / 255.0),
        ]
    )
    return data_augmentation

class MLP(layers.Layer):
    """각 시프트 블록에 대한 MLP 레이어를 가져옵니다.

    Args:
        mlp_expand_ratio (int): 첫 번째 특성 맵이 확장되는 비율입니다.
        mlp_dropout_rate (float): Dropout 비율입니다.
    """

    def __init__(self, mlp_expand_ratio, mlp_dropout_rate, **kwargs):
        super().__init__(**kwargs)
        self.mlp_expand_ratio = mlp_expand_ratio
        self.mlp_dropout_rate = mlp_dropout_rate

    def build(self, input_shape):
        input_channels = input_shape[-1]
        initial_filters = int(self.mlp_expand_ratio * input_channels)

        self.mlp = keras.Sequential(
            [
                layers.Dense(
                    units=initial_filters,
                    activation=tf.nn.gelu,
                ),
                layers.Dropout(rate=self.mlp_dropout_rate),
                layers.Dense(units=input_channels),
                layers.Dropout(rate=self.mlp_dropout_rate),
            ]
        )

    def call(self, x):
        x = self.mlp(x)
        return x

class DropPath(layers.Layer):
    """Stochastic Depth 레이어라고도 하는 Drop Path입니다.

    Refernece:
        - https://keras.io/examples/vision/cct/#stochastic-depth-for-regularization
        - github.com:rwightman/pytorch-image-models
    """

    def __init__(self, drop_path_prob, **kwargs):
        super().__init__(**kwargs)
        self.drop_path_prob = drop_path_prob

    def call(self, x, training=False):
        if training:
            keep_prob = 1 - self.drop_path_prob
            shape = (tf.shape(x)[0],) + (1,) * (len(tf.shape(x)) - 1)
            random_tensor = keep_prob + tf.random.uniform(shape, 0, 1)
            random_tensor = tf.floor(random_tensor)
            return (x / keep_prob) * random_tensor
        return x

out[:, g * 0:g * 1, :, :-1] = x[:, g * 0:g * 1, :, 1:]  # 왼쪽으로 시프트
out[:, g * 1:g * 2, :, 1:] = x[:, g * 1:g * 2, :, :-1]  # 오른쪽으로 시프트
out[:, g * 2:g * 3, :-1, :] = x[:, g * 2:g * 3, 1:, :]  # 위쪽으로 시프트
out[:, g * 3:g * 4, 1:, :] = x[:, g * 3:g * 4, :-1, :]  # 아래쪽으로 시프트

out[:, g * 4:, :, :] = x[:, g * 4:, :, :]  # 시프트 없음

class ShiftViTBlock(layers.Layer):
    """ShiftViT Block 유닛

    Args:
        shift_pixel (int): 이동할 픽셀 수입니다. 기본값은 1입니다.
        mlp_expand_ratio (int): MLP 특성이 확장되는 비율입니다. 기본값은 2입니다.
        mlp_dropout_rate (float): MLP에서 사용되는 dropout 비율입니다.
        num_div (int): 특성 맵 채널의 분할 수입니다. 총, 4/num_div의 채널이 시프트됩니다. 기본값은 12입니다.
        epsilon (float): Epsilon 상수입니다.
        drop_path_prob (float): drop path에 대한 drop 확률입니다.
    """

    def __init__(
        self,
        epsilon,
        drop_path_prob,
        mlp_dropout_rate,
        num_div=12,
        shift_pixel=1,
        mlp_expand_ratio=2,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.shift_pixel = shift_pixel
        self.mlp_expand_ratio = mlp_expand_ratio
        self.mlp_dropout_rate = mlp_dropout_rate
        self.num_div = num_div
        self.epsilon = epsilon
        self.drop_path_prob = drop_path_prob

    def build(self, input_shape):
        self.H = input_shape[1]
        self.W = input_shape[2]
        self.C = input_shape[3]
        self.layer_norm = layers.LayerNormalization(epsilon=self.epsilon)
        self.drop_path = (
            DropPath(drop_path_prob=self.drop_path_prob)
            if self.drop_path_prob > 0.0
            else layers.Activation("linear")
        )
        self.mlp = MLP(
            mlp_expand_ratio=self.mlp_expand_ratio,
            mlp_dropout_rate=self.mlp_dropout_rate,
        )

    def get_shift_pad(self, x, mode):
        """선택한 모드에 따라 채널을 시프트합니다."""
        if mode == "left":
            offset_height = 0
            offset_width = 0
            target_height = 0
            target_width = self.shift_pixel
        elif mode == "right":
            offset_height = 0
            offset_width = self.shift_pixel
            target_height = 0
            target_width = self.shift_pixel
        elif mode == "up":
            offset_height = 0
            offset_width = 0
            target_height = self.shift_pixel
            target_width = 0
        else:
            offset_height = self.shift_pixel
            offset_width = 0
            target_height = self.shift_pixel
            target_width = 0
        crop = tf.image.crop_to_bounding_box(
            x,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=self.H - target_height,
            target_width=self.W - target_width,
        )
        shift_pad = tf.image.pad_to_bounding_box(
            crop,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=self.H,
            target_width=self.W,
        )
        return shift_pad

    def call(self, x, training=False):
        # 특성 맵 분할
        x_splits = tf.split(x, num_or_size_splits=self.C // self.num_div, axis=-1)

        # 특성 맵 시프트
        x_splits[0] = self.get_shift_pad(x_splits[0], mode="left")
        x_splits[1] = self.get_shift_pad(x_splits[1], mode="right")
        x_splits[2] = self.get_shift_pad(x_splits[2], mode="up")
        x_splits[3] = self.get_shift_pad(x_splits[3], mode="down")

        # 시프트된 특성 맵과 시프트되지 않은 특성 맵 연결(Concatenate)하기
        x = tf.concat(x_splits, axis=-1)

        # residual 연결 추가
        shortcut = x
        x = shortcut + self.drop_path(self.mlp(self.layer_norm(x)), training=training)
        return x

class PatchMerging(layers.Layer):
    """Patch Merging 레이어.

    Args:
        epsilon (float): epsilon 상수입니다.
    """

    def __init__(self, epsilon, **kwargs):
        super().__init__(**kwargs)
        self.epsilon = epsilon

    def build(self, input_shape):
        filters = 2 * input_shape[-1]
        self.reduction = layers.Conv2D(
            filters=filters, kernel_size=2, strides=2, padding="same", use_bias=False
        )
        self.layer_norm = layers.LayerNormalization(epsilon=self.epsilon)

    def call(self, x):
        # 특성 맵에 패치 병합 알고리즘 적용하기
        x = self.layer_norm(x)
        x = self.reduction(x)
        return x

# 참고 : 이 레이어는 모델 단계마다 쌓는 깊이가 다릅니다.
class StackedShiftBlocks(layers.Layer):
    """stacked ShiftViT 블록을 포함하는 레이어입니다.

    Args:
        epsilon (float): epsilon 상수입니다.
        mlp_dropout_rate (float): MLP 블록에서 사용되는 dropout 비율입니다.
        num_shift_blocks (int): 이 스테이지의 shift vit 블록 수입니다.
        stochastic_depth_rate (float): 선택한 최대 drop path 비율입니다.
        is_merge (boolean): shift vit 이후, Patch Merge 레이어의 사용을 결정하는 플래그입니다.
        num_div (int): 특성 맵의 채널 division 입니다. 기본값은 12입니다.
        shift_pixel (int): shift 할 픽셀 수입니다. 기본값은 1입니다.
        mlp_expand_ratio (int): MLP의 초기 dense 레이어가 확장되는 비율입니다.
    """

    def __init__(
        self,
        epsilon,
        mlp_dropout_rate,
        num_shift_blocks,
        stochastic_depth_rate,
        is_merge,
        num_div=12,
        shift_pixel=1,
        mlp_expand_ratio=2,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.epsilon = epsilon
        self.mlp_dropout_rate = mlp_dropout_rate
        self.num_shift_blocks = num_shift_blocks
        self.stochastic_depth_rate = stochastic_depth_rate
        self.is_merge = is_merge
        self.num_div = num_div
        self.shift_pixel = shift_pixel
        self.mlp_expand_ratio = mlp_expand_ratio

    def build(self, input_shapes):
        # stochastic depth probabilities을 계산합니다.
        # Reference: https://keras.io/examples/vision/cct/#the-final-cct-model
        dpr = [
            x
            for x in np.linspace(
                start=0, stop=self.stochastic_depth_rate, num=self.num_shift_blocks
            )
        ]

        # 시프트 블록을 ShiftViT 블록 목록으로 빌드합니다.
        self.shift_blocks = list()
        for num in range(self.num_shift_blocks):
            self.shift_blocks.append(
                ShiftViTBlock(
                    num_div=self.num_div,
                    epsilon=self.epsilon,
                    drop_path_prob=dpr[num],
                    mlp_dropout_rate=self.mlp_dropout_rate,
                    shift_pixel=self.shift_pixel,
                    mlp_expand_ratio=self.mlp_expand_ratio,
                )
            )
        if self.is_merge:
            self.patch_merge = PatchMerging(epsilon=self.epsilon)

    def call(self, x, training=False):
        for shift_block in self.shift_blocks:
            x = shift_block(x, training=training)
        if self.is_merge:
            x = self.patch_merge(x)
        return x

    # 커스텀 레이어이므로, 트레이닝 후 모델을 쉽게 저장하고 로드할 수 있도록
    # get_config()를 덮어써야 합니다.
    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "epsilon": self.epsilon,
                "mlp_dropout_rate": self.mlp_dropout_rate,
                "num_shift_blocks": self.num_shift_blocks,
                "stochastic_depth_rate": self.stochastic_depth_rate,
                "is_merge": self.is_merge,
                "num_div": self.num_div,
                "shift_pixel": self.shift_pixel,
                "mlp_expand_ratio": self.mlp_expand_ratio,
            }
        )
        return config

class ShiftViTModel(keras.Model):
    """ShiftViT 모델.

    Args:
        data_augmentation (keras.Model): 데이터 보강 모델입니다.
        projected_dim (int): 이미지의 패치가 프로젝션 될 차원입니다.
        patch_size (int): 이미지의 패치 크기입니다.
        num_shift_blocks_per_stages (list[int]): 스테이지 당 모든 shift 블록의 개수 목록입니다.
        epsilon (float): 엡실론 상수.
        mlp_dropout_rate (float): MLP 블록에 사용되는 드롭아웃 비율입니다.
        stochastic_depth_rate (float): 최대 드랍률 확률.
        num_div (int): 특성 맵의 채널 분할 수. 기본값은 12입니다.
        shift_pixel (int): shift 할 픽셀 수. 기본값은 1입니다.
        mlp_expand_ratio (int): 초기 mlp dense 레이어가 확장되는 비율입니다. 기본값은 2입니다.
    """

    def __init__(
        self,
        data_augmentation,
        projected_dim,
        patch_size,
        num_shift_blocks_per_stages,
        epsilon,
        mlp_dropout_rate,
        stochastic_depth_rate,
        num_div=12,
        shift_pixel=1,
        mlp_expand_ratio=2,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.data_augmentation = data_augmentation
        self.patch_projection = layers.Conv2D(
            filters=projected_dim,
            kernel_size=patch_size,
            strides=patch_size,
            padding="same",
        )
        self.stages = list()
        for index, num_shift_blocks in enumerate(num_shift_blocks_per_stages):
            if index == len(num_shift_blocks_per_stages) - 1:
                # 이것은 마지막 단계이므로, 여기서는 패치 병합을 사용하지 마세요.
                is_merge = False
            else:
                is_merge = True
            # 스테이지를 구축하세요.
            self.stages.append(
                StackedShiftBlocks(
                    epsilon=epsilon,
                    mlp_dropout_rate=mlp_dropout_rate,
                    num_shift_blocks=num_shift_blocks,
                    stochastic_depth_rate=stochastic_depth_rate,
                    is_merge=is_merge,
                    num_div=num_div,
                    shift_pixel=shift_pixel,
                    mlp_expand_ratio=mlp_expand_ratio,
                )
            )
        self.global_avg_pool = layers.GlobalAveragePooling2D()

        self.classifier = layers.Dense(config.num_classes)

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "data_augmentation": self.data_augmentation,
                "patch_projection": self.patch_projection,
                "stages": self.stages,
                "global_avg_pool": self.global_avg_pool,
                "classifier": self.classifier,
            }
        )
        return config

    def _calculate_loss(self, data, training=False):
        (images, labels) = data

        # 이미지를 보강하세요
        augmented_images = self.data_augmentation(images, training=training)

        # 패치를 만들고 패치를 프로젝션합니다.
        projected_patches = self.patch_projection(augmented_images)

        # 스테이지를 통과하세요
        x = projected_patches
        for stage in self.stages:
            x = stage(x, training=training)

        # 로짓을 얻으세요.
        x = self.global_avg_pool(x)
        logits = self.classifier(x)

        # 손실을 계산하여 반환하세요.
        total_loss = self.compiled_loss(labels, logits)
        return total_loss, labels, logits

    def train_step(self, inputs):
        with tf.GradientTape() as tape:
            total_loss, labels, logits = self._calculate_loss(
                data=inputs, training=True
            )

        # 그라디언트를 적용합니다.
        train_vars = [
            self.data_augmentation.trainable_variables,
            self.patch_projection.trainable_variables,
            self.global_avg_pool.trainable_variables,
            self.classifier.trainable_variables,
        ]
        train_vars = train_vars + [stage.trainable_variables for stage in self.stages]

        # 그라디언트를 최적화합니다.
        grads = tape.gradient(total_loss, train_vars)
        trainable_variable_list = []
        for (grad, var) in zip(grads, train_vars):
            for g, v in zip(grad, var):
                trainable_variable_list.append((g, v))
        self.optimizer.apply_gradients(trainable_variable_list)

        # 메트릭 업데이트
        self.compiled_metrics.update_state(labels, logits)
        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        _, labels, logits = self._calculate_loss(data=data, training=False)

        # 메트릭 업데이트
        self.compiled_metrics.update_state(labels, logits)
        return {m.name: m.result() for m in self.metrics}

    def call(self, images):
        augmented_images = self.data_augmentation(images)
        x = self.patch_projection(augmented_images)
        for stage in self.stages:
            x = stage(x, training=False)
        x = self.global_avg_pool(x)
        logits = self.classifier(x)
        return logits

model = ShiftViTModel(
    data_augmentation=get_augmentation_model(),
    projected_dim=config.projected_dim,
    patch_size=config.patch_size,
    num_shift_blocks_per_stages=config.num_shift_blocks_per_stages,
    epsilon=config.epsilon,
    mlp_dropout_rate=config.mlp_dropout_rate,
    stochastic_depth_rate=config.stochastic_depth_rate,
    num_div=config.num_div,
    shift_pixel=config.shift_pixel,
    mlp_expand_ratio=config.mlp_expand_ratio,
)

# 일부 코드는 다음에서 가져왔습니다.
# https://www.kaggle.com/ashusma/training-rfcx-tensorflow-tpu-effnet-b2.
class WarmUpCosine(keras.optimizers.schedules.LearningRateSchedule):
    """워밍업 코사인 감쇠 스케쥴을 사용하는 LearningRateSchedule입니다."""

    def __init__(self, lr_start, lr_max, warmup_steps, total_steps):
        """
        Args:
            lr_start: 초기 학습률
            lr_max: 워밍업 단계에서 lr이 증가해야 하는 최대 학습률
            warmup_steps: 모델이 워밍업하는 단계 수
            total_steps: 모델 트레이닝을 위한 총 단계 수
        """
        super().__init__()
        self.lr_start = lr_start
        self.lr_max = lr_max
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.pi = tf.constant(np.pi)

    def __call__(self, step):
        # 총 단계 수가 워밍업 단계 수보다 작으면,
        # 값 오류를 throw합니다.
        if self.total_steps < self.warmup_steps:
            raise ValueError(
                f"Total number of steps {self.total_steps} must be"
                + f"larger or equal to warmup steps {self.warmup_steps}."
            )

        # `cos_annealed_lr`은 초기 단계에서 워밍업 단계까지 1로 증가하는 그래프입니다.
        # 그 후, 이 그래프는 최종 단계 마크에서 -1로 감소합니다.
        cos_annealed_lr = tf.cos(
            self.pi
            * (tf.cast(step, tf.float32) - self.warmup_steps)
            / tf.cast(self.total_steps - self.warmup_steps, tf.float32)
        )

        # `cos_annealed_lr` 그래프의 평균을 1로 옮깁니다.
        # 이제 그래프가 0에서 2로 이동합니다.
        # 그래프를 0.5로 정규화하여, 이제 0에서 1로 이동합니다.
        # 정규화된 그래프를 `lr_max`로 조정하여, 0에서 `lr_max`로 이동합니다.
        learning_rate = 0.5 * self.lr_max * (1 + cos_annealed_lr)

        # warmup_steps가 0보다 큰지 확인합니다.
        if self.warmup_steps > 0:
            # lr_max가 lr_start보다 큰지 확인합니다. 그렇지 않으면, 값 오류를 throw합니다.
            if self.lr_max < self.lr_start:
                raise ValueError(
                    f"lr_start {self.lr_start} must be smaller or"
                    + f"equal to lr_max {self.lr_max}."
                )

            # warumup 스케쥴에서 학습률이 증가해야 하는 기울기를 계산합니다.
            # 기울기 공식은 m = ((b-a)/steps)입니다.
            slope = (self.lr_max - self.lr_start) / self.warmup_steps

            # 직선 공식(y = mx+c)을 이용하여 워밍업 스케쥴을 빌드합니다.
            warmup_rate = slope * tf.cast(step, tf.float32) + self.lr_start

            # 현재 단계가 워밍업 단계보다 작으면, 선 그래프를 구합니다.
            # 현재 단계가 워밍업 단계보다 크면, 스케일된 코사인 그래프를 구합니다.
            learning_rate = tf.where(
                step < self.warmup_steps, warmup_rate, learning_rate
            )

        # 현재 단계가 총 단계보다 큰 경우 0을 반환하고, 그렇지 않으면 계산된 그래프를 반환합니다.
        return tf.where(
            step > self.total_steps, 0.0, learning_rate, name="learning_rate"
        )

    def get_config(self):
        config = {
            "lr_start": self.lr_start,
            "lr_max": self.lr_max,
            "total_steps": self.total_steps,
            "warmup_steps": self.warmup_steps,
        }
        return config

# 모델 저장 시 입력 모양을 사용할 수 있도록, 샘플 데이터를 모델에 전달합니다.
sample_ds, _ = next(iter(train_ds))
model(sample_ds, training=False)

# 트레이닝에 필요한 총 스텝 수를 계산합니다.
total_steps = int((len(x_train) / config.batch_size) * config.epochs)

# 워밍업을 위한 스텝 수를 계산하세요.
warmup_epoch_percentage = 0.15
warmup_steps = int(total_steps * warmup_epoch_percentage)

# 워밍업 코사인 스케쥴을 초기화합니다.
scheduled_lrs = WarmUpCosine(
    lr_start=1e-5,
    lr_max=1e-3,
    warmup_steps=warmup_steps,
    total_steps=total_steps,
)

# 옵티마이저를 얻습니다.
optimizer = tfa.optimizers.AdamW(
    learning_rate=scheduled_lrs, weight_decay=config.weight_decay
)

# 모델을 컴파일하고 사전 트레이닝합니다.
model.compile(
    optimizer=optimizer,
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[
        keras.metrics.SparseCategoricalAccuracy(name="accuracy"),
        keras.metrics.SparseTopKCategoricalAccuracy(5, name="top-5-accuracy"),
    ],
)

# 모델 트레이닝
history = model.fit(
    train_ds,
    epochs=config.epochs,
    validation_data=val_ds,
    callbacks=[
        keras.callbacks.EarlyStopping(
            monitor="val_accuracy",
            patience=5,
            mode="auto",
        )
    ],
)

# 테스트 데이터 세트로 모델을 평가합니다.
print("TESTING")
loss, acc_top1, acc_top5 = model.evaluate(test_ds)
print(f"Loss: {loss:0.2f}")
print(f"Top 1 test accuracy: {acc_top1*100:0.2f}%")
print(f"Top 5 test accuracy: {acc_top5*100:0.2f}%")

model.save("ShiftViT")

!wget -q 'https://tinyurl.com/2p9483sw' -O inference_set.zip
!unzip -q inference_set.zip

# 모델이 저장될 때 커스텀 객체는 포함되지 않습니다.
# 로딩 시, 이러한 객체는 모델 재구성을 위해 전달되어야 합니다.
saved_model = tf.keras.models.load_model(
    "ShiftViT",
    custom_objects={"WarmUpCosine": WarmUpCosine, "AdamW": tfa.optimizers.AdamW},
)

def process_image(img_path):
    # 문자열 경로에서 이미지 파일을 읽습니다.
    img = tf.io.read_file(img_path)

    # jpeg를 uint8 텐서로 디코딩
    img = tf.io.decode_jpeg(img, channels=3)

    # 모델에서 허용하는 입력 크기에 맞게 이미지 크기를 조정합니다.
    # `resize()`에 전달된 입력의 dtype을 보존하기 위해, `method`를 `nearest`로 사용합니다.
    img = tf.image.resize(
        img, [config.input_shape[0], config.input_shape[1]], method="nearest"
    )
    return img


def create_tf_dataset(image_dir):
    data_dir = pathlib.Path(image_dir)

    # 이미지 디렉토리를 사용하여, tf.data 데이터세트 생성
    predict_ds = tf.data.Dataset.list_files(str(data_dir / "*.jpg"), shuffle=False)

    # map을 사용하여 문자열 경로를 uint8 이미지 텐서로 변환하고,
    # `num_parallel_calls'를 설정하면 여러 이미지를 병렬로 처리하는 데 도움이 됩니다.
    predict_ds = predict_ds.map(process_image, num_parallel_calls=AUTO)

    # 더 나은 대기 시간과 처리량을 위해 Prefetch Dataset을 만듭니다.
    predict_ds = predict_ds.batch(config.tf_ds_batch_size).prefetch(AUTO)
    return predict_ds


def predict(predict_ds):
    # ShiftViT 모델은 로짓(정규화되지 않은 예측)을 반환합니다.
    logits = saved_model.predict(predict_ds)

    # softmax()를 호출하여 예측을 정규화합니다.
    probabilities = tf.nn.softmax(logits)
    return probabilities


def get_predicted_class(probabilities):
    pred_label = np.argmax(probabilities)
    predicted_class = config.label_map[pred_label]
    return predicted_class


def get_confidence_scores(probabilities):
    # 확률 점수의 인덱스를 내림차순으로 정렬합니다.
    labels = np.argsort(probabilities)[::-1]
    confidences = {
        config.label_map[label]: np.round((probabilities[label]) * 100, 2)
        for label in labels
    }
    return confidences

img_dir = "inference_set"
predict_ds = create_tf_dataset(img_dir)
probabilities = predict(predict_ds)
print(f"probabilities: {probabilities[0]}")
confidences = get_confidence_scores(probabilities[0])
print(confidences)

plt.figure(figsize=(10, 10))
for images in predict_ds:
    for i in range(min(6, probabilities.shape[0])):
        ax = plt.subplot(3, 3, i + 1)
        plt.imshow(images[i].numpy().astype("uint8"))
        predicted_class = get_predicted_class(probabilities[i])
        plt.title(predicted_class)
        plt.axis("off")