diff --git a/extraction/core/functions/distributors.py b/extraction/core/functions/distributors.py
index 862b858053f9c3cf57e422a0648fcc4bc26d72c9..421835910510a7edd5eb0cd97bf7f140e9d511a6 100644
--- a/extraction/core/functions/distributors.py
+++ b/extraction/core/functions/distributors.py
@@ -1,21 +1,19 @@
 import numpy as np
 
 
-def trap_apply(cell_fun, cell_masks, trap_image, **kwargs):
+def trap_apply(cell_fun, cell_masks, *args, **kwargs):
     """
     Apply a cell_function to a mask, trap_image pair
 
     :param cell_fun: function to apply to a cell (from extraction/cell.py)
     :param cell_masks: (numpy 3d array) cells' segmentation mask
-    :param trap_image: the image for the trap in which the cell is (all
+    :param trap_image: (Optional) the image for the trap in which the cell is (all
     channels)
     :**kwargs: parameters to pass if needed for custom functions
     """
 
     cells_iter = (*range(cell_masks.shape[2]),)
-    return [
-        cell_fun(cell_masks[..., i], trap_image, **kwargs) for i in cells_iter
-    ]
+    return [cell_fun(cell_masks[..., i], *args, **kwargs) for i in cells_iter]
 
 
 def reduce_z(trap_image, fun):