Merge branch 'deisseroth-lab-performance'

cellpose/dynamics.py

@@ -8,7 +8,7 @@
 from numba import njit, float32, int32, vectorize
 from . import utils, metrics
 
-@njit('(float64[:], int32[:], int32[:], int32, int32, int32, int32)')
+@njit('(float64[:], int32[:], int32[:], int32, int32, int32, int32)', nogil=True)
 def _extend_centers(T,y,x,ymed,xmed,Lx, niter):
     """ run diffusion from center of mask (ymed, xmed) on mask pixels (y, x)
 
@@ -166,7 +166,7 @@ def masks_to_flows(masks):
 
     return mu, mu_c
 
-@njit('(float32[:,:,:,:],float32[:,:,:,:], int32[:,:], int32)')
+@njit('(float32[:,:,:,:],float32[:,:,:,:], int32[:,:], int32)', nogil=True)
 def steps3D(p, dP, inds, niter):
     """ run dynamics of pixels to recover masks in 3D
     
@@ -207,7 +207,7 @@ def steps3D(p, dP, inds, niter):
             p[2,z,y,x] = min(shape[2]-1, max(0, p[2,z,y,x] - dP[2,p0,p1,p2]))
     return p
 
-@njit('(float32[:,:,:], float32[:,:,:], int32[:,:], int32)')
+@njit('(float32[:,:,:], float32[:,:,:], int32[:,:], int32)', nogil=True)
 def steps2D(p, dP, inds, niter):
     """ run dynamics of pixels to recover masks in 2D
     

cellpose/models.py

@@ -1,4 +1,4 @@
-import os, sys, time, shutil, tempfile, datetime, pathlib, gc
+import os, sys, time, shutil, tempfile, datetime, pathlib
 import numpy as np
 from tqdm import trange, tqdm
 from urllib.parse import urlparse
@@ -536,8 +536,7 @@ def _run_net(self, imgs, rsz=1.0, augment=True, tile=True, bsize=224):
             1D array summarizing the style of the image,
             if tiled it is averaged over tiles
 
-        """
-
+        """   
         shape = imgs.shape
         # rescale image for flow computation
         if not isinstance(rsz, list):
@@ -663,8 +662,6 @@ def _run_tiled(self, imgi, augment=True, bsize=224):
                         stylei = stylea[i*ntiles:(i+1)*ntiles].sum(axis=0)
                         stylei /= (stylei**2).sum()**0.5
                         styles.append(stylei)
-                del IMGa 
-                gc.collect()
             return yf, np.array(styles)
         else:
             IMG, ysub, xsub, Ly, Lx = transforms.make_tiles(imgi, bsize=bsize, augment=augment)
@@ -690,9 +687,6 @@ def _run_tiled(self, imgi, augment=True, bsize=224):
             yf = transforms.average_tiles(y, ysub, xsub, Ly, Lx)
             yf = yf[:,:imgi.shape[1],:imgi.shape[2]]
             styles /= (styles**2).sum()**0.5
-
-            del IMG
-            gc.collect()
             return yf, styles
 
     def train(self, train_data, train_labels, test_data=None, test_labels=None, channels=None,