Merge pull request #1417 from MouseLand/refactor_eval

Refactor eval

Author: carsen-stringer

cellpose/models.py

@@ -268,17 +268,9 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
             x = x[np.newaxis, ...]
         nimg = x.shape[0]
 
-        image_scaling = None
-        Ly_0 = x.shape[1]
-        Lx_0 = x.shape[2]
-        Lz_0 = None
-        if do_3D or stitch_threshold > 0:
-            Lz_0 = x.shape[0]
-        if diameter is not None:
+        image_scaling = 1.0
+        if diameter is not None and diameter > 0:
             image_scaling = 30. / diameter
-            x = transforms.resize_image(x,
-                                        Ly=int(x.shape[1] * image_scaling),
-                                        Lx=int(x.shape[2] * image_scaling))
 
 
         # normalize image
@@ -306,40 +298,17 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
         if do_normalization:
             x = transforms.normalize_img(x, **normalize_params)
 
-        # ajust the anisotropy when diameter is specified and images are resized:
-        if isinstance(anisotropy, (float, int)) and image_scaling:
-            anisotropy = image_scaling * anisotropy
-
         dP, cellprob, styles = self._run_net(
-            x, 
+            x,
+            resample=resample,
+            rescale=image_scaling,
             augment=augment, 
             batch_size=batch_size, 
             tile_overlap=tile_overlap, 
             bsize=bsize,
             do_3D=do_3D, 
             anisotropy=anisotropy)
 
-        if do_3D:
-            torch.cuda.empty_cache()
-            gc.collect()
-
-            if resample:
-                # resize XY then YZ and then put channels first
-                dP = transforms.resize_image(dP.transpose(1, 2, 3, 0), Ly=Ly_0, Lx=Lx_0, no_channels=False)
-                dP = transforms.resize_image(dP.transpose(1, 0, 2, 3), Lx=Lx_0, Ly=Lz_0, no_channels=False)
-                dP = dP.transpose(3, 1, 0, 2)
-
-                # resize cellprob:
-                cellprob = transforms.resize_image(cellprob, Ly=Ly_0, Lx=Lx_0, no_channels=True)
-                cellprob = transforms.resize_image(cellprob.transpose(1, 0, 2), Lx=Lx_0, Ly=Lz_0, no_channels=True)
-                cellprob = cellprob.transpose(1, 0, 2)
-
-        # 2d case:
-        if resample and not do_3D:
-            # 2D images have N = 1 in batch dimension:
-            dP = transforms.resize_image(dP.transpose(1, 2, 3, 0), Ly=Ly_0, Lx=Lx_0, no_channels=False).transpose(3, 0, 1, 2)
-            cellprob = transforms.resize_image(cellprob, Ly=Ly_0, Lx=Lx_0, no_channels=True)
-        
         if do_3D and flow3D_smooth:
             if isinstance(flow3D_smooth, (int, float)):
                 flow3D_smooth = [flow3D_smooth]*3 
@@ -350,15 +319,21 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
                 dP = gaussian_filter(dP, [0, *flow3D_smooth])
             else: 
                 models_logger.warning(f"Could not do flow smoothing with {flow3D_smooth} either because its len was not 3 or no items were > 0, skipping flow3D_smoothing")
+            torch.cuda.empty_cache()
+            gc.collect()
 
         if compute_masks:
             # use user niter if specified, otherwise scale niter (200) with diameter
             niter_scale = 1 if image_scaling is None else image_scaling
             niter = int(200/niter_scale) if niter is None or niter == 0 else niter
-            masks = self._compute_masks((Lz_0 or nimg, Ly_0, Lx_0), dP, cellprob, flow_threshold=flow_threshold,
-                            cellprob_threshold=cellprob_threshold, min_size=min_size,
-                        max_size_fraction=max_size_fraction, niter=niter,
-                        stitch_threshold=stitch_threshold, do_3D=do_3D)
+            masks = self._compute_masks(x.shape, dP, cellprob, 
+                                        flow_threshold=flow_threshold,
+                                        cellprob_threshold=cellprob_threshold, 
+                                        min_size=min_size,
+                                        max_size_fraction=max_size_fraction, 
+                                        niter=niter,
+                                        stitch_threshold=stitch_threshold, 
+                                        do_3D=do_3D)
         else:
             masks = np.zeros(0) #pass back zeros if not compute_masks
 
@@ -368,38 +343,54 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
 
 
     def _run_net(self, x, 
-                augment=False, 
-                batch_size=8, tile_overlap=0.1,
-                bsize=256, anisotropy=1.0, do_3D=False):
+                 rescale=1.0,
+                 resample=True,
+                 augment=False, 
+                 batch_size=8, 
+                 tile_overlap=0.1,
+                 bsize=256, 
+                 anisotropy=1.0, 
+                 do_3D=False):
         """ run network on image x """
         tic = time.time()
         shape = x.shape
         nimg = shape[0]
 
-
         if do_3D:
             Lz, Ly, Lx = shape[:-1]
-            if anisotropy is not None and anisotropy != 1.0:
+            if rescale != 1.0 or (anisotropy is not None and anisotropy != 1.0):
                 models_logger.info(f"resizing 3D image with anisotropy={anisotropy}")
+                anisotropy = 1.0 if anisotropy is None else anisotropy
+                if rescale != 1.0:
+                    x = transforms.resize_image(x, Ly=int(Ly*rescale),
+                                                Lx=int(Lx*rescale))
                 x = transforms.resize_image(x.transpose(1,0,2,3),
-                                        Ly=int(Lz*anisotropy), 
-                                        Lx=int(Lx)).transpose(1,0,2,3)
+                                        Ly=int(Lz*anisotropy*rescale), 
+                                        Lx=int(Lx*rescale)).transpose(1,0,2,3)
             yf, styles = run_3D(self.net, x,
                                 batch_size=batch_size, augment=augment,  
                                 tile_overlap=tile_overlap, 
                                 bsize=bsize
                                 )
+            if resample:
+                if rescale != 1.0 or Lz != yf.shape[0]:
+                    models_logger.info("resizing 3D flows and cellprobl to original image size")
+                    if rescale != 1.0:
+                        yf = transforms.resize_image(yf, Ly=Ly, Lx=Lx)
+                    if Lz != yf.shape[0]:
+                        yf = transforms.resize_image(yf.transpose(1, 0, 2, 3), Ly=Lz, Lx=Lx).transpose(1, 0, 2, 3)
             cellprob = yf[..., -1]
             dP = yf[..., :-1].transpose((3, 0, 1, 2))
         else:
             yf, styles = run_net(self.net, x, bsize=bsize, augment=augment,
                                 batch_size=batch_size,  
                                 tile_overlap=tile_overlap, 
-                                )
+                                rsz=rescale if rescale !=1.0 else None)
+            if resample:
+                if rescale != 1.0:
+                    yf = transforms.resize_image(yf, shape[1], shape[2])
             cellprob = yf[..., -1]
             dP = yf[..., -3:-1].transpose((3, 0, 1, 2))
-            if yf.shape[-1] > 3:
-                styles = yf[..., :-3]
 
         styles = styles.squeeze()