Merge pull request #3737 from DanteOz/colsplitter-on-value

Add split value argument to ColSplitter
fastai · Jul 17, 2022 · caf805e · caf805e
2 parents fa188e5 + 8ae46d8
commit caf805e
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 21 deletions.
diff --git a/fastai/data/transforms.py b/fastai/data/transforms.py
@@ -160,11 +160,14 @@ def _inner(o): return FuncSplitter(_func)(o)
     return _inner
 
 # Cell
-def ColSplitter(col='is_valid'):
+def ColSplitter(col='is_valid', on=None):
     "Split `items` (supposed to be a dataframe) by value in `col`"
     def _inner(o):
         assert isinstance(o, pd.DataFrame), "ColSplitter only works when your items are a pandas DataFrame"
-        valid_idx = (o.iloc[:,col] if isinstance(col, int) else o[col]).values.astype('bool')
+        c = o.iloc[:,col] if isinstance(col, int) else o[col]
+        if on is None:      valid_idx = c.values.astype('bool')
+        elif is_listy(on):  valid_idx = c.isin(on)
+        else:               valid_idx = c == on
         return IndexSplitter(mask2idxs(valid_idx))(o)
     return _inner