diff --git a/petric.py b/petric.py
index ec62dc7..ac821f2 100755
--- a/petric.py
+++ b/petric.py
@@ -92,8 +92,8 @@ class QualityMetrics(ImageQualityCallback):
     """From https://github.com/SyneRBI/PETRIC/wiki#metrics-and-thresholds"""
     def __init__(self, reference_image, whole_object_mask, background_mask, **kwargs):
         super().__init__(reference_image, **kwargs)
-        self.whole_object_indices = np.where(whole_object_mask == 1)
-        self.background_indices = np.where(background_mask == 1)
+        self.whole_object_indices = np.where(whole_object_mask.as_array())
+        self.background_indices = np.where(background_mask.as_array())
         self.ref_im_arr = reference_image.as_array()
         self.norm = self.ref_im_arr[self.background_indices].mean()
 
@@ -238,7 +238,7 @@ def get_image(fname):
         if data.reference_image is not None:
             metrics_with_timeout.callbacks.append(
                 QualityMetrics(data.reference_image, data.whole_object_mask, data.background_mask,
-                               tb_summary_writer=metrics_with_timeout.tb, roi_mask_dict=data.voi_masks))
+                               tb_summary_writer=metrics_with_timeout.tb, voi_mask_dict=data.voi_masks))
         metrics_with_timeout.reset() # timeout from now
         algo = Submission(data)
         try: