watershed transform to_categorical should use num_classes=distance_bins

vanvalenlab · Dec 18, 2018 · 9020a2d · 9020a2d
1 parent de0bf95
commit 9020a2d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepcell/image_generators.py b/deepcell/image_generators.py
@@ -120,7 +120,8 @@ def _transform_masks(y, transform, data_format=None, **kwargs):
                 mask, distance_bins, erosion)
 
         # convert to one hot notation
-        y_transform = to_categorical(np.expand_dims(y_transform, axis=-1))
+        y_transform = np.expand_dims(y_transform, axis=-1)
+        y_transform = to_categorical(y_transform, num_classes=distance_bins)
         if data_format == 'channels_first':
             y_transform = np.rollaxis(y_transform, y.ndim - 1, 1)