diff --git a/.gitignore b/.gitignore
index cba523c..92022f8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,6 +14,7 @@ ModelNet40/
 ModelNet40.zip.1
 ModelNet40.zip
 results_bak
+dont_track/*
 binvox
 viewvox
 *.log
\ No newline at end of file
diff --git a/ModelNet10_class_imbalance.png b/ModelNet10_class_imbalance.png
index 1b5596b..66471f6 100644
Binary files a/ModelNet10_class_imbalance.png and b/ModelNet10_class_imbalance.png differ
diff --git a/ModelNet40.png b/ModelNet40.png
deleted file mode 100644
index 1b5596b..0000000
Binary files a/ModelNet40.png and /dev/null differ
diff --git a/ModelNet40_class_imbalance.png b/ModelNet40_class_imbalance.png
new file mode 100644
index 0000000..479ccb4
Binary files /dev/null and b/ModelNet40_class_imbalance.png differ
diff --git a/count_files.py b/count_files.py
index f3816c9..c279037 100644
--- a/count_files.py
+++ b/count_files.py
@@ -32,7 +32,7 @@ def make_file_description_data_frame(path='./ModelNet10/', show_missed=True):
     return pd.DataFrame(paths)
 
 
-def plot_unbalanced_classes(df, title, display_df=False):
+def plot_unbalanced_classes(df, title, display_df=False, save=False):
     grouped = df.groupby(['class', 'type'])['type'].count()
     unstacked = grouped.unstack()
     plt.figure()
@@ -40,6 +40,8 @@ def plot_unbalanced_classes(df, title, display_df=False):
     unstacked.plot(kind='bar', title=title, figsize=(10, 6), yticks=np.arange(0, unstacked['train'].max(), 100))
     plt.ylabel('Number of Samples')
     # plt.grid()
+    if save:
+        plt.savefig(title + '.png')
     plt.show()
     if display_df:
         display(unstacked)