diff --git a/app/mains/functions.py b/app/mains/functions.py
index 87cd14f..decfac8 100644
--- a/app/mains/functions.py
+++ b/app/mains/functions.py
@@ -65,11 +65,15 @@ def scale_to_0_255(data):
 
 
 def display_images(
-    data, layer_number, embedding_method: str, dataset_name: str, suptitle: str
+    data, label, layer_number, embedding_method: str, dataset_name: str, suptitle: str
 ):
     num_data = 3
+    label = np.argmax(label, axis=1)
+    indices = [np.where(label == i)[0][0] for i in [1, 2, 3]] # 1,2,3のラベルを持つデータの最初のインデックスを取得
+    
     img_idx = 1
-    for n in range(3):
+    for n in indices:
+        print(n)
         data_to_display = data[n]
         # data_to_display = scale_to_0_255(data_to_display)
         if data_to_display.shape[2] == 6:
diff --git a/app/mains/layers.py b/app/mains/layers.py
index 8e19ea4..9e9c4df 100644
--- a/app/mains/layers.py
+++ b/app/mains/layers.py
@@ -544,6 +544,7 @@ def fit(self, X, Y):
                     )
                     display_images(
                         X,
+                        Y,
                         n + 2,
                         layer.embedding,
                         self.data_set_name,
diff --git a/app/mains/main_LeNet.py b/app/mains/main_LeNet.py
index 537e5b7..cb98475 100644
--- a/app/mains/main_LeNet.py
+++ b/app/mains/main_LeNet.py
@@ -41,7 +41,6 @@ def main_LeNet(
     train_X, train_Y, test_X, test_Y, channel, image_size = select_datasets(
         num_train, num_test, datasets
     )
-    display_images(train_X, 1, "LeNet_train", datasets, f"Input Layer1")
     # LeNet-5 model definition
     activation = "relu"
     # activation = 'tanh'
@@ -174,6 +173,7 @@ def main_LeNet(
         )
         display_images(
             block_outputs[0],
+            train_Y,
             2,
             "LeNet",
             datasets,