Merge pull request #709 from sophiamaedler/main

replace np.arrange with slice to improve performance
MouseLand · May 19, 2023 · aa16d39 · aa16d39
2 parents ca88d6a + 37a1037
commit aa16d39
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/cellpose/core.py b/cellpose/core.py
@@ -539,9 +539,9 @@ def _run_tiled(self, imgi, augment=False, bsize=224, tile_overlap=0.1, return_co
             nout = self.nclasses + 32*return_conv
             y = np.zeros((IMG.shape[0], nout, ly, lx))
             for k in range(niter):
-                irange = np.arange(batch_size*k, min(IMG.shape[0], batch_size*k+batch_size))
+                irange = slice(batch_size*k, min(IMG.shape[0], batch_size*k+batch_size))
                 y0, style = self.network(IMG[irange], return_conv=return_conv)
-                y[irange] = y0.reshape(len(irange), y0.shape[-3], y0.shape[-2], y0.shape[-1])
+                y[irange] = y0.reshape(irange.stop-irange.start, y0.shape[-3], y0.shape[-2], y0.shape[-1])
                 if k==0:
                     styles = style[0]
                 styles += style.sum(axis=0)