Merge pull request janelia-flyem#78 from jni/minor-fixes

Some bug fixes and speed improvements

gala/classify.py

@@ -187,10 +187,11 @@ def get_classifier(name='random forest', *args, **kwargs):
 
 class DefaultRandomForest(RandomForestClassifier):
     def __init__(self, n_estimators=100, criterion='entropy', max_depth=20,
-            bootstrap=False, random_state=None):
+            bootstrap=False, random_state=None, n_jobs=-1):
         super(DefaultRandomForest, self).__init__(
             n_estimators=n_estimators, criterion=criterion,
-            max_depth=max_depth, bootstrap=bootstrap, random_state=random_state)
+            max_depth=max_depth, bootstrap=bootstrap,
+            random_state=random_state, n_jobs=n_jobs)
 
 
 class VigraRandomForest(object):

gala/morpho.py

@@ -26,6 +26,8 @@
 from skimage import measure, util
 import skimage.morphology
 
+from sklearn.externals import joblib
+
 zero3d = array([0,0,0])
 
 def complement(a):
@@ -228,8 +230,7 @@ def watershed(a, seeds=None, connectivity=1, mask=None, smooth_thresh=0.0,
     if not seeded:
         seeds = regional_minima(a, connectivity)
     if minimum_seed_size > 0:
-        seeds = remove_small_connected_components(seeds, minimum_seed_size,
-                                                  in_place=True)
+        seeds = remove_small_connected_components(seeds, minimum_seed_size)
         seeds = relabel_from_one(seeds)[0]
     if smooth_seeds:
         seeds = binary_opening(seeds, sel)
@@ -269,7 +270,7 @@ def watershed(a, seeds=None, connectivity=1, mask=None, smooth_thresh=0.0,
                                     (br[nidxs] == level)).astype(bool) ])
     return juicy_center(ws)
 
-def watershed_sequence(a, seeds=None, mask=None, axis=0, **kwargs):
+def watershed_sequence(a, seeds=None, mask=None, axis=0, n_jobs=1, **kwargs):
     """Perform a watershed on a plane-by-plane basis.
 
     See documentation for `watershed` for available kwargs.
@@ -294,6 +295,9 @@ def watershed_sequence(a, seeds=None, mask=None, axis=0, **kwargs):
         Which axis defines the plane sequence. For example, if the input image
         is 3D and axis=1, then the output will be the watershed on a[:, 0, :], 
         a[:, 1, :], a[:, 2, :], ... and so on.
+    n_jobs : int, optional
+        Use joblib to distribute each plane over given number of processing
+        cores. If -1, `multiprocessing.cpu_count` is used.
 
     Returns
     -------
@@ -314,8 +318,9 @@ def watershed_sequence(a, seeds=None, mask=None, axis=0, **kwargs):
         seeds = it.repeat(None)
     if mask is None:
         mask = it.repeat(None)
-    ws = [watershed(i, seeds=s, mask=m, **kwargs)
-                                        for i, s, m in zip(a, seeds, mask)]
+    ws = joblib.Parallel(n_jobs=n_jobs)(
+        joblib.delayed(watershed)(i, seeds=s, mask=m, **kwargs)
+        for i, s, m in zip(a, seeds, mask))
     counts = list(map(np.max, ws[:-1]))
     counts = np.concatenate((np.array([0]), counts))
     counts = np.cumsum(counts)