updating docstrings to google format

Author: carsen-stringer

cellpose/core.py

@@ -80,8 +80,7 @@ def assign_device(use_torch=True, gpu=False, device=0):
         device (int or str, optional): The device index or name to be used. Defaults to 0.
 
     Returns:
-        torch.device: The assigned device.
-        bool: True if GPU is used, False otherwise.
+        torch.device, bool (True if GPU is used, False otherwise)
     """
 
     if isinstance(device, str):
@@ -212,9 +211,9 @@ def run_net(net, imgi, batch_size=8, augment=False, tile_overlap=0.1, bsize=224,
         bsize (int, optional): Size of tiles to use in pixels [bsize x bsize]. Defaults to 224.
 
     Returns:
-        y (np.ndarray): output of network, if tiled it is averaged in tile overlaps. Size of [Ly x Lx x 3] or [Lz x Ly x Lx x 3].
-            y[...,0] is Y flow; y[...,1] is X flow; y[...,2] is cell probability.
-        style (np.ndarray): 1D array of size 256 summarizing the style of the image, if tiled it is averaged over tiles.
+        Tuple[numpy.ndarray, numpy.ndarray]: outputs of network y and style. If tiled `y` is averaged in tile overlaps. Size of [Ly x Lx x 3] or [Lz x Ly x Lx x 3].
+            y[...,0] is Y flow; y[...,1] is X flow; y[...,2] is cell probability. 
+            style is a 1D array of size 256 summarizing the style of the image, if tiled `style` is averaged over tiles.
     """
     # run network
     nout = net.nout
@@ -300,9 +299,9 @@ def run_3D(net, imgs, batch_size=8, augment=False,
         progress (QProgressBar, optional): pyqt progress bar. Defaults to None.
 
     Returns:
-        y (np.ndarray): output of network, if tiled it is averaged in tile overlaps. Size of [Ly x Lx x 3] or [Lz x Ly x Lx x 3].
-            y[...,0] is Y flow; y[...,1] is X flow; y[...,2] is cell probability.
-        style (np.ndarray): 1D array of size 256 summarizing the style of the image, if tiled it is averaged over tiles.
+        Tuple[numpy.ndarray, numpy.ndarray]: outputs of network y and style. If tiled `y` is averaged in tile overlaps. Size of [Ly x Lx x 3] or [Lz x Ly x Lx x 3].
+            y[...,0] is Z flow; y[...,1] is Y flow; y[...,2] is X flow; y[...,3] is cell probability. 
+            style is a 1D array of size 256 summarizing the style of the image, if tiled `style` is averaged over tiles.
     """
     sstr = ["YX", "ZY", "ZX"]
     pm = [(0, 1, 2, 3), (1, 0, 2, 3), (2, 0, 1, 3)]

cellpose/denoise.py

@@ -554,10 +554,10 @@ def eval(self, x, batch_size=8, channels=None, channel_axis=None, z_axis=None,
             interp (bool, optional): interpolate during 2D dynamics (not available in 3D) . Defaults to True.
             
         Returns:
-            masks (list, np.ndarray): labelled image(s), where 0=no masks; 1,2,...=mask labels
-            flows (list): list of lists: flows[k][0] = XY flow in HSV 0-255; flows[k][1] = XY(Z) flows at each pixel; flows[k][2] = cell probability (if > cellprob_threshold, pixel used for dynamics); flows[k][3] = final pixel locations after Euler integration 
-            styles (list, np.ndarray): style vector summarizing each image of size 256.
-            imgs (list of 2D/3D arrays): Restored images
+            A tuple containing (masks, flows, styles, imgs); masks: labelled image(s), where 0=no masks; 1,2,...=mask labels; 
+            flows: list of lists: flows[k][0] = XY flow in HSV 0-255; flows[k][1] = XY(Z) flows at each pixel; flows[k][2] = cell probability (if > cellprob_threshold, pixel used for dynamics); flows[k][3] = final pixel locations after Euler integration;
+            styles: style vector summarizing each image of size 256;
+            imgs: Restored images.
         """
 
         if isinstance(normalize, dict):
@@ -729,9 +729,9 @@ def eval(self, x, batch_size=8, channels=None, channel_axis=None, z_axis=None,
             diameter (float, optional):  diameter for each image, 
                 if diameter is None, set to diam_mean or diam_train if available. Defaults to None.
             tile_overlap (float, optional): fraction of overlap of tiles when computing flows. Defaults to 0.1.
-            
+              
         Returns:
-            imgs (list of 2D/3D arrays): Restored images
+            list: A list of 2D/3D arrays of restored images
 
         """
         if isinstance(x, list) or x.squeeze().ndim == 5:
@@ -836,7 +836,7 @@ def _eval(self, net, x, batch_size=8, normalize=True, rescale=None,
             tile_overlap (float, optional): fraction of overlap of tiles when computing flows. Defaults to 0.1.
             
         Returns:
-            imgs (list of 2D/3D arrays): Restored images
+            list: A list of 2D/3D arrays of restored images
 
         """
         if isinstance(normalize, dict):

cellpose/dynamics.py

@@ -105,9 +105,7 @@ def get_centers(masks, slices):
         slices (ndarray): The slices of the masks.
 
     Returns:
-        tuple containing
-            - centers (ndarray): The centers of the masks.
-            - ext (ndarray): The extents of the masks.
+        A tuple containing the centers of the masks and the extents of the masks.
     """
     centers = np.zeros((len(slices), 2), "int32")
     ext = np.zeros((len(slices),), "int32")
@@ -131,16 +129,20 @@ def get_centers(masks, slices):
 def masks_to_flows_gpu(masks, device=torch.device("cpu"), niter=None):
     """Convert masks to flows using diffusion from center pixel.
 
-    Center of masks where diffusion starts is defined using COM.
+    Center of masks where diffusion starts is defined by pixel closest to median within the mask.
 
     Args:
         masks (int, 2D or 3D array): Labelled masks. 0=NO masks; 1,2,...=mask labels.
+        device (torch.device, optional): The device to run the computation on. Defaults to torch.device("cpu").
+        niter (int, optional): Number of iterations for the diffusion process. Defaults to None.
 
     Returns:
-        tuple containing
-            - mu (float, 3D or 4D array): Flows in Y = mu[-2], flows in X = mu[-1].
-                If masks are 3D, flows in Z = mu[0].
-            - meds_p (float, 2D or 3D array): cell centers
+        np.ndarray: A 4D array representing the flows for each pixel in Z, X, and Y.
+       
+
+    Returns:
+        A tuple containing (mu, meds_p). mu is float 3D or 4D array of flows in (Z)XY. 
+        meds_p are cell centers.
     """
     if device is None:
         device = torch.device('cuda') if torch.cuda.is_available() else torch.device('mps') if torch.backends.mps.is_available() else None
@@ -200,11 +202,12 @@ def masks_to_flows_gpu_3d(masks, device=None, niter=None):
 
     Args:
         masks (int, 2D or 3D array): Labelled masks. 0=NO masks; 1,2,...=mask labels.
+        device (torch.device, optional): The device to run the computation on. Defaults to None.
+        niter (int, optional): Number of iterations for the diffusion process. Defaults to None.
 
     Returns:
-        tuple containing
-            - mu (float, 3D or 4D array): Flows in Y = mu[-2], flows in X = mu[-1]. If masks are 3D, flows in Z = mu[0].
-            - mu_c (float, 2D or 3D array): zeros
+        np.ndarray: A 4D array representing the flows for each pixel in Z, X, and Y.
+        
     """
     if device is None:
         device = torch.device('cuda') if torch.cuda.is_available() else torch.device('mps') if torch.backends.mps.is_available() else None
@@ -264,24 +267,22 @@ def masks_to_flows_gpu_3d(masks, device=None, niter=None):
     # put into original image
     mu0 = np.zeros((3, Lz0, Ly0, Lx0))
     mu0[:, z.cpu().numpy() - 1, y.cpu().numpy() - 1, x.cpu().numpy() - 1] = mu
-    mu_c = np.zeros_like(mu0)
-    return mu0, mu_c
+    return mu0
 
 
-def masks_to_flows_cpu(masks, device=None, niter=None):
+def masks_to_flows_cpu(masks, niter=None, device=None):
     """Convert masks to flows using diffusion from center pixel.
 
     Center of masks where diffusion starts is defined to be the closest pixel to the mean of all pixels that is inside the mask.
     Result of diffusion is converted into flows by computing the gradients of the diffusion density map.
 
     Args:
         masks (int, 2D or 3D array): Labelled masks 0=NO masks; 1,2,...=mask labels
-
+        niter (int, optional): Number of iterations for computing flows. Defaults to None.
+    
     Returns:
-        tuple containing
-            - mu (float, 3D or 4D array): Flows in Y = mu[-2], flows in X = mu[-1].
-                If masks are 3D, flows in Z = mu[0].
-            - meds (float, 2D or 3D array): cell centers
+        A tuple containing (mu, meds_p). mu is float 3D or 4D array of flows in (Z)XY. 
+        meds_p are cell centers.
     """
     Ly, Lx = masks.shape
     mu = np.zeros((2, Ly, Lx), np.float64)
@@ -327,8 +328,7 @@ def masks_to_flows(masks, device=torch.device("cpu"), niter=None):
         masks (int, 2D or 3D array): Labelled masks 0=NO masks; 1,2,...=mask labels
 
     Returns:
-        mu (float, 3D or 4D array): Flows in Y = mu[-2], flows in X = mu[-1].
-                If masks are 3D, flows in Z = mu[0].
+        np.ndarray: mu is float 3D or 4D array of flows in (Z)XY.
     """
     if masks.max() == 0:
         dynamics_logger.warning("empty masks!")
@@ -583,9 +583,8 @@ def follow_flows(dP, inds, niter=200, interp=True, device=torch.device("cpu")):
         device (torch.device, optional): Device to use for computation. Default is None.
 
     Returns:
-        tuple containing:
-            - p (np.ndarray): Final locations of each pixel after dynamics; [axis x Ly x Lx] or [axis x Lz x Ly x Lx].
-            - inds (np.ndarray): Indices of pixels used for dynamics; [axis x Ly x Lx] or [axis x Lz x Ly x Lx].
+        A tuple containing (p, inds): p (np.ndarray): Final locations of each pixel after dynamics; [axis x Ly x Lx] or [axis x Lz x Ly x Lx]; 
+        inds (np.ndarray): Indices of pixels used for dynamics; [axis x Ly x Lx] or [axis x Lz x Ly x Lx].
     """
     shape = np.array(dP.shape[1:]).astype(np.int32)
     ndim = len(inds)

cellpose/io.py

@@ -457,12 +457,8 @@ def load_train_test_data(train_dir, test_dir=None, image_filter=None,
         look_one_level_down (bool, optional): Whether to look for data in subdirectories of train_dir and test_dir. Defaults to False.
 
     Returns:
-        images (list): A list of training images.
-        labels (list): A list of labels corresponding to the training images.
-        image_names (list): A list of names of the training images.
-        test_images (list, optional): A list of testing images. None if test_dir is not provided.
-        test_labels (list, optional): A list of labels corresponding to the testing images. None if test_dir is not provided.
-        test_image_names (list, optional): A list of names of the testing images. None if test_dir is not provided.
+        images, labels, image_names, test_images, test_labels, test_image_names
+
     """
     images, labels, image_names = load_images_labels(train_dir, mask_filter,
                                                      image_filter, look_one_level_down)

cellpose/metrics.py

@@ -245,7 +245,7 @@ def flow_error(maski, dP_net, device=None):
         dP_net (np.ndarray, float): ND flows where dP_net.shape[1:] = maski.shape.
 
     Returns:
-        flow_errors (np.ndarray, float): Mean squared error between predicted flows and flows from masks.
+        A tuple containing (flow_errors, dP_masks): flow_errors (np.ndarray, float): Mean squared error between predicted flows and flows from masks; 
         dP_masks (np.ndarray, float): ND flows produced from the predicted masks.
     """
     if dP_net.shape[1:] != maski.shape:

cellpose/models.py

@@ -162,15 +162,12 @@ def eval(self, x, batch_size=8, channels=[0, 0], channel_axis=None, invert=False
             do_3D (bool, optional): Set to True to run 3D segmentation on 4D image input. Defaults to False.
 
         Returns:
-            tuple containing
-                - masks (list of 2D arrays or single 3D array): Labelled image, where 0=no masks; 1,2,...=mask labels.
-                - flows (list of lists 2D arrays or list of 3D arrays): 
-                    - flows[k][0] = XY flow in HSV 0-255
-                    - flows[k][1] = XY flows at each pixel
-                    - flows[k][2] = cell probability (if > cellprob_threshold, pixel used for dynamics)
-                    - flows[k][3] = final pixel locations after Euler integration
-                - styles (list of 1D arrays of length 256 or single 1D array): Style vector summarizing each image, also used to estimate size of objects in image.
-                - diams (list of diameters or float): List of diameters or float (if do_3D=True).
+            A tuple containing (masks, flows, styles, diams): masks (list of 2D arrays or single 3D array): Labelled image, where 0=no masks; 1,2,...=mask labels;
+            flows (list of lists 2D arrays or list of 3D arrays): flows[k][0] = XY flow in HSV 0-255; flows[k][1] = XY flows at each pixel; 
+            flows[k][2] = cell probability (if > cellprob_threshold, pixel used for dynamics); 
+            flows[k][3] = final pixel locations after Euler integration; 
+            styles (list of 1D arrays of length 256 or single 1D array): Style vector summarizing each image, also used to estimate size of objects in image;
+            diams (list of diameters or float): List of diameters or float (if do_3D=True).
 
         """
 
@@ -435,10 +432,11 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
             progress (QProgressBar, optional): pyqt progress bar. Defaults to None.
 
         Returns:
-            A tuple containing:
-                - masks (list, np.ndarray): labelled image(s), where 0=no masks; 1,2,...=mask labels
-                - flows (list): list of lists: flows[k][0] = XY flow in HSV 0-255; flows[k][1] = XY(Z) flows at each pixel; flows[k][2] = cell probability (if > cellprob_threshold, pixel used for dynamics); flows[k][3] = final pixel locations after Euler integration 
-                - styles (list, np.ndarray): style vector summarizing each image of size 256.
+            masks (list of 2D arrays or single 3D array): Labelled image, where 0=no masks; 1,2,...=mask labels;
+            flows (list of lists 2D arrays or list of 3D arrays): flows[k][0] = XY flow in HSV 0-255; flows[k][1] = XY flows at each pixel; 
+            flows[k][2] = cell probability (if > cellprob_threshold, pixel used for dynamics); 
+            flows[k][3] = final pixel locations after Euler integration; 
+            styles (list of 1D arrays of length 256 or single 1D array): Style vector summarizing each image, also used to estimate size of objects in image.
             
         """
         if isinstance(x, list) or x.squeeze().ndim == 5:
@@ -733,9 +731,9 @@ def eval(self, x, channels=None, channel_axis=None, normalize=True, invert=False
 
 
         Returns:
-            A tuple containing:
-                - diam (np.ndarray): Final estimated diameters from images x or styles style after running both steps.
-                - diam_style (np.ndarray): Estimated diameters from style alone.
+            A tuple containing (diam, diam_style):
+            diam (np.ndarray): Final estimated diameters from images x or styles style after running both steps;
+            diam_style (np.ndarray): Estimated diameters from style alone.
         """
         if isinstance(x, list):
             self.timing = []

cellpose/train.py

@@ -371,9 +371,8 @@ def train_seg(net, train_data=None, train_labels=None, train_files=None,
         model_name (str, optional): String - name of the network. Defaults to None.
 
     Returns:
-        Path: path to saved model weights
-        np.ndarray: training losses
-        np.ndarray: test losses
+        tuple: A tuple containing the path to the saved model weights, training losses, and test losses.
+       
     """
     device = net.device
 

docs/train.rst

@@ -1,6 +1,10 @@
 Training
 ---------------------------
 
+.. warning::
+    MPS support for pytorch is incomplete, and so training on Macs with MPS may give NaN's, 
+    if so please use the CPU instead
+
 At the beginning of training, cellpose computes the flow field representation for each 
 mask image (``dynamics.labels_to_flows``).