File size: 5,138 Bytes
9060565
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
import sys, os, argparse, glob, pathlib, time
import numpy as np
from natsort import natsorted
from tqdm import tqdm
from cellpose import utils, models, io, core, version_str, transforms


def main():
    parser = argparse.ArgumentParser(description='cellpose parameters')

    input_img_args = parser.add_argument_group("input image arguments")
    input_img_args.add_argument('--dir', default=[], type=str,
                                help='folder containing data to run or train on.')
    input_img_args.add_argument(
        '--image_path', default=[], type=str, help=
        'if given and --dir not given, run on single image instead of folder (cannot train with this option)'
    )
    input_img_args.add_argument(
        '--look_one_level_down', action='store_true',
        help='run processing on all subdirectories of current folder')
    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=None, type=int,
        help='axis of image which corresponds to image channels')
    input_img_args.add_argument('--z_axis', default=None, type=int,
                                help='axis of image which corresponds to Z dimension')
    input_img_args.add_argument(
        '--chan', default=0, type=int, help=
        'channel to segment; 0: GRAY, 1: RED, 2: GREEN, 3: BLUE. Default: %(default)s')
    input_img_args.add_argument(
        '--chan2', default=0, type=int, help=
        'nuclear channel (if cyto, optional); 0: NONE, 1: RED, 2: GREEN, 3: BLUE. Default: %(default)s'
    )
    input_img_args.add_argument('--invert', action='store_true',
                                help='invert grayscale channel')
    input_img_args.add_argument(
        '--all_channels', action='store_true', help=
        'use all channels in image if using own model and images with special channels')
    input_img_args.add_argument("--anisotropy", required=False, default=1.0, type=float,
                                help="anisotropy of volume in 3D")
    

    # algorithm settings
    algorithm_args = parser.add_argument_group("algorithm arguments")
    algorithm_args.add_argument('--sharpen_radius', required=False, default=0.0,
                                type=float, help='high-pass filtering radius. Default: %(default)s')
    algorithm_args.add_argument('--tile_norm', required=False, default=0, type=int,
                                help='tile normalization block size. Default: %(default)s')
    algorithm_args.add_argument('--nimg_per_tif', required=False, default=10, type=int,
                                help='number of crops in XY to save per tiff. Default: %(default)s')
    algorithm_args.add_argument('--crop_size', required=False, default=512, type=int,
                                help='size of random crop to save. Default: %(default)s')

    args = parser.parse_args()

    # find images
    if len(args.img_filter) > 0:
        imf = args.img_filter
    else:
        imf = None

    if len(args.dir) > 0:
        image_names = io.get_image_files(args.dir, "_masks", imf=imf,
                                     look_one_level_down=args.look_one_level_down)
        dirname = args.dir
    else:
        if os.path.exists(args.image_path):
            image_names = [args.image_path]
            dirname = os.path.split(args.image_path)[0]
        else:
            raise ValueError(f"ERROR: no file found at {args.image_path}")
        
    np.random.seed(0)
    nimg_per_tif = args.nimg_per_tif
    crop_size = args.crop_size
    os.makedirs(os.path.join(dirname, 'train/'), exist_ok=True)
    pm = [(0, 1, 2, 3), (2, 0, 1, 3), (1, 0, 2, 3)]
    npm = ["YX", "ZY", "ZX"]
    for name in image_names:
        name0 = os.path.splitext(os.path.split(name)[-1])[0]
        img0 = io.imread(name)
        img0 = transforms.convert_image(img0, channels=[args.chan, args.chan2], channel_axis=args.channel_axis, z_axis=args.z_axis)
        for p in range(3):
            img = img0.transpose(pm[p]).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]]
            if args.anisotropy > 1.0 and p > 0:
                imgs = transforms.resize_image(imgs, Ly=int(args.anisotropy * Ly), Lx=Lx)
            for k, img in enumerate(imgs):
                if args.tile_norm:
                    img = transforms.normalize99_tile(img, blocksize=args.tile_norm)
                if args.sharpen_radius:
                    img = transforms.smooth_sharpen_img(img,
                                                        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())


if __name__ == '__main__':
    main()