diff --git a/active_plugins/runcellpose.py b/active_plugins/runcellpose.py index e838af9..1b7da41 100644 --- a/active_plugins/runcellpose.py +++ b/active_plugins/runcellpose.py @@ -6,9 +6,13 @@ import numpy import os -from cellpose import models, io, core, utils -from skimage.transform import resize +import skimage import importlib.metadata +import subprocess +import uuid +import shutil +import logging +import sys ################################# # @@ -19,10 +23,11 @@ from cellprofiler_core.image import Image from cellprofiler_core.module.image_segmentation import ImageSegmentation from cellprofiler_core.object import Objects -from cellprofiler_core.setting import Binary +from cellprofiler_core.setting import Binary, ValidationError from cellprofiler_core.setting.choice import Choice from cellprofiler_core.setting.do_something import DoSomething from cellprofiler_core.setting.subscriber import ImageSubscriber +from cellprofiler_core.preferences import get_default_output_directory from cellprofiler_core.setting.text import ( Integer, ImageName, @@ -34,7 +39,7 @@ CUDA_LINK = "https://pytorch.org/get-started/locally/" Cellpose_link = " https://doi.org/10.1038/s41592-020-01018-x" Omnipose_link = "https://doi.org/10.1101/2021.11.03.467199" -cellpose_ver = importlib.metadata.version("cellpose") +LOGGER = logging.getLogger(__name__) __doc__ = f"""\ RunCellpose @@ -78,8 +83,11 @@ """ -model_dic = models.MODEL_NAMES -model_dic.append("custom") +CELLPOSE_DOCKER_NO_PRETRAINED = "cellprofiler/runcellpose_no_pretrained:0.1" +CELLPOSE_DOCKER_IMAGE_WITH_PRETRAINED = "cellprofiler/runcellpose_with_pretrained:0.1" + +MODEL_NAMES = ['cyto','nuclei','tissuenet','livecell', 'cyto2', 'general', + 'CP', 'CPx', 'TN1', 'TN2', 'TN3', 'LC1', 'LC2', 'LC3', 'LC4', 'custom'] class RunCellpose(ImageSegmentation): @@ -87,7 +95,7 @@ class RunCellpose(ImageSegmentation): module_name = "RunCellpose" - variable_revision_number = 3 + variable_revision_number = 4 doi = { "Please cite the following when using RunCellPose:": "https://doi.org/10.1038/s41592-020-01018-x", @@ -97,9 +105,42 @@ class RunCellpose(ImageSegmentation): def create_settings(self): super(RunCellpose, self).create_settings() + self.docker_or_python = Choice( + text="Run CellPose in docker or local python environment", + choices=["Docker", "Python"], + value="Docker", + doc="""\ +If Docker is selected, ensure that Docker Desktop is open and running on your +computer. On first run of the RunCellpose plugin, the Docker container will be +downloaded. However, this slow downloading process will only have to happen +once. + +If Python is selected, the Python environment in which CellProfiler and Cellpose +are installed will be used. +""", + ) + + self.docker_image = Choice( + text="Select Cellpose docker image", + choices=[CELLPOSE_DOCKER_IMAGE_WITH_PRETRAINED, CELLPOSE_DOCKER_NO_PRETRAINED], + value=CELLPOSE_DOCKER_IMAGE_WITH_PRETRAINED, + doc="""\ +Select which Docker image to use for running Cellpose. + +If you are not using a custom model, you can select +**"{CELLPOSE_DOCKER_IMAGE_WITH_PRETRAINED}"**. If you are using a custom model, +you can use either **"{CELLPOSE_DOCKER_NO_PRETRAINED}"** or +**"{CELLPOSE_DOCKER_IMAGE_WITH_PRETRAINED}"**, but the latter will be slightly +larger (~500 MB) due to including all of the pretrained models. +""".format( + **{"CELLPOSE_DOCKER_NO_PRETRAINED": CELLPOSE_DOCKER_NO_PRETRAINED, + "CELLPOSE_DOCKER_IMAGE_WITH_PRETRAINED": CELLPOSE_DOCKER_IMAGE_WITH_PRETRAINED} +), + ) + self.expected_diameter = Integer( text="Expected object diameter", - value=15, + value=30, minval=0, doc="""\ The average diameter of the objects to be detected. Setting this to 0 will attempt to automatically detect object size. @@ -113,8 +154,8 @@ def create_settings(self): self.mode = Choice( text="Detection mode", - choices=model_dic, - value="cyto2", + choices=MODEL_NAMES, + value=MODEL_NAMES[0], doc="""\ CellPose comes with models for detecting nuclei or cells. Alternatively, you can supply a custom-trained model generated using the command line or Cellpose GUI. Custom models can be useful if working with unusual cell types. @@ -154,7 +195,7 @@ def create_settings(self): self.use_averaging = Binary( text="Use averaging", - value=True, + value=False, doc="""\ If enabled, CellPose will run it's 4 inbuilt models and take a consensus to determine the results. If disabled, only a single model will be called to produce results. Disabling averaging is faster to run but less accurate.""", @@ -299,6 +340,8 @@ def set_directory_fn(path): def settings(self): return [ self.x_name, + self.docker_or_python, + self.docker_image, self.expected_diameter, self.mode, self.y_name, @@ -322,10 +365,15 @@ def settings(self): ] def visible_settings(self): - if float(cellpose_ver[0:3]) >= 0.6 and int(cellpose_ver[0]) < 2: - vis_settings = [self.mode, self.omni, self.x_name] - else: - vis_settings = [self.mode, self.x_name] + vis_settings = [self.docker_or_python] + + if self.docker_or_python.value == "Docker": + vis_settings += [self.docker_image] + + vis_settings += [self.mode, self.x_name] + + if self.docker_or_python.value == "Python": + vis_settings += [self.omni] if self.mode.value != "nuclei": vis_settings += [self.supply_nuclei] @@ -357,8 +405,9 @@ def visible_settings(self): vis_settings += [self.use_averaging, self.use_gpu] - if self.use_gpu.value: - vis_settings += [self.gpu_test, self.manual_GPU_memory_share] + if self.docker_or_python.value == 'Python': + if self.use_gpu.value: + vis_settings += [self.gpu_test, self.manual_GPU_memory_share] return vis_settings @@ -375,48 +424,16 @@ def validate_module(self, pipeline): "Failed to load custom file: %s " % model_path, self.model_file_name, ) - try: - model = models.CellposeModel( - pretrained_model=model_path, gpu=self.use_gpu.value - ) - except: - raise ValidationError( - "Failed to load custom model: %s " % model_path, - self.model_file_name, - ) + if self.docker_or_python.value == "Python": + try: + model = models.CellposeModel(pretrained_model=model_path, gpu=self.use_gpu.value) + except: + raise ValidationError( + "Failed to load custom model: %s " + % model_path, self.model_file_name, + ) def run(self, workspace): - if float(cellpose_ver[0:3]) >= 0.6 and int(cellpose_ver[0]) < 2: - if self.mode.value != "custom": - model = models.Cellpose( - model_type=self.mode.value, gpu=self.use_gpu.value - ) - else: - model_file = self.model_file_name.value - model_directory = self.model_directory.get_absolute_path() - model_path = os.path.join(model_directory, model_file) - model = models.CellposeModel( - pretrained_model=model_path, gpu=self.use_gpu.value - ) - - else: - if self.mode.value != "custom": - model = models.CellposeModel( - model_type=self.mode.value, gpu=self.use_gpu.value - ) - else: - model_file = self.model_file_name.value - model_directory = self.model_directory.get_absolute_path() - model_path = os.path.join(model_directory, model_file) - model = models.CellposeModel( - pretrained_model=model_path, gpu=self.use_gpu.value - ) - - if self.use_gpu.value and model.torch: - from torch import cuda - - cuda.set_per_process_memory_fraction(self.manual_GPU_memory_share.value) - x_name = self.x_name.value y_name = self.y_name.value images = workspace.image_set @@ -424,10 +441,11 @@ def run(self, workspace): dimensions = x.dimensions x_data = x.pixel_data anisotropy = 0.0 - if self.do_3D.value: anisotropy = x.spacing[0] / x.spacing[1] + diam = self.expected_diameter.value if self.expected_diameter.value > 0 else None + if x.multichannel: raise ValueError( "Color images are not currently supported. Please provide greyscale images." @@ -450,59 +468,142 @@ def run(self, workspace): else: channels = [0, 0] - diam = ( - self.expected_diameter.value if self.expected_diameter.value > 0 else None - ) + if self.docker_or_python.value == "Python": + from cellpose import models, io, core, utils + self.cellpose_ver = importlib.metadata.version('cellpose') + if float(self.cellpose_ver[0:3]) >= 0.6 and int(self.cellpose_ver[0])<2: + if self.mode.value != 'custom': + model = models.Cellpose(model_type= self.mode.value, + gpu=self.use_gpu.value) + else: + model_file = self.model_file_name.value + model_directory = self.model_directory.get_absolute_path() + model_path = os.path.join(model_directory, model_file) + model = models.CellposeModel(pretrained_model=model_path, gpu=self.use_gpu.value) - try: - if float(cellpose_ver[0:3]) >= 0.7 and int(cellpose_ver[0]) < 2: - y_data, flows, *_ = model.eval( - x_data, - channels=channels, - diameter=diam, - net_avg=self.use_averaging.value, - do_3D=self.do_3D.value, - anisotropy=anisotropy, - flow_threshold=self.flow_threshold.value, - cellprob_threshold=self.cellprob_threshold.value, - stitch_threshold=self.stitch_threshold.value, - min_size=self.min_size.value, - omni=self.omni.value, - invert=self.invert.value, - ) else: - y_data, flows, *_ = model.eval( - x_data, - channels=channels, - diameter=diam, - net_avg=self.use_averaging.value, - do_3D=self.do_3D.value, - anisotropy=anisotropy, - flow_threshold=self.flow_threshold.value, - cellprob_threshold=self.cellprob_threshold.value, - stitch_threshold=self.stitch_threshold.value, - min_size=self.min_size.value, - invert=self.invert.value, + if self.mode.value != 'custom': + model = models.CellposeModel(model_type= self.mode.value, + gpu=self.use_gpu.value) + else: + model_file = self.model_file_name.value + model_directory = self.model_directory.get_absolute_path() + model_path = os.path.join(model_directory, model_file) + model = models.CellposeModel(pretrained_model=model_path, gpu=self.use_gpu.value) + + if self.use_gpu.value and model.torch: + from torch import cuda + cuda.set_per_process_memory_fraction(self.manual_GPU_memory_share.value) + + try: + if float(self.cellpose_ver[0:3]) >= 0.7 and int(self.cellpose_ver[0])<2: + y_data, flows, *_ = model.eval( + x_data, + channels=channels, + diameter=diam, + net_avg=self.use_averaging.value, + do_3D=self.do_3D.value, + anisotropy=anisotropy, + flow_threshold=self.flow_threshold.value, + cellprob_threshold=self.cellprob_threshold.value, + stitch_threshold=self.stitch_threshold.value, + min_size=self.min_size.value, + omni=self.omni.value, + invert=self.invert.value, + ) + else: + y_data, flows, *_ = model.eval( + x_data, + channels=channels, + diameter=diam, + net_avg=self.use_averaging.value, + do_3D=self.do_3D.value, + anisotropy=anisotropy, + flow_threshold=self.flow_threshold.value, + cellprob_threshold=self.cellprob_threshold.value, + stitch_threshold=self.stitch_threshold.value, + min_size=self.min_size.value, + invert=self.invert.value, ) - if self.remove_edge_masks: - y_data = utils.remove_edge_masks(y_data) + if self.remove_edge_masks: + y_data = utils.remove_edge_masks(y_data) + + except Exception as a: + print(f"Unable to create masks. Check your module settings. {a}") + finally: + if self.use_gpu.value and model.torch: + # Try to clear some GPU memory for other worker processes. + try: + cuda.empty_cache() + except Exception as e: + print(f"Unable to clear GPU memory. You may need to restart CellProfiler to change models. {e}") + + elif self.docker_or_python.value == "Docker": + # Define how to call docker + docker_path = "docker" if sys.platform.lower().startswith("win") else "/usr/local/bin/docker" + # Create a UUID for this run + unique_name = str(uuid.uuid4()) + # Directory that will be used to pass images to the docker container + temp_dir = os.path.join(get_default_output_directory(), ".cellprofiler_temp", unique_name) + temp_img_dir = os.path.join(temp_dir, "img") + + os.makedirs(temp_dir, exist_ok=True) + os.makedirs(temp_img_dir, exist_ok=True) + + temp_img_path = os.path.join(temp_img_dir, unique_name+".tiff") + if self.mode.value == "custom": + model_file = self.model_file_name.value + model_directory = self.model_directory.get_absolute_path() + model_path = os.path.join(model_directory, model_file) + temp_model_dir = os.path.join(temp_dir, "model") + + os.makedirs(temp_model_dir, exist_ok=True) + # Copy the model + shutil.copy(model_path, os.path.join(temp_model_dir, model_file)) + + # Save the image to the Docker mounted directory + skimage.io.imsave(temp_img_path, x_data) + + cmd = f""" + {docker_path} run --rm -v {temp_dir}:/data + {self.docker_image.value} + {'--gpus all' if self.use_gpu.value else ''} + cellpose + --dir /data/img + {'--pretrained_model ' + self.mode.value if self.mode.value != 'custom' else '--pretrained_model /data/model/' + model_file} + --chan {channels[0]} + --chan2 {channels[1]} + --diameter {diam} + {'--net_avg' if self.use_averaging.value else ''} + {'--do_3D' if self.do_3D.value else ''} + --anisotropy {anisotropy} + --flow_threshold {self.flow_threshold.value} + --cellprob_threshold {self.cellprob_threshold.value} + --stitch_threshold {self.stitch_threshold.value} + --min_size {self.min_size.value} + {'--invert' if self.invert.value else ''} + {'--exclude_on_edges' if self.remove_edge_masks.value else ''} + --verbose + """ - y = Objects() - y.segmented = y_data + try: + subprocess.run(cmd.split(), text=True) + cellpose_output = numpy.load(os.path.join(temp_img_dir, unique_name + "_seg.npy"), allow_pickle=True).item() - except Exception as a: - print(f"Unable to create masks. Check your module settings. {a}") - finally: - if self.use_gpu.value and model.torch: - # Try to clear some GPU memory for other worker processes. + y_data = cellpose_output["masks"] + flows = cellpose_output["flows"] + finally: + # Delete the temporary files try: - cuda.empty_cache() - except Exception as e: - print( - f"Unable to clear GPU memory. You may need to restart CellProfiler to change models. {e}" - ) + shutil.rmtree(temp_dir) + except: + LOGGER.error("Unable to delete temporary directory, files may be in use by another program.") + LOGGER.error("Temp folder is subfolder {tempdir} in your Default Output Folder.\nYou may need to remove it manually.") + + y = Objects() + y.segmented = y_data y.parent_image = x.parent_image objects = workspace.object_set objects.add_objects(y, y_name) @@ -510,7 +611,7 @@ def run(self, workspace): if self.save_probabilities.value: # Flows come out sized relative to CellPose's inbuilt model size. # We need to slightly resize to match the original image. - size_corrected = resize(flows[2], y_data.shape) + size_corrected = skimage.transform.resize(flows[2], y_data.shape) prob_image = Image( size_corrected, parent_image=x.parent_image, @@ -571,10 +672,9 @@ def display(self, workspace, figure): def do_check_gpu(self): import importlib.util - - torch_installed = importlib.util.find_spec("torch") is not None - # if the old version of cellpose <2.0, then use istorch kwarg - if float(cellpose_ver[0:3]) >= 0.7 and int(cellpose_ver[0]) < 2: + torch_installed = importlib.util.find_spec('torch') is not None + #if the old version of cellpose <2.0, then use istorch kwarg + if float(self.cellpose_ver[0:3]) >= 0.7 and int(self.cellpose_ver[0])<2: GPU_works = core.use_gpu(istorch=torch_installed) else: # if new version of cellpose, use use_torch kwarg GPU_works = core.use_gpu(use_torch=torch_installed) @@ -595,4 +695,7 @@ def upgrade_settings(self, setting_values, variable_revision_number, module_name if variable_revision_number == 2: setting_values = setting_values + ["0.0", False, "15", "1.0", False, False] variable_revision_number = 3 + if variable_revision_number == 3: + setting_values = [setting_values[0]] + ["Python",CELLPOSE_DOCKER_IMAGE_WITH_PRETRAINED] + setting_values[1:] + variable_revision_number = 4 return setting_values, variable_revision_number diff --git a/dockerfiles/RunCellpose/Dockerfile b/dockerfiles/RunCellpose/Dockerfile new file mode 100644 index 0000000..c46a0ab --- /dev/null +++ b/dockerfiles/RunCellpose/Dockerfile @@ -0,0 +1,7 @@ +FROM pytorch/pytorch:1.13.0-cuda11.6-cudnn8-runtime + +RUN pip install cellpose==2.2 + +# Include if you wish the image to contain Cellpose pretrained models +COPY download_cellpose_models.py / +RUN python /download_cellpose_models.py diff --git a/dockerfiles/RunCellpose/download_cellpose_models.py b/dockerfiles/RunCellpose/download_cellpose_models.py new file mode 100644 index 0000000..249f5c8 --- /dev/null +++ b/dockerfiles/RunCellpose/download_cellpose_models.py @@ -0,0 +1,8 @@ +import cellpose +from cellpose.models import MODEL_NAMES + +for model in MODEL_NAMES: + for model_index in range(4): + model_name = cellpose.models.model_path(model, model_index) + if model in ("cyto", "nuclei", "cyto2"): + size_model_name = cellpose.models.size_model_path(model) \ No newline at end of file diff --git a/documentation/CP-plugins-documentation/RunCellPose.md b/documentation/CP-plugins-documentation/RunCellPose.md index 168fa7c..744e667 100644 --- a/documentation/CP-plugins-documentation/RunCellPose.md +++ b/documentation/CP-plugins-documentation/RunCellPose.md @@ -1,6 +1,10 @@ -# RunCellPose +# RunCellpose -## Using RunCellPose with a GPU +RunCellpose is one of the modules that has additional dependencies that are not packaged with the built CellProfiler. +Therefore, you must additionally download RunCellpose's dependencies. +See [Using Plugins](using_plugins.md) for more information. + +## Using RunCellpose with a GPU If you want to use a GPU to run the model (this is recommended for speed), you'll need a compatible version of PyTorch and a supported GPU. General instructions are available at this [link](https://pytorch.org/get-started/locally/). diff --git a/documentation/CP-plugins-documentation/_toc.yml b/documentation/CP-plugins-documentation/_toc.yml index 1ecd291..64f99d2 100644 --- a/documentation/CP-plugins-documentation/_toc.yml +++ b/documentation/CP-plugins-documentation/_toc.yml @@ -6,6 +6,8 @@ parts: - caption: Overview chapters: - file: using_plugins + sections: RunCellpose + - file: config_examples - file: supported_plugins - file: unsupported_plugins - file: contributing_plugins diff --git a/documentation/CP-plugins-documentation/using_plugins.md b/documentation/CP-plugins-documentation/using_plugins.md index 9c05ecd..d2fbfa6 100644 --- a/documentation/CP-plugins-documentation/using_plugins.md +++ b/documentation/CP-plugins-documentation/using_plugins.md @@ -14,11 +14,15 @@ Please report any installation issues or bugs related to plugins in the [CellPro If the plugin you would like to use does not have any additional dependencies outside of those required for running CellProfiler (this is most plugins), using plugins is very simple. See [Installing plugins without dependencies](#installing-plugins-without-dependencies). -If the plugin you would like to use has dependencies, you have two separate options for installation. -The first option requires building CellProfiler from source, but plugin installation is simpler. +If the plugin you would like to use has dependencies, you have three separate options for installation. +- The first option requires building CellProfiler from source, but plugin installation is simpler. See [Installing plugins with dependencies, using CellProfiler from source](#installing-plugins-with-dependencies-using-cellprofiler-from-source). -The second option allows you to use pre-built CellProfiler, but plugin installation is more complex. +- The second option allows you to use pre-built CellProfiler, but plugin installation is more complex. See [Installing plugins with dependencies, using pre-built CellProfiler](#installing-plugins-with-dependencies-using-pre-built-cellprofiler). +- The third option uses Docker to bypass installation requirements. +It is the simplest option that only requires download of Docker Desktop; the module that has dependencies will automatically download a Docker that has all of the dependencies upon run and access that Docker while running the plugin. +It is currently only supported for the RunCellpose plugin but will be available in other plugins soon. +See [Using Docker to Bypass Installation Requirements](#using-docker-to-bypass-installation-requirements). ### Installing plugins without dependencies @@ -157,4 +161,18 @@ These are all the folders you need to copy over: 7. **Open and use CellProfiler.** When you try to run your plugin in your pipeline, if you have missed copying over any specific requirements, it will give you an error message that will tell you what dependency is missing in the terminal window that opens with CellProfiler on Windows machines. -This information is not available in Mac machines. \ No newline at end of file +This information is not available in Mac machines. + +### Using Docker to bypass installation requirements + +1. **Download Docker** +Download Docker Desktop from [Docker.com](https://www.docker.com/products/docker-desktop/). + +2. **Run Docker Desktop** +Open Docker Desktop. +Docker Desktop will need to be open every time you use a plugin with Docker. + +3. **Select "Run with Docker"** +In your plugin, select `Docker` for "Run module in docker or local python environment" setting. +On the first run of the plugin, the Docker container will be downloaded, however, this slow downloading process will only have to happen +once. \ No newline at end of file