diff --git a/cellpose/gui/gui.py b/cellpose/gui/gui.py index 955b9b88..7621aed7 100644 --- a/cellpose/gui/gui.py +++ b/cellpose/gui/gui.py @@ -2056,9 +2056,9 @@ def compute_segmentation(self, custom=False, model_name=None, load_model=True): do_3D = False if stitch_threshold > 0. else do_3D if self.restore == "filter": - data = self.stack_filtered.copy().squeeze() + data = self.stack_filtered.copy().squeeze(axis=0) else: - data = self.stack.copy().squeeze() + data = self.stack.copy().squeeze(axis=0) flow_threshold = self.segmentation_settings.flow_threshold cellprob_threshold = self.segmentation_settings.cellprob_threshold diff --git a/cellpose/gui/make_train.py b/cellpose/gui/make_train.py index add945e0..907a312e 100644 --- a/cellpose/gui/make_train.py +++ b/cellpose/gui/make_train.py @@ -1,7 +1,7 @@ import os, argparse import numpy as np from cellpose import io, transforms - +import cv2 def main(): parser = argparse.ArgumentParser(description='Make slices of XYZ image data for training. Assumes image is ZXYC unless specified otherwise using --channel_axis and --z_axis') @@ -19,7 +19,7 @@ def main(): input_img_args.add_argument('--img_filter', default=[], type=str, help='end string for images to run on') input_img_args.add_argument( - '--channel_axis', default=-1, type=int, + '--channel_axis', default=None, type=int, help='axis of image which corresponds to image channels') input_img_args.add_argument('--z_axis', default=0, type=int, help='axis of image which corresponds to Z dimension') @@ -37,6 +37,9 @@ def main(): 'deprecated') input_img_args.add_argument("--anisotropy", required=False, default=1.0, type=float, help="anisotropy of volume in 3D") + input_img_args.add_argument( + '--seg_masks', action='store_true', help= + 'use 3D masks saved in a _seg.npy file to create 2D _seg.npy files') # algorithm settings @@ -83,14 +86,30 @@ def main(): z_axis=args.z_axis, do_3D=True) except ValueError: print('Error converting image. Did you provide the correct --channel_axis and --z_axis ?') - + + masks0 = None + if args.seg_masks: + try: + masks0 = np.load(name.replace(".tif", "_seg.npy"), allow_pickle=True).item()["masks"].squeeze() + except FileNotFoundError: + print(f'Warning: no segmentation masks found for {name}') + for p in range(3): img = img0.transpose(pm[p]).copy() + if masks0 is not None: + masks = masks0.transpose(pm[p][:3]).copy() print(npm[p], img[0].shape) Ly, Lx = img.shape[1:3] - imgs = img[np.random.permutation(img.shape[0])[:args.nimg_per_tif]] + irand = np.random.permutation(img.shape[0])[:args.nimg_per_tif] + imgs = img[irand] + if masks0 is not None: + masks = masks[irand] if args.anisotropy > 1.0 and p > 0: imgs = transforms.resize_image(imgs, Ly=int(args.anisotropy * Ly), Lx=Lx) + if masks0 is not None: + masks = transforms.resize_image(masks, Ly=int(args.anisotropy*Ly), Lx=Lx, + no_channels=True, interpolation=cv2.INTER_NEAREST).astype(masks.dtype) + for k, img in enumerate(imgs): if args.tile_norm: img = transforms.normalize99_tile(img, blocksize=args.tile_norm) @@ -99,8 +118,21 @@ def main(): sharpen_radius=args.sharpen_radius) ly = 0 if Ly - crop_size <= 0 else np.random.randint(0, Ly - crop_size) lx = 0 if Lx - crop_size <= 0 else np.random.randint(0, Lx - crop_size) - io.imsave(os.path.join(dirname, f'train/{name0}_{npm[p]}_{k}.tif'), - img[ly:ly + args.crop_size, lx:lx + args.crop_size].squeeze()) + fname = os.path.join(dirname, f'train/{name0}_{npm[p]}_{k}.tif') + img_crop = img[ly:ly + crop_size, lx:lx + crop_size].squeeze() + if masks0 is not None: + masks_crop = masks[k][ly:ly + crop_size, lx:lx + crop_size].squeeze() + try: + from skimage.measure import label + masks_crop = label(masks_crop) + except ImportError: + print("skimage not found, cannot relabel masks. Run `pip install scikit-image` to relabel and save masks.") + flows = [np.zeros(masks_crop.shape, dtype="uint8"), + np.zeros((2, *masks_crop.shape), dtype="uint8"), + np.zeros(masks_crop.shape, dtype="float32")] + io.masks_flows_to_seg(img_crop, masks_crop, flows, fname) + + io.imsave(fname, img_crop) if __name__ == '__main__': diff --git a/docs/do3d.rst b/docs/do3d.rst index 1b54843c..a06cbac9 100644 --- a/docs/do3d.rst +++ b/docs/do3d.rst @@ -126,6 +126,8 @@ See the help message for more information: --all_channels use all channels in image if using own model and images with special channels --anisotropy ANISOTROPY anisotropy of volume in 3D + --seg_masks use 3D masks saved in a _seg.npy file to create 2D _seg.npy files + algorithm arguments: --sharpen_radius SHARPEN_RADIUS