python_main.py

# RUN CELLPOSE
#!/usr/bin/python3

from cellpose import models, io, plot
import subprocess
import argparse
import os
import shutil

def workflow_main():
    #set up of argparse arguments for processing command line input
    parser = argparse.ArgumentParser(description='cellpose parameters')
    parser.add_argument('--filedir', required=True,
                        default=[], type=str, help='folder containing data to run or train on')
    parser.add_argument('--savedir', required=False,
                        default=[], type=str, help='folder where you want segmented files to be saved')
    # settings for running cellpose
    parser.add_argument('--pretrained_model', required=True,
                        default='cyto', type=str, help='model to use (either "cyto" or "nuclei")')
    parser.add_argument('--chan', required=False,
                        default=0, type=int, help='channel to segment; 0: GRAY, 1: RED, 2: GREEN, 3: BLUE')
    parser.add_argument('--chan2', required=False,
                        default=0, type=int,
                        help='nuclear channel (if cyto, optional); 0: NONE, 1: RED, 2: GREEN, 3: BLUE')
    parser.add_argument('--diameter', required=False,
                        default=30., type=float, help='estimated cell diameter, that cellpose will use to rescale image to model used for training')
    parser.add_argument('--flow_threshold', required=False,
                        default=0.4, type=float,
                        help='parameter describes the maximum allowed error of the flows for each mask')
    parser.add_argument('--cellprob_threshold', required=False,
                        default=0.0, type=float,
                        help='Threshold that defines which pixels are incorporated into analysis. Lower to increase number of masks returned')
    parser.add_argument('--save_png', action='store_true', help='save masks as png')
    parser.add_argument('--save_tif', action='store_true', help='save masks as tif')

    #set up variable to call parser arguments
    args = parser.parse_args()

    #parses through the png and tif, makes saving tif the default option
    pngortif = "--save_tif"
    if args.save_png:
            pngortif = "--save_png"
    else:
            pngortif = "--save_tif"

    #runs cellpose using subprocess and the argparse inputs
    subprocess.run(['python3', '-m', 'cellpose',
                    '--dir', args.filedir,
                    '--pretrained_model', args.pretrained_model,
                    '--diameter', str(args.diameter),
                    '--flow_threshold', str(args.flow_threshold),
                    '--cellprob_threshold', str(args.cellprob_threshold),
                    '--chan', str(args.chan),
                    '--chan2', str(args.chan2),
                    pngortif])

    #for the case that the savedir tag is used
    if args.savedir:
        #constructs the directory to save the new files in if it doesn't already exist
        if not os.path.exists(args.savedir):
            os.makedirs(args.savedir)

        #the string parts of the cellpose generated files that need to be moved
        patterns = ['_seg.npy', '_cp_masks', '_cp_outlines']

        #lists files in original file directory, including the new files generated by Cellpose
        #adds files with the above strings to a separate list
        postfilenames = os.listdir(args.filedir)
        movefiles = []
        for file in postfilenames:
            for p in patterns:
                if p in file:
                    movefiles.append(file)

        #generate absolute file path for the files to be moved
        finalmove = []
        for file in movefiles:
            x = os.path.abspath(os.path.join(args.filedir, file))
            finalmove.append(x)

        #move files to the directory that the user dictates
        for f in finalmove:
            shutil.move(f, args.savedir)


if __name__ == '__main__':
    workflow_main()