diff --git a/src/DecisionTree.py b/src/DecisionTree.py
index 750e3efb..a3d37e06 100644
--- a/src/DecisionTree.py
+++ b/src/DecisionTree.py
@@ -9,7 +9,7 @@ from sklearn.tree import export_graphviz
 import pydotplus
 
 # Read the dataset
-train_data_m = pd.read_csv("datasets/treedataset2.csv")
+train_data_m = pd.read_csv("datasets/bees.csv")
 
 feature_cols = ['Weather','Temperature','WindSpeed','PollenAvailability','BeePopulation','TimeOfDay','FlowerDistance','Humidity']
 print(train_data_m.head())
@@ -36,5 +36,4 @@ export_graphviz(clf, out_file=dot_data,
                 filled=True, rounded=True,
                 special_characters=True, feature_names = feature_cols,class_names=['0','1'])
 graph = pydotplus.graph_from_dot_data(dot_data.getvalue())
-graph.write_png('tree.png')
-Image(graph.create_png())
\ No newline at end of file
+graph.write_png('tree.png')
\ No newline at end of file
diff --git a/src/datasets/testdataset.csv b/src/datasets/testdataset.csv
index cd8ea134..32f35535 100644
--- a/src/datasets/testdataset.csv
+++ b/src/datasets/testdataset.csv
@@ -1,4 +1,4 @@
-Weather,Temperature,Wind Speed,Pollen Availability,Bee Population,Time of Day,Flower Distance,Humidity,Fly_out_for_honey
+Weather,Temperature,WindSpeed,PollenAvailability,BeePopulation,TimeOfDay,FlowerDistance,Humidity,FlyOutForHoney
 Cloudy,Medium,Strong,Low,High,Morning,Close,Normal,No
 Cloudy,Low,Strong,High,Moderate,Morning,Close,High,Yes
 Sunny,High,Strong,Medium,Moderate,Morning,Close,Normal,Yes
diff --git a/src/datasets/treedataset.csv b/src/datasets/treedataset.csv
index b51936e6..1ee7f1b0 100644
--- a/src/datasets/treedataset.csv
+++ b/src/datasets/treedataset.csv
@@ -1,4 +1,4 @@
-Weather,Temperature,Wind Speed,Pollen Availability,Bee Population,Time of Day,Flower Distance,Humidity,Fly_out_for_honey
+Weather,Temperature,WindSpeed,PollenAvailability,BeePopulation,TimeOfDay,FlowerDistance,Humidity,FlyOutForHoney
 Cloudy,High,Strong,Medium,Moderate,Evening,Close,High,Yes
 Rainy,Low,Strong,Medium,Moderate,Afternoon,Medium,High,No
 Cloudy,Medium,Strong,Medium,Moderate,Evening,Medium,Normal,Yes
diff --git a/src/main.py b/src/main.py
index 81660f83..239a84ee 100644
--- a/src/main.py
+++ b/src/main.py
@@ -100,7 +100,6 @@ bee_dir = 'west'
 field = Field(tilemap_types, tilemap_values)
 
 
-
 def Manhattan_dis(start, end):
     start_x, start_y = start
     end_x, end_y = end
diff --git a/src/tree.png b/src/tree.png
index 5439e159..39fe60e8 100644
Binary files a/src/tree.png and b/src/tree.png differ