Upload 9 files

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+img1.png filter=lfs diff=lfs merge=lfs -text
+img2.png filter=lfs diff=lfs merge=lfs -text
+img3.png filter=lfs diff=lfs merge=lfs -text

BasePredictor.py

@@ -0,0 +1,120 @@
+import torch
+import numpy as np
+import time, os
+import tifffile as tif
+
+from datetime import datetime
+from zipfile import ZipFile
+from pytz import timezone
+
+from train_tools.data_utils.transforms import get_pred_transforms
+
+
+class BasePredictor:
+    def __init__(
+        self,
+        model,
+        device,
+        input_path,
+        output_path,
+        make_submission=False,
+        exp_name=None,
+        algo_params=None,
+    ):
+        self.model = model
+        self.device = device
+        self.input_path = input_path
+        self.output_path = output_path
+        self.make_submission = make_submission
+        self.exp_name = exp_name
+
+        # Assign algoritm-specific arguments
+        if algo_params:
+            self.__dict__.update((k, v) for k, v in algo_params.items())
+
+        # Prepare inference environments
+        self._setups()
+
+    @torch.no_grad()
+    def conduct_prediction(self):
+        self.model.to(self.device)
+        self.model.eval()
+        total_time = 0
+        total_times = []
+
+        for img_name in self.img_names:
+            img_data = self._get_img_data(img_name)
+            img_data = img_data.to(self.device)
+
+            start = time.time()
+
+            pred_mask = self._inference(img_data)
+            pred_mask = self._post_process(pred_mask.squeeze(0).cpu().numpy())
+            self.write_pred_mask(
+                pred_mask, self.output_path, img_name, self.make_submission
+            )
+
+            end = time.time()
+
+            time_cost = end - start
+            total_times.append(time_cost)
+            total_time += time_cost
+            print(
+                f"Prediction finished: {img_name}; img size = {img_data.shape}; costing: {time_cost:.2f}s"
+            )
+
+        print(f"\n Total Time Cost: {total_time:.2f}s")
+
+        if self.make_submission:
+            fname = "%s.zip" % self.exp_name
+
+            os.makedirs("./submissions", exist_ok=True)
+            submission_path = os.path.join("./submissions", fname)
+
+            with ZipFile(submission_path, "w") as zipObj2:
+                pred_names = sorted(os.listdir(self.output_path))
+                for pred_name in pred_names:
+                    pred_path = os.path.join(self.output_path, pred_name)
+                    zipObj2.write(pred_path)
+
+            print("\n>>>>> Submission file is saved at: %s\n" % submission_path)
+
+        return time_cost
+
+    def write_pred_mask(self, pred_mask, output_dir, image_name, submission=False):
+
+        # All images should contain at least 5 cells
+        if submission:
+            if not (np.max(pred_mask) > 5):
+                print("[!Caution] Only %d Cells Detected!!!\n" % np.max(pred_mask))
+
+        file_name = image_name.split(".")[0]
+        file_name = file_name + "_label.tiff"
+        file_path = os.path.join(output_dir, file_name)
+
+        tif.imwrite(file_path, pred_mask, compression="zlib")
+
+    def _setups(self):
+        self.pred_transforms = get_pred_transforms()
+        os.makedirs(self.output_path, exist_ok=True)
+
+        now = datetime.now(timezone("Asia/Seoul"))
+        dt_string = now.strftime("%m%d_%H%M")
+        self.exp_name = (
+            self.exp_name + dt_string if self.exp_name is not None else dt_string
+        )
+
+        self.img_names = sorted(os.listdir(self.input_path))
+
+    def _get_img_data(self, img_name):
+        img_path = os.path.join(self.input_path, img_name)
+        img_data = self.pred_transforms(img_path)
+        img_data = img_data.unsqueeze(0)
+
+        return img_data
+
+    def _inference(self, img_data):
+        raise NotImplementedError
+
+    def _post_process(self, pred_mask):
+        raise NotImplementedError

MEDIARFormer.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn as nn
+
+from segmentation_models_pytorch import MAnet
+from segmentation_models_pytorch.base.modules import Activation
+
+__all__ = ["MEDIARFormer"]
+
+
+class MEDIARFormer(MAnet):
+    """MEDIAR-Former Model"""
+
+    def __init__(
+        self,
+        encoder_name="mit_b5",  # Default encoder
+        encoder_weights="imagenet",  # Pre-trained weights
+        decoder_channels=(1024, 512, 256, 128, 64),  # Decoder configuration
+        decoder_pab_channels=256,  # Decoder Pyramid Attention Block channels
+        in_channels=3,  # Number of input channels
+        classes=3,  # Number of output classes
+    ):
+        # Initialize the MAnet model with provided parameters
+        super().__init__(
+            encoder_name=encoder_name,
+            encoder_weights=encoder_weights,
+            decoder_channels=decoder_channels,
+            decoder_pab_channels=decoder_pab_channels,
+            in_channels=in_channels,
+            classes=classes,
+        )
+
+        # Remove the default segmentation head as it's not used in this architecture
+        self.segmentation_head = None
+
+        # Modify all activation functions in the encoder and decoder from ReLU to Mish
+        _convert_activations(self.encoder, nn.ReLU, nn.Mish(inplace=True))
+        _convert_activations(self.decoder, nn.ReLU, nn.Mish(inplace=True))
+
+        # Add custom segmentation heads for different segmentation tasks
+        self.cellprob_head = DeepSegmentationHead(
+            in_channels=decoder_channels[-1], out_channels=1
+        )
+        self.gradflow_head = DeepSegmentationHead(
+            in_channels=decoder_channels[-1], out_channels=2
+        )
+
+    def forward(self, x):
+        """Forward pass through the network"""
+        # Ensure the input shape is correct
+        self.check_input_shape(x)
+
+        # Encode the input and then decode it
+        features = self.encoder(x)
+        decoder_output = self.decoder(*features)
+
+        # Generate masks for cell probability and gradient flows
+        cellprob_mask = self.cellprob_head(decoder_output)
+        gradflow_mask = self.gradflow_head(decoder_output)
+
+        # Concatenate the masks for output
+        masks = torch.cat([gradflow_mask, cellprob_mask], dim=1)
+
+        return masks
+
+
+class DeepSegmentationHead(nn.Sequential):
+    """Custom segmentation head for generating specific masks"""
+
+    def __init__(
+        self, in_channels, out_channels, kernel_size=3, activation=None, upsampling=1
+    ):
+        # Define a sequence of layers for the segmentation head
+        layers = [
+            nn.Conv2d(
+                in_channels,
+                in_channels // 2,
+                kernel_size=kernel_size,
+                padding=kernel_size // 2,
+            ),
+            nn.Mish(inplace=True),
+            nn.BatchNorm2d(in_channels // 2),
+            nn.Conv2d(
+                in_channels // 2,
+                out_channels,
+                kernel_size=kernel_size,
+                padding=kernel_size // 2,
+            ),
+            nn.UpsamplingBilinear2d(scale_factor=upsampling)
+            if upsampling > 1
+            else nn.Identity(),
+            Activation(activation) if activation else nn.Identity(),
+        ]
+        super().__init__(*layers)
+
+
+def _convert_activations(module, from_activation, to_activation):
+    """Recursively convert activation functions in a module"""
+    for name, child in module.named_children():
+        if isinstance(child, from_activation):
+            setattr(module, name, to_activation)
+        else:
+            _convert_activations(child, from_activation, to_activation)

Predictor.py

@@ -0,0 +1,234 @@
+import torch
+import numpy as np
+import os, sys
+from monai.inferers import sliding_window_inference
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), "../../")))
+
+from BasePredictor import BasePredictor
+from utils import compute_masks
+
+__all__ = ["Predictor"]
+
+
+class Predictor(BasePredictor):
+    def __init__(
+        self,
+        model,
+        device,
+        input_path,
+        output_path,
+        make_submission=False,
+        exp_name=None,
+        algo_params=None,
+    ):
+        super(Predictor, self).__init__(
+            model,
+            device,
+            input_path,
+            output_path,
+            make_submission,
+            exp_name,
+            algo_params,
+        )
+        self.hflip_tta = HorizontalFlip()
+        self.vflip_tta = VerticalFlip()
+
+    @torch.no_grad()
+    def _inference(self, img_data):
+        """Conduct model prediction"""
+
+        img_data = img_data.to(self.device)
+        img_base = img_data
+        outputs_base = self._window_inference(img_base)
+        outputs_base = outputs_base.cpu().squeeze()
+        img_base.cpu()
+
+        if not self.use_tta:
+            pred_mask = outputs_base
+            return pred_mask
+
+        else:
+            # HorizontalFlip TTA
+            img_hflip = self.hflip_tta.apply_aug_image(img_data, apply=True)
+            outputs_hflip = self._window_inference(img_hflip)
+            outputs_hflip = self.hflip_tta.apply_deaug_mask(outputs_hflip, apply=True)
+            outputs_hflip = outputs_hflip.cpu().squeeze()
+            img_hflip = img_hflip.cpu()
+
+            # VertricalFlip TTA
+            img_vflip = self.vflip_tta.apply_aug_image(img_data, apply=True)
+            outputs_vflip = self._window_inference(img_vflip)
+            outputs_vflip = self.vflip_tta.apply_deaug_mask(outputs_vflip, apply=True)
+            outputs_vflip = outputs_vflip.cpu().squeeze()
+            img_vflip = img_vflip.cpu()
+
+            # Merge Results
+            pred_mask = torch.zeros_like(outputs_base)
+            pred_mask[0] = (outputs_base[0] + outputs_hflip[0] - outputs_vflip[0]) / 3
+            pred_mask[1] = (outputs_base[1] - outputs_hflip[1] + outputs_vflip[1]) / 3
+            pred_mask[2] = (outputs_base[2] + outputs_hflip[2] + outputs_vflip[2]) / 3
+
+        return pred_mask
+
+    def _window_inference(self, img_data, aux=False):
+        """Inference on RoI-sized window"""
+        outputs = sliding_window_inference(
+            img_data,
+            roi_size=512,
+            sw_batch_size=4,
+            predictor=self.model if not aux else self.model_aux,
+            padding_mode="constant",
+            mode="gaussian",
+            overlap=0.6,
+        )
+
+        return outputs
+
+    def _post_process(self, pred_mask):
+        """Generate cell instance masks."""
+        dP, cellprob = pred_mask[:2], self._sigmoid(pred_mask[-1])
+        H, W = pred_mask.shape[-2], pred_mask.shape[-1]
+
+        if np.prod(H * W) < (5000 * 5000):
+            pred_mask = compute_masks(
+                dP,
+                cellprob,
+                use_gpu=True,
+                flow_threshold=0.4,
+                device=self.device,
+                cellprob_threshold=0.5,
+            )[0]
+
+        else:
+            print("\n[Whole Slide] Grid Prediction starting...")
+            roi_size = 2000
+
+            # Get patch grid by roi_size
+            if H % roi_size != 0:
+                n_H = H // roi_size + 1
+                new_H = roi_size * n_H
+            else:
+                n_H = H // roi_size
+                new_H = H
+
+            if W % roi_size != 0:
+                n_W = W // roi_size + 1
+                new_W = roi_size * n_W
+            else:
+                n_W = W // roi_size
+                new_W = W
+
+            # Allocate values on the grid
+            pred_pad = np.zeros((new_H, new_W), dtype=np.uint32)
+            dP_pad = np.zeros((2, new_H, new_W), dtype=np.float32)
+            cellprob_pad = np.zeros((new_H, new_W), dtype=np.float32)
+
+            dP_pad[:, :H, :W], cellprob_pad[:H, :W] = dP, cellprob
+
+            for i in range(n_H):
+                for j in range(n_W):
+                    print("Pred on Grid (%d, %d) processing..." % (i, j))
+                    dP_roi = dP_pad[
+                        :,
+                        roi_size * i : roi_size * (i + 1),
+                        roi_size * j : roi_size * (j + 1),
+                    ]
+                    cellprob_roi = cellprob_pad[
+                        roi_size * i : roi_size * (i + 1),
+                        roi_size * j : roi_size * (j + 1),
+                    ]
+
+                    pred_mask = compute_masks(
+                        dP_roi,
+                        cellprob_roi,
+                        use_gpu=True,
+                        flow_threshold=0.4,
+                        device=self.device,
+                        cellprob_threshold=0.5,
+                    )[0]
+
+                    pred_pad[
+                        roi_size * i : roi_size * (i + 1),
+                        roi_size * j : roi_size * (j + 1),
+                    ] = pred_mask
+
+            pred_mask = pred_pad[:H, :W]
+
+        return pred_mask
+
+    def _sigmoid(self, z):
+        return 1 / (1 + np.exp(-z))
+
+
+"""
+Adapted from the following references:
+[1] https://github.com/qubvel/ttach/blob/master/ttach/transforms.py
+
+"""
+
+
+def hflip(x):
+    """flip batch of images horizontally"""
+    return x.flip(3)
+
+
+def vflip(x):
+    """flip batch of images vertically"""
+    return x.flip(2)
+
+
+class DualTransform:
+    identity_param = None
+
+    def __init__(
+        self, name: str, params,
+    ):
+        self.params = params
+        self.pname = name
+
+    def apply_aug_image(self, image, *args, **params):
+        raise NotImplementedError
+
+    def apply_deaug_mask(self, mask, *args, **params):
+        raise NotImplementedError
+
+
+class HorizontalFlip(DualTransform):
+    """Flip images horizontally (left -> right)"""
+
+    identity_param = False
+
+    def __init__(self):
+        super().__init__("apply", [False, True])
+
+    def apply_aug_image(self, image, apply=False, **kwargs):
+        if apply:
+            image = hflip(image)
+        return image
+
+    def apply_deaug_mask(self, mask, apply=False, **kwargs):
+        if apply:
+            mask = hflip(mask)
+        return mask
+
+
+class VerticalFlip(DualTransform):
+    """Flip images vertically (up -> down)"""
+
+    identity_param = False
+
+    def __init__(self):
+        super().__init__("apply", [False, True])
+
+    def apply_aug_image(self, image, apply=False, **kwargs):
+        if apply:
+            image = vflip(image)
+
+        return image
+
+    def apply_deaug_mask(self, mask, apply=False, **kwargs):
+        if apply:
+            mask = vflip(mask)
+
+        return mask

requirements.txt

@@ -0,0 +1,32 @@
+transformers 
+datasets 
+torch 
+torchvision 
+accelerate 
+opencv-python 
+pillow
+fastremap==1.14.1
+monai==1.3.0
+numba==0.57.1
+numpy==1.24.3
+pandas==2.0.3
+pytz==2023.3.post1
+scipy==1.12.0
+segmentation_models_pytorch==0.3.3
+tifffile==2023.4.12
+torch==2.1.2
+tqdm==4.65.0
+wandb==0.16.2
+scikit-image
+matplotlib
+segmentation-models-pytorch==0.3.1
+TensorFlow
+stardist
+csbdeep
+matplotlib
+scikit-image
+numpy
+scipy
+cellpose
+natsort
+

trained_model_200.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a15e5c6dfecd0b01c900ed7951aefb66d0c646eccb4bc12bc2fecb78715fcf6
+size 14952483

utils.py

@@ -0,0 +1,429 @@
+"""
+Copyright © 2022 Howard Hughes Medical Institute, 
+Authored by Carsen Stringer and Marius Pachitariu.
+
+Redistribution and use in source and binary forms, with or without 
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, 
+   this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, 
+   this list of conditions and the following disclaimer in the documentation 
+   and/or other materials provided with the distribution.
+
+3. Neither the name of HHMI nor the names of its contributors may be used to 
+   endorse or promote products derived from this software without specific 
+   prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 
+ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 
+LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
+CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 
+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 
+INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 
+CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 
+ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 
+POSSIBILITY OF SUCH DAMAGE.
+
+--------------------------------------------------------------------------
+MEDIAR Prediction uses CellPose's Gradient Flow Tracking.
+
+This code is adapted from the following codes:
+[1] https://github.com/MouseLand/cellpose/blob/main/cellpose/utils.py
+[2] https://github.com/MouseLand/cellpose/blob/main/cellpose/dynamics.py
+[3] https://github.com/MouseLand/cellpose/blob/main/cellpose/metrics.py
+"""
+
+import torch
+from torch.nn.functional import grid_sample
+import numpy as np
+import fastremap
+
+from skimage import morphology
+from scipy.ndimage import mean, find_objects
+from scipy.ndimage.filters import maximum_filter1d
+
+torch_GPU = torch.device("cuda")
+torch_CPU = torch.device("cpu")
+
+
+def labels_to_flows(labels, use_gpu=False, device=None, redo_flows=False):
+    """
+    Convert labels (list of masks or flows) to flows for training model
+    """
+
+    # Labels b x 1 x h x w
+    labels = labels.cpu().numpy().astype(np.int16)
+    nimg = len(labels)
+
+    if labels[0].ndim < 3:
+        labels = [labels[n][np.newaxis, :, :] for n in range(nimg)]
+
+    # Flows need to be recomputed
+    if labels[0].shape[0] == 1 or labels[0].ndim < 3 or redo_flows:
+        # compute flows; labels are fixed here to be unique, so they need to be passed back
+        # make sure labels are unique!
+        labels = [fastremap.renumber(label, in_place=True)[0] for label in labels]
+        veci = [
+            masks_to_flows(labels[n][0], use_gpu=use_gpu, device=device)
+            for n in range(nimg)
+        ]
+
+        # concatenate labels, distance transform, vector flows, heat (boundary and mask are computed in augmentations)
+        flows = [
+            np.concatenate((labels[n], labels[n] > 0.5, veci[n]), axis=0).astype(
+                np.float32
+            )
+            for n in range(nimg)
+        ]
+
+    return np.array(flows)
+
+
+def compute_masks(
+    dP,
+    cellprob,
+    p=None,
+    niter=200,
+    cellprob_threshold=0.4,
+    flow_threshold=0.4,
+    interp=True,
+    resize=None,
+    use_gpu=False,
+    device=None,
+):
+    """compute masks using dynamics from dP, cellprob, and boundary"""
+
+    cp_mask = cellprob > cellprob_threshold
+    cp_mask = morphology.remove_small_holes(cp_mask, area_threshold=16)
+    cp_mask = morphology.remove_small_objects(cp_mask, min_size=16)
+
+    if np.any(cp_mask):  # mask at this point is a cell cluster binary map, not labels
+        # follow flows
+        if p is None:
+            p, inds = follow_flows(
+                dP * cp_mask / 5.0,
+                niter=niter,
+                interp=interp,
+                use_gpu=use_gpu,
+                device=device,
+            )
+            if inds is None:
+                shape = resize if resize is not None else cellprob.shape
+                mask = np.zeros(shape, np.uint16)
+                p = np.zeros((len(shape), *shape), np.uint16)
+                return mask, p
+
+        # calculate masks
+        mask = get_masks(p, iscell=cp_mask)
+        
+        # flow thresholding factored out of get_masks
+        shape0 = p.shape[1:]
+        if mask.max() > 0 and flow_threshold is not None and flow_threshold > 0:
+            # make sure labels are unique at output of get_masks
+            mask = remove_bad_flow_masks(
+                mask, dP, threshold=flow_threshold, use_gpu=use_gpu, device=device
+            )
+        else:  # nothing to compute, just make it compatible
+            shape = resize if resize is not None else cellprob.shape
+            mask = np.zeros(shape, np.uint16)
+            p = np.zeros((len(shape), *shape), np.uint16)
+
+    return mask, p
+
+
+def _extend_centers_gpu(
+    neighbors, centers, isneighbor, Ly, Lx, n_iter=200, device=torch.device("cuda")
+):
+    if device is not None:
+        device = device
+    nimg = neighbors.shape[0] // 9
+    pt = torch.from_numpy(neighbors).to(device)
+
+    T = torch.zeros((nimg, Ly, Lx), dtype=torch.double, device=device)
+    meds = torch.from_numpy(centers.astype(int)).to(device).long()
+    isneigh = torch.from_numpy(isneighbor).to(device)
+    for i in range(n_iter):
+        T[:, meds[:, 0], meds[:, 1]] += 1
+        Tneigh = T[:, pt[:, :, 0], pt[:, :, 1]]
+        Tneigh *= isneigh
+        T[:, pt[0, :, 0], pt[0, :, 1]] = Tneigh.mean(axis=1)
+    del meds, isneigh, Tneigh
+    T = torch.log(1.0 + T)
+    # gradient positions
+    grads = T[:, pt[[2, 1, 4, 3], :, 0], pt[[2, 1, 4, 3], :, 1]]
+    del pt
+    dy = grads[:, 0] - grads[:, 1]
+    dx = grads[:, 2] - grads[:, 3]
+    del grads
+    mu_torch = np.stack((dy.cpu().squeeze(), dx.cpu().squeeze()), axis=-2)
+    return mu_torch
+
+
+def diameters(masks):
+    _, counts = np.unique(np.int32(masks), return_counts=True)
+    counts = counts[1:]
+    md = np.median(counts ** 0.5)
+    if np.isnan(md):
+        md = 0
+    md /= (np.pi ** 0.5) / 2
+    return md, counts ** 0.5
+
+
+def masks_to_flows_gpu(masks, device=None):
+    if device is None:
+        device = torch.device("cuda")
+
+    Ly0, Lx0 = masks.shape
+    Ly, Lx = Ly0 + 2, Lx0 + 2
+
+    masks_padded = np.zeros((Ly, Lx), np.int64)
+    masks_padded[1:-1, 1:-1] = masks
+
+    # get mask pixel neighbors
+    y, x = np.nonzero(masks_padded)
+    neighborsY = np.stack((y, y - 1, y + 1, y, y, y - 1, y - 1, y + 1, y + 1), axis=0)
+    neighborsX = np.stack((x, x, x, x - 1, x + 1, x - 1, x + 1, x - 1, x + 1), axis=0)
+    neighbors = np.stack((neighborsY, neighborsX), axis=-1)
+
+    # get mask centers
+    slices = find_objects(masks)
+
+    centers = np.zeros((masks.max(), 2), "int")
+    for i, si in enumerate(slices):
+        if si is not None:
+            sr, sc = si
+
+            ly, lx = sr.stop - sr.start + 1, sc.stop - sc.start + 1
+            yi, xi = np.nonzero(masks[sr, sc] == (i + 1))
+            yi = yi.astype(np.int32) + 1  # add padding
+            xi = xi.astype(np.int32) + 1  # add padding
+            ymed = np.median(yi)
+            xmed = np.median(xi)
+            imin = np.argmin((xi - xmed) ** 2 + (yi - ymed) ** 2)
+            xmed = xi[imin]
+            ymed = yi[imin]
+            centers[i, 0] = ymed + sr.start
+            centers[i, 1] = xmed + sc.start
+
+    # get neighbor validator (not all neighbors are in same mask)
+    neighbor_masks = masks_padded[neighbors[:, :, 0], neighbors[:, :, 1]]
+    isneighbor = neighbor_masks == neighbor_masks[0]
+    ext = np.array(
+        [[sr.stop - sr.start + 1, sc.stop - sc.start + 1] for sr, sc in slices]
+    )
+    n_iter = 2 * (ext.sum(axis=1)).max()
+    # run diffusion
+    mu = _extend_centers_gpu(
+        neighbors, centers, isneighbor, Ly, Lx, n_iter=n_iter, device=device
+    )
+
+    # normalize
+    mu /= 1e-20 + (mu ** 2).sum(axis=0) ** 0.5
+
+    # put into original image
+    mu0 = np.zeros((2, Ly0, Lx0))
+    mu0[:, y - 1, x - 1] = mu
+    mu_c = np.zeros_like(mu0)
+    return mu0, mu_c
+
+
+def masks_to_flows(masks, use_gpu=False, device=None):
+    if masks.max() == 0 or (masks != 0).sum() == 1:
+        # dynamics_logger.warning('empty masks!')
+        return np.zeros((2, *masks.shape), "float32")
+
+    if use_gpu:
+        if use_gpu and device is None:
+            device = torch_GPU
+        elif device is None:
+            device = torch_CPU
+        masks_to_flows_device = masks_to_flows_gpu
+
+    if masks.ndim == 3:
+        Lz, Ly, Lx = masks.shape
+        mu = np.zeros((3, Lz, Ly, Lx), np.float32)
+        for z in range(Lz):
+            mu0 = masks_to_flows_device(masks[z], device=device)[0]
+            mu[[1, 2], z] += mu0
+        for y in range(Ly):
+            mu0 = masks_to_flows_device(masks[:, y], device=device)[0]
+            mu[[0, 2], :, y] += mu0
+        for x in range(Lx):
+            mu0 = masks_to_flows_device(masks[:, :, x], device=device)[0]
+            mu[[0, 1], :, :, x] += mu0
+        return mu
+    elif masks.ndim == 2:
+        mu, mu_c = masks_to_flows_device(masks, device=device)
+        return mu
+
+    else:
+        raise ValueError("masks_to_flows only takes 2D or 3D arrays")
+
+
+def steps2D_interp(p, dP, niter, use_gpu=False, device=None):
+    shape = dP.shape[1:]
+    if use_gpu:
+        if device is None:
+            device = torch_GPU
+        shape = (
+            np.array(shape)[[1, 0]].astype("float") - 1
+        )  # Y and X dimensions (dP is 2.Ly.Lx), flipped X-1, Y-1
+        pt = (
+            torch.from_numpy(p[[1, 0]].T).float().to(device).unsqueeze(0).unsqueeze(0)
+        )  # p is n_points by 2, so pt is [1 1 2 n_points]
+        im = (
+            torch.from_numpy(dP[[1, 0]]).float().to(device).unsqueeze(0)
+        )  # covert flow numpy array to tensor on GPU, add dimension
+        # normalize pt between  0 and  1, normalize the flow
+        for k in range(2):
+            im[:, k, :, :] *= 2.0 / shape[k]
+            pt[:, :, :, k] /= shape[k]
+
+        # normalize to between -1 and 1
+        pt = pt * 2 - 1
+
+        # here is where the stepping happens
+        for t in range(niter):
+            # align_corners default is False, just added to suppress warning
+            dPt = grid_sample(im, pt, align_corners=False)
+
+            for k in range(2):  # clamp the final pixel locations
+                pt[:, :, :, k] = torch.clamp(
+                    pt[:, :, :, k] + dPt[:, k, :, :], -1.0, 1.0
+                )
+
+        # undo the normalization from before, reverse order of operations
+        pt = (pt + 1) * 0.5
+        for k in range(2):
+            pt[:, :, :, k] *= shape[k]
+
+        p = pt[:, :, :, [1, 0]].cpu().numpy().squeeze().T
+        return p
+
+    else:
+        assert print("ho")
+
+
+def follow_flows(dP, mask=None, niter=200, interp=True, use_gpu=True, device=None):
+    shape = np.array(dP.shape[1:]).astype(np.int32)
+    niter = np.uint32(niter)
+
+    p = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing="ij")
+    p = np.array(p).astype(np.float32)
+
+    inds = np.array(np.nonzero(np.abs(dP[0]) > 1e-3)).astype(np.int32).T
+
+    if inds.ndim < 2 or inds.shape[0] < 5:
+        return p, None
+
+    if not interp:
+        assert print("woo")
+
+    else:
+        p_interp = steps2D_interp(
+            p[:, inds[:, 0], inds[:, 1]], dP, niter, use_gpu=use_gpu, device=device
+        )
+        p[:, inds[:, 0], inds[:, 1]] = p_interp
+
+    return p, inds
+
+
+def flow_error(maski, dP_net, use_gpu=False, device=None):
+    if dP_net.shape[1:] != maski.shape:
+        print("ERROR: net flow is not same size as predicted masks")
+        return
+
+    # flows predicted from estimated masks
+    dP_masks = masks_to_flows(maski, use_gpu=use_gpu, device=device)
+    # difference between predicted flows vs mask flows
+    flow_errors = np.zeros(maski.max())
+    for i in range(dP_masks.shape[0]):
+        flow_errors += mean(
+            (dP_masks[i] - dP_net[i] / 5.0) ** 2,
+            maski,
+            index=np.arange(1, maski.max() + 1),
+        )
+
+    return flow_errors, dP_masks
+
+
+def remove_bad_flow_masks(masks, flows, threshold=0.4, use_gpu=False, device=None):
+    merrors, _ = flow_error(masks, flows, use_gpu, device)
+    badi = 1 + (merrors > threshold).nonzero()[0]
+    masks[np.isin(masks, badi)] = 0
+    return masks
+
+
+def get_masks(p, iscell=None, rpad=20):
+    pflows = []
+    edges = []
+    shape0 = p.shape[1:]
+    dims = len(p)
+
+    for i in range(dims):
+        pflows.append(p[i].flatten().astype("int32"))
+        edges.append(np.arange(-0.5 - rpad, shape0[i] + 0.5 + rpad, 1))
+
+    h, _ = np.histogramdd(tuple(pflows), bins=edges)
+    hmax = h.copy()
+    for i in range(dims):
+        hmax = maximum_filter1d(hmax, 5, axis=i)
+
+    seeds = np.nonzero(np.logical_and(h - hmax > -1e-6, h > 10))
+    Nmax = h[seeds]
+    isort = np.argsort(Nmax)[::-1]
+    for s in seeds:
+        s = s[isort]
+
+    pix = list(np.array(seeds).T)
+
+    shape = h.shape
+    if dims == 3:
+        expand = np.nonzero(np.ones((3, 3, 3)))
+    else:
+        expand = np.nonzero(np.ones((3, 3)))
+    for e in expand:
+        e = np.expand_dims(e, 1)
+
+    for iter in range(5):
+        for k in range(len(pix)):
+            if iter == 0:
+                pix[k] = list(pix[k])
+            newpix = []
+            iin = []
+            for i, e in enumerate(expand):
+                epix = e[:, np.newaxis] + np.expand_dims(pix[k][i], 0) - 1
+                epix = epix.flatten()
+                iin.append(np.logical_and(epix >= 0, epix < shape[i]))
+                newpix.append(epix)
+            iin = np.all(tuple(iin), axis=0)
+            for p in newpix:
+                p = p[iin]
+            newpix = tuple(newpix)
+            igood = h[newpix] > 2
+            for i in range(dims):
+                pix[k][i] = newpix[i][igood]
+            if iter == 4:
+                pix[k] = tuple(pix[k])
+
+    M = np.zeros(h.shape, np.uint32)
+    for k in range(len(pix)):
+        M[pix[k]] = 1 + k
+
+    for i in range(dims):
+        pflows[i] = pflows[i] + rpad
+    M0 = M[tuple(pflows)]
+
+    # remove big masks
+    uniq, counts = fastremap.unique(M0, return_counts=True)
+    big = np.prod(shape0) * 0.9
+    bigc = uniq[counts > big]
+    if len(bigc) > 0 and (len(bigc) > 1 or bigc[0] != 0):
+        M0 = fastremap.mask(M0, bigc)
+    fastremap.renumber(M0, in_place=True)  # convenient to guarantee non-skipped labels
+    M0 = np.reshape(M0, shape0)
+    return M0