diff --git a/AI_brain/image_recognition.py b/AI_brain/image_recognition.py
index 58ba878..a36ce3b 100644
--- a/AI_brain/image_recognition.py
+++ b/AI_brain/image_recognition.py
@@ -7,14 +7,18 @@ import random
 
 
 class VacuumRecognizer:
-    model = keras.models.load_model("D:/Image_dataset/model.h5")
+    model = keras.models.load_model("D:/Image_dataset/model_color.h5")
+    #model = keras.models.load_model("D:/Image_dataset/model.h5")
+
 
     def recognize(self, image_path) -> str:
         class_names = ['Banana', 'Cat', 'Earings', 'Plant']
 
-        img = cv2.imread(image_path, flags=cv2.IMREAD_GRAYSCALE)
+
+        img = cv2.imread(image_path)
+        #img = cv2.imread(image_path, flags=cv2.IMREAD_GRAYSCALE)
         # print(img.shape)
-        cv2.imshow("lala", img)
+        cv2.imshow("img", img)
         cv2.waitKey(0)
         img = (np.expand_dims(img, 0))
 
@@ -29,10 +33,15 @@ class VacuumRecognizer:
 
 
 image_paths = []
-image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Banana/')
-image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Cat/')
-image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Earings/')
-image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Plant/')
+image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Banana/')
+image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Cat/')
+image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Earings/')
+image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Plant/')
+
+# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/train/Banana/')
+# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Cat/')
+# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Earings/')
+# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Plant/')
 uio = VacuumRecognizer()
 
 
diff --git a/AI_brain/train.py b/AI_brain/train.py
new file mode 100644
index 0000000..1f8b8ff
--- /dev/null
+++ b/AI_brain/train.py
@@ -0,0 +1,53 @@
+import numpy as np
+import tensorflow as tf
+from tensorflow import keras
+
+
+def norm_pix(image, label):
+    return image / 255, label
+
+
+directoryTRAIN = "D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/train"
+directoryVALIDATION = "D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/validation"
+
+train_ds = tf.keras.utils.image_dataset_from_directory(directoryTRAIN,
+                                                       seed=123, batch_size=16,
+                                                       image_size=(200, 200), color_mode='rgb')#color_mode='grayscale')
+
+
+val_ds = tf.keras.utils.image_dataset_from_directory(directoryVALIDATION,
+                                                     seed=123, batch_size=16,
+                                                     image_size=(200, 200), color_mode='rgb')#color_mode='grayscale')
+
+
+train_ds = train_ds.map(norm_pix)
+val_ds = val_ds.map(norm_pix)
+
+model = keras.Sequential([
+    #keras.layers.Conv2D(64, (3, 3), activation='relu', input_shape=(200, 200, 1)),
+    keras.layers.Conv2D(64, (3, 3), activation='relu', input_shape=(200, 200, 3)),
+    keras.layers.MaxPool2D((2, 2)),
+
+    keras.layers.Conv2D(128, (3, 3), activation='relu'),
+    keras.layers.MaxPool2D((2, 2)),
+
+    keras.layers.Conv2D(256, (3, 3), activation='relu'),
+    keras.layers.MaxPool2D((2, 2)),
+
+    keras.layers.Flatten(),
+    keras.layers.Dense(1024, activation='relu'),
+    keras.layers.Dense(4, activation='softmax')
+])
+
+print(model.summary())
+
+model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
+
+trainHistory = model.fit(train_ds, epochs=8, validation_data=val_ds)
+
+#model.save("model.h5")
+model.save("model_color.h5")
+
+(loss, accuracy) = model.evaluate(val_ds)
+print(loss)
+print(accuracy)