allow multiple process_fn arguments

Author: subpic

generators.py

@@ -198,9 +198,11 @@ def _data_generation(self, ids_batch):
                     # if needed, process each image, and add to X_list (inputs list)
                     if params.process_fn not in [None, False]:                
                         data_list = []
-                        for i, row in enumerate(ids_batch.itertuples()):                    
-                            arg = [] if args_name is None else [getattr(row, args_name)]
-                            data_i = params.process_fn(data[i], *arg)
+                        for i, row in enumerate(ids_batch.itertuples()):
+                            args = []
+                            if args_name is not None:
+                                args = [getattr(row, name) for name in force_list(args_name)]         
+                            data_i = params.process_fn(data[i], *args)
                             data_list.append(force_list(data_i))
 
                         # transpose list, sublists become batches

tests/test_generators.py

@@ -146,6 +146,22 @@ def preproc(im, *arg):
     assert np.array_equal(np.squeeze(g[0][0][1]), np.arange(1,5))
     assert np.array_equal(np.squeeze(g[0][1]), np.arange(1,5))
 
+def test_multi_process_args_DataGeneratorHDF5():
+    def preproc(im, arg1, arg2):
+        return np.zeros(1) + arg1 + arg2
+
+    gen_params_local = gen_params.copy()
+    gen_params_local.process_fn = preproc
+    gen_params_local.process_args  = {'filename': ['filename_args','filename_args']}
+    gen_params_local.batch_size = 4
+
+    ids_local = ids.copy()
+    ids_local['filename_args'] = range(len(ids_local))
+
+    g = gr.DataGeneratorDisk(ids_local, **gen_params_local)
+    x = g[0]
+    assert np.array_equal(np.squeeze(x[0][0].T), np.arange(4)*2)
+    
 def test_callable_outputs_DataGeneratorHDF5():
     d = {'features': [1, 2, 3, 4, 5],
          'mask': [1, 0, 1, 1, 0]}

tests/tests.ipynb

@@ -50,12 +50,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def preproc(im, arg):\n",
-    "    return np.zeros(1) + arg\n",
+    "def preproc(im, arg1, arg2):\n",
+    "    return np.zeros(1) + arg1 + arg2\n",
     "\n",
     "gen_params_local = gen_params.copy()\n",
     "gen_params_local.process_fn = preproc\n",
-    "gen_params_local.process_args  = {'filename': 'filename_args'}\n",
+    "gen_params_local.process_args  = {'filename': ['filename_args','filename_args']}\n",
     "gen_params_local.batch_size = 4\n",
     "\n",
     "ids_local = ids.copy()\n",
@@ -64,7 +64,7 @@
     "g = gr.DataGeneratorDisk(ids_local, **gen_params_local)\n",
     "x = g[0]\n",
     "# gen.pretty(g)\n",
-    "assert np.array_equal(np.squeeze(x[0].T), np.arange(4))"
+    "assert np.array_equal(np.squeeze(x[0][0].T), np.arange(4)*2)"
    ]
   },
   {