diff --git a/AI_brain/image_recognition.py b/AI_brain/image_recognition.py
index 58ba878..d5c097b 100644
--- a/AI_brain/image_recognition.py
+++ b/AI_brain/image_recognition.py
@@ -7,36 +7,37 @@ import random
 
 
 class VacuumRecognizer:
-    model = keras.models.load_model("D:/Image_dataset/model.h5")
+    model = keras.models.load_model('AI_brain\model.h5')
 
     def recognize(self, image_path) -> str:
         class_names = ['Banana', 'Cat', 'Earings', 'Plant']
 
         img = cv2.imread(image_path, flags=cv2.IMREAD_GRAYSCALE)
-        # print(img.shape)
-        cv2.imshow("lala", img)
         cv2.waitKey(0)
         img = (np.expand_dims(img, 0))
 
         predictions = self.model.predict(img)[0].tolist()
 
-        print(class_names)
-        print(predictions)
-        print(max(predictions))
-        print(predictions.index(max(predictions)))
+        # print(img.shape)
+        # cv2.imshow("test_show", img)
+        # print(class_names)
+        # print(predictions)
+        # print(max(predictions))
+        # print(predictions.index(max(predictions)))
 
         return class_names[predictions.index(max(predictions))]
 
 
-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/')
-uio = VacuumRecognizer()
 
+#For testing the neuron model
+'''image_paths = []
+image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Banana')
+image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Cat')
+image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Earings')
+image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Plant')
+uio = VacuumRecognizer()
 
 for image_path in image_paths:
     dirs = os.listdir(image_path)
-    for i in range(10):
-        print(uio.recognize(image_path + dirs[random.randint(0, len(dirs)-1)]))
\ No newline at end of file
+    for i in range(3):
+        print(uio.recognize(image_path + '\\' + dirs[random.randint(0, len(dirs)-1)]))'''
\ No newline at end of file
diff --git a/domain/world.py b/domain/world.py
index d3a7e8a..caf17f8 100644
--- a/domain/world.py
+++ b/domain/world.py
@@ -15,14 +15,14 @@ class World:
         self.doc_station = None
 
     def add_entity(self, entity: Entity):
-        if entity.type == "PEEL":
+        if entity.type == "DOC_STATION":
+            self.doc_station = entity
+        elif entity.type == "PEEL":
             self.dust[entity.x][entity.y].append(entity)
         elif entity.type == "EARRING":
             self.dust[entity.x][entity.y].append(entity)
         elif entity.type == "VACUUM":
             self.vacuum = entity
-        elif entity.type == "DOC_STATION":
-            self.doc_station = entity
         elif entity.type == "CAT":
             self.cat = entity
             self.obstacles[entity.x][entity.y].append(entity)