diff --git a/image_processing/build_dataset.py b/image_processing/build_dataset.py
index 67718f795a80d675e817fe5940ceae8e532bd195..1357c665905a519e5fc94ab6bcb70e15dffad823 100644
--- a/image_processing/build_dataset.py
+++ b/image_processing/build_dataset.py
@@ -5,6 +5,7 @@ import pandas as pd
 import re
 import numpy as np
 import matplotlib.image as mpimg
+from pyRawMSDataReader.pyRawMSDataReader.WiffFileReader_py import WiffFileReader
 
 
 from build_image import build_image_ms1_wiff_charge_filtered, build_image_ms1_wiff_charge_filtered_apex_only
diff --git a/image_processing/build_image.py b/image_processing/build_image.py
index 7d7d67d52bf8066091a8ba56d512ef5d434e48eb..1088a7a66725ca426023ae006de2b5756942ce9e 100644
--- a/image_processing/build_image.py
+++ b/image_processing/build_image.py
@@ -1,10 +1,6 @@
 import numpy as np
-import matplotlib.pyplot as plt
-import matplotlib.colors as colors
 import pyopenms as oms
 from matplotlib import image as mpimg
-from torch.onnx.symbolic_opset12 import argmax
-
 from pyRawMSDataReader.pyRawMSDataReader.WiffFileReader_py import WiffFileReader