diff --git a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB11_AER.png b/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB11_AER.png
deleted file mode 100644
index 58d5872d687a503a5f0a1cb4b6020741cf84bd06..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB11_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB11_ANA.png b/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB11_ANA.png
deleted file mode 100644
index 31cfcc6325456f3eb1e1d0c46b7f33c599647b39..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB11_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB6_AER.png b/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB6_AER.png
deleted file mode 100644
index 81bc8bbc17e7d2d4a8b2cf5ed40cc09803e946f5..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB6_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB6_ANA.png b/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB6_ANA.png
deleted file mode 100644
index eca7f6fc42f34b28234ed08257da517b64679ebc..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB6_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB7_AER.png b/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB7_AER.png
deleted file mode 100644
index 0ae9f44d78cea9f073e02c84175fe011d892ac94..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB7_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB7_ANA.png b/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB7_ANA.png
deleted file mode 100644
index 87a1c33a674866b249e93646bc9189f0ce7ebd3e..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB7_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB8_AER.png b/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB8_AER.png
deleted file mode 100644
index d1e13cd0e9d41e9883d5e4e47beecc3d031e1e33..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB8_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB8_ANA.png b/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB8_ANA.png
deleted file mode 100644
index a2d1cf085e3b694fb10f849625d3f5bcf2171dc8..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB8_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB9_AER.png b/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB9_AER.png
deleted file mode 100644
index a8cb2f16cb2e3f9868f000584e7e50836a17ce5d..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB9_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB9_ANA.png b/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB9_ANA.png
deleted file mode 100644
index 4853cfc0173a669e344d973bfa5f8e7f2d71b6ae..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Enterobacter asburiae/ENTASB9_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC13_AER.png b/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC13_AER.png
deleted file mode 100644
index 54cebe28b98af4618016981d83d0c646491a482a..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC13_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC13_ANA.png b/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC13_ANA.png
deleted file mode 100644
index abcfd824e2ee4196b96d1c476ab2467947a1196b..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC13_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC1_AER.png b/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC1_AER.png
deleted file mode 100644
index f2347d99dccbce18c98fac6c1535e3b1a681f888..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC1_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC1_ANA.png b/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC1_ANA.png
deleted file mode 100644
index 86ff33a87c92264a60bf55bbb018d11e0f46855f..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC1_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC3_AER.png b/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC3_AER.png
deleted file mode 100644
index aec8ea2bb1ceefbaa1a6d44fde649d258c26d46b..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC3_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC3_ANA.png b/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC3_ANA.png
deleted file mode 100644
index 841b962ee6068bd6163bdb368f3e0d0478dcb126..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC3_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC4_AER.png b/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC4_AER.png
deleted file mode 100644
index b8eeee69c5d09314b0d225bfe3a0087d2f4ed129..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC4_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC4_ANA.png b/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC4_ANA.png
deleted file mode 100644
index c23610827a9278c5ae360072aa633b8e4ec83168..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC4_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC7_AER.png b/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC7_AER.png
deleted file mode 100644
index c0d2aeeee4fc254afac2ad2c5457f9e96f237828..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC7_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC7_ANA.png b/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC7_ANA.png
deleted file mode 100644
index b64b22c001968c3e12360e39c4d8cf925fcf9f3b..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella michiganensis/KLEMIC7_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY18_AER.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY18_AER.png
deleted file mode 100644
index 4701cc0763863527c9297e4f47fdf19757e56e14..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY18_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY18_ANA.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY18_ANA.png
deleted file mode 100644
index db9c81950e8b097370ae0074b3d9ca38b7b08cbf..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY18_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY23_AER.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY23_AER.png
deleted file mode 100644
index c5bb56da3ff03c167a6537ac9026c7b64a93b255..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY23_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY23_ANA.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY23_ANA.png
deleted file mode 100644
index ff5cec37caa5483c61d096c2200c38ccf789d8ea..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY23_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY31_AER.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY31_AER.png
deleted file mode 100644
index 956c4f8fe73e0c6941517101e85742c81659750e..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY31_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY31_ANA.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY31_ANA.png
deleted file mode 100644
index a36fb58c338383ce47bb4488c478856c6db798d3..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY31_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY32_AER.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY32_AER.png
deleted file mode 100644
index b289d8fa712fbfca6e5913be26e5f9952d4409e4..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY32_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY32_ANA.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY32_ANA.png
deleted file mode 100644
index 266aaed46240b3a9993fe9ec67444d0c77681229..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY32_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY36_AER.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY36_AER.png
deleted file mode 100644
index bd98f438bfa9bea40492a19aaae8e624bc1313f2..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY36_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY36_ANA.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY36_ANA.png
deleted file mode 100644
index 2b1a8eec04c591dde96d271bdc72d2dd860772dd..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY36_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY37_AER.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY37_AER.png
deleted file mode 100644
index e9d5fd2f16727f6be535ae7e46057fa165979d63..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY37_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY37_ANA.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY37_ANA.png
deleted file mode 100644
index cf973296acf346a7672e3a2d301b6ed2a22af16e..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY37_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY42_AER.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY42_AER.png
deleted file mode 100644
index 8aff3ffa3f21c50f9f4e078ad8fc2f3b22aef985..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY42_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY42_ANA.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY42_ANA.png
deleted file mode 100644
index a89fa7a7da7e1416701aecec01cbd2eaf2eb24bf..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY42_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY7_AER.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY7_AER.png
deleted file mode 100644
index 4692438381428d9b305070ad95a3e4c144ffb65d..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY7_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY7_ANA.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY7_ANA.png
deleted file mode 100644
index b43859a6bac30f17246f5edfdff4de3766058864..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY7_ANA.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY8_AER.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY8_AER.png
deleted file mode 100644
index 9a61a3912e24cc1969b1d4f4fed264e71f28d9cf..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY8_AER.png and /dev/null differ
diff --git a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY8_ANA.png b/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY8_ANA.png
deleted file mode 100644
index c770292a2648b69d9cc123a318efa72d30799960..0000000000000000000000000000000000000000
Binary files a/data/processed_data/png_image/data_training/Klebsiella oxytoca/KLEOXY8_ANA.png and /dev/null differ
diff --git a/image_processing/build_dataset.py b/image_processing/build_dataset.py
index 6f7c4f71cbd15ed21771b20a99dfcf8d8bde6ad6..1123acdadc24df556cb0b7d65c3f3ae3d68bec34 100644
--- a/image_processing/build_dataset.py
+++ b/image_processing/build_dataset.py
@@ -38,7 +38,7 @@ def create_antibio_dataset(path='../data/label_raw/230804_strain_peptides_antibi
         l = split_before_number(s)
         species = l[0]
         nb = l[1]
-        return '{}-{}-{}_100vW_100SPD.mzML'.format(species,nb,analyse)
+        return '{}-{}-{}-d200.mzML'.format(species,nb,analyse)
 
     df['path_ana'] = df['sample_name'].map(lambda x: create_fname(x,analyse='ANA'))
     df['path_aer'] = df['sample_name'].map(lambda x: create_fname(x, analyse='AER'))
@@ -64,14 +64,16 @@ def create_dataset():
             name = label[label['path_aer'] == path.split("/")[-1]]['sample_name'].values[0]
             analyse = 'AER'
         if species is not None:
-            directory_path = '../data/processed_data/{}'.format(species)
-            if not os.path.isdir(directory_path):
-                os.makedirs(directory_path)
+            directory_path_png = '../data/processed_data/png_image/{}'.format(species)
+            directory_path_npy = '../data/processed_data/npy_image/{}'.format(species)
+            if not os.path.isdir(directory_path_png):
+                os.makedirs(directory_path_png)
+            if not os.path.isdir(directory_path_npy):
+                os.makedirs(directory_path_npy)
             mat = build_image_ms1(path, 1)
-            mpimg.imsave(directory_path + "/" + name + '_' + analyse + '.png', mat)
-            np.save(directory_path + "/" + name + '_' + analyse + '.npy', mat)
+            mpimg.imsave(directory_path_png + "/" + name + '_' + analyse + '.png', mat)
+            np.save(directory_path_npy + "/" + name + '_' + analyse + '.npy', mat)
 
 
-#TODO : train val test split
 if __name__ =='__main__' :
     create_dataset()
\ No newline at end of file
diff --git a/main.py b/main.py
index cd8556e94654b87dc3fbba68d36479e7e4ff12a8..ba90ffaf5e1bf17edab389546101ad77cc31b7ad 100644
--- a/main.py
+++ b/main.py
@@ -59,7 +59,10 @@ def test(model, data_test, loss_function, epoch):
     return losses,acc
 
 def run(args):
-    model = Classification_model(n_class=9)
+    data_train, data_test = load_data(base_dir=args.dataset_dir, batch_size=args.batch_size)
+    model = Classification_model(n_class=len(data_train.dataset.dataset.classes))
+    if args.pretrain_path is not None :
+        load_model(model,args.pretrain_path)
     if torch.cuda.is_available():
         model = model.cuda()
     best_acc = 0
@@ -67,9 +70,6 @@ def run(args):
     train_loss=[]
     val_acc=[]
     val_loss=[]
-    if args.pretrain_path is not None :
-        load_model(model,args.pretrain_path)
-    data_train, data_test = load_data(base_dir=args.dataset_dir, batch_size=args.batch_size)
     loss_function = nn.CrossEntropyLoss()
     optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
 
@@ -121,7 +121,7 @@ def make_prediction(model, data):
                          columns=[i for i in classes])
     plt.figure(figsize=(12, 7))
     sn.heatmap(df_cm, annot=True)
-    plt.savefig('output.png')
+    plt.savefig('confusion_matrix.png')
 
 
 def save_model(model, path):