allow onehot with Nones

mpl-extensions · Nov 18, 2022 · 8f3aed4 · 8f3aed4
1 parent 63fdf18
commit 8f3aed4
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/mpl_image_labeller/_util.py b/mpl_image_labeller/_util.py
@@ -52,7 +52,9 @@ def list_to_onehot(labels, classes):
     arr = np.zeros((len(labels), len(classes)), dtype=bool)
     for i, l in enumerate(labels):
 
-        if isinstance(l, str) or not isinstance(l, Iterable):
+        if l is None:
+            continue
+        elif isinstance(l, str) or not isinstance(l, Iterable):
             # str, or number, or something like that
             arr[i, lookup[l]] = True
         else: