Skip to content

Commit 59bb394

Browse files
ErinWeisbartErinWeisbart
committed
support omnipose in docker
1 parent ecf0167 commit 59bb394

File tree

7 files changed

+84
-25
lines changed

7 files changed

+84
-25
lines changed

‎active_plugins/runcellpose.py‎

Lines changed: 56 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -81,14 +81,16 @@
8181
"""
8282

8383
"Select Cellpose Docker Image"
84-
CELLPOSE_DOCKERS = {'v2': ["cellprofiler/runcellpose_no_pretrained:2.3.2",
84+
CELLPOSE_DOCKERS = {'omnipose': ["gnodar01/runcellpose_omnipose_no_pretrained:0.0.1"], # TODO make official
85+
'v2': ["cellprofiler/runcellpose_no_pretrained:2.3.2",
8586
"cellprofiler/runcellpose_with_pretrained:2.3.2",
8687
"cellprofiler/runcellpose_with_pretrained:2.2"],
87-
'v3': ["erinweisbart/cellpose:3.1.1.2"], #TODO
88-
'v4': ["erinweisbart/cellpose:4.0.5"]} #TODO
88+
'v3': ["erinweisbart/cellpose:3.1.1.2"], # TODO make official
89+
'v4': ["erinweisbart/cellpose:4.0.5"]} # TODO make official
8990

9091
"Detection mode"
91-
MODEL_NAMES = {'v2':['cyto','nuclei','tissuenet','livecell', 'cyto2', 'general',
92+
MODEL_NAMES = {'omnipose':['bact_phase_omni','bact_fluor_omni','cyto2_omni','worm_omni','plant_omni','bact_phase_affinity','cyto','nuclei','custom'],
93+
'v2':['cyto','nuclei','tissuenet','livecell', 'cyto2', 'general',
9294
'CP', 'CPx', 'TN1', 'TN2', 'TN3', 'LC1', 'LC2', 'LC3', 'LC4', 'custom'],
9395
'v3':[ "cyto3", "nuclei", "cyto2_cp3", "tissuenet_cp3", "livecell_cp3", "yeast_PhC_cp3",
9496
"yeast_BF_cp3", "bact_phase_cp3", "bact_fluor_cp3", "deepbacs_cp3", "cyto2", "cyto", "custom"],
@@ -160,6 +162,16 @@ def create_settings(self):
160162
choices=['omnipose', 'v2', 'v3', 'v4'],
161163
value='v3',
162164
doc="Select the version of Cellpose you want to use.")
165+
166+
self.docker_image_omnipose = Choice(
167+
text="Select Cellpose docker image",
168+
choices=CELLPOSE_DOCKERS['omnipose'],
169+
value=CELLPOSE_DOCKERS['omnipose'][0],
170+
doc="""\
171+
Select which Docker image to use for running Cellpose.
172+
If you are not using a custom model, you should select a Docker image **with pretrained**. If you are using a custom model,
173+
you can use any of the available Dockers, but those with pretrained models will be slightly larger (~500 MB).
174+
""")
163175

164176
self.docker_image_v2 = Choice(
165177
text="Select Cellpose docker image",
@@ -211,7 +223,15 @@ def create_settings(self):
211223
Note that automatic diameter mode does not work when running on 3D images.
212224
""",
213225
)
214-
226+
self.mode_omnipose = Choice(
227+
text="Detection mode",
228+
choices=MODEL_NAMES['omnipose'],
229+
value=MODEL_NAMES['omnipose'][0],
230+
doc="""\
231+
CellPose comes with models for detecting nuclei or cells. Alternatively, you can supply a custom-trained model
232+
generated using the command line or Cellpose GUI. Custom models can be useful if working with unusual cell types.
233+
""",
234+
)
215235
self.mode_v2 = Choice(
216236
text="Detection mode",
217237
choices=MODEL_NAMES['v2'],
@@ -244,7 +264,7 @@ def create_settings(self):
244264
text="Use Omnipose for mask reconstruction",
245265
value=False,
246266
doc="""\
247-
If enabled, use omnipose mask recontruction features will be used (Omnipose installation required and CellPose >= 1.0) """,
267+
If enabled, use omnipose mask recontruction features will be used (Omnipose installation required and CellPose >= 1.0, <2) """,
248268
)
249269

250270
self.do_3D = Binary(
@@ -466,11 +486,13 @@ def settings(self):
466486
self.rescale,
467487
self.docker_or_python,
468488
self.cellpose_version,
489+
self.docker_image_omnipose,
469490
self.docker_image_v2,
470491
self.docker_image_v3,
471492
self.docker_image_v4,
472493
self.specify_diameter,
473494
self.expected_diameter,
495+
self.mode_omnipose,
474496
self.mode_v2,
475497
self.mode_v3,
476498
self.mode_v4,
@@ -500,23 +522,25 @@ def settings(self):
500522
]
501523

502524
def visible_settings(self):
503-
vis_settings = [self.rescale, self.cellpose_version]
525+
vis_settings = [self.rescale, self.cellpose_version,self.docker_or_python]
504526

505-
if self.cellpose_version.value == 'omnipose': # omnipose only supports Python, not Docker
506-
self.docker_or_python.value = "Python"
527+
if self.cellpose_version.value == 'omnipose':
507528
vis_settings += [self.omni]
508-
else:
509-
vis_settings += [self.docker_or_python]
510529

511530
if self.docker_or_python.value == "Docker":
512-
if self.cellpose_version.value == 'v2':
531+
if self.cellpose_version.value == 'omnipose':
532+
vis_settings += [self.docker_image_omnipose]
533+
elif self.cellpose_version.value == 'v2':
513534
vis_settings += [self.docker_image_v2]
514535
elif self.cellpose_version.value == 'v3':
515536
vis_settings += [self.docker_image_v3]
516537
elif self.cellpose_version.value == 'v4':
517538
vis_settings += [self.docker_image_v4]
518-
519-
if self.cellpose_version.value == 'v2':
539+
540+
if self.cellpose_version.value == 'omnipose':
541+
vis_settings += [self.mode_omnipose]
542+
self.mode = self.mode_omnipose
543+
elif self.cellpose_version.value == 'v2':
520544
vis_settings += [self.mode_v2]
521545
self.mode = self.mode_v2
522546
elif self.cellpose_version.value == 'v3':
@@ -542,19 +566,24 @@ def visible_settings(self):
542566
if self.specify_diameter.value:
543567
vis_settings += [self.expected_diameter]
544568

569+
if self.cellpose_version.value != 'omnipose':
570+
vis_settings += [
571+
self.cellprob_threshold,
572+
self.min_size,
573+
]
574+
545575
vis_settings += [
546-
self.cellprob_threshold,
547-
self.min_size,
548576
self.flow_threshold,
549577
self.y_name,
550578
self.save_probabilities,
551579
]
580+
552581
if self.save_probabilities.value:
553582
vis_settings += [self.probabilities_name]
554583
if self.docker_or_python.value == 'Python':
555584
vis_settings += [self.probability_rescale_setting]
556585

557-
if self.cellpose_version.value in ['v2','v3']:
586+
if self.cellpose_version.value in ['omnipose','v2','v3']:
558587
vis_settings += [self.invert]
559588

560589
vis_settings += [self.do_3D, self.stitch_threshold, self.remove_edge_masks]
@@ -609,7 +638,8 @@ def run(self, workspace):
609638
x_data = x.pixel_data
610639

611640
if self.cellpose_version.value == 'omnipose':
612-
self.mode = self.mode_v2
641+
self.mode = self.mode_omnipose
642+
self.docker_image = self.docker_image_omnipose
613643
self.denoise.value = False # Denoising only supported in v3
614644
if self.cellpose_version.value == 'v2':
615645
self.mode = self.mode_v2
@@ -901,18 +931,19 @@ def run(self, workspace):
901931
if self.cellpose_version.value == 'v3':
902932
if self.denoise.value:
903933
cmd += ['--restore_type', self.denoise_type.value]
904-
if self.cellpose_version.value in ['v2','v3']:
905-
cmd += ['--chan', str(channels[0]), '--chan2', str(channels[1]), '--diameter', str(diam)]
934+
if self.cellpose_version.value in ['omnipose','v2','v3']:
935+
cmd += ['--chan', str(channels[0]), '--chan2', str(channels[1]), '--diameter', str(self.expected_diameter)]
906936
if self.cellpose_version.value in ['v4']:
907937
if self.specify_diameter.value:
908938
cmd += ['--diameter', str(diam)]
909939
if self.use_averaging.value:
910940
cmd += ['--net_avg']
911941
if self.do_3D.value:
912942
cmd += ['--do_3D']
913-
cmd += ['--anisotropy', str(anisotropy), '--flow_threshold', str(self.flow_threshold.value), '--cellprob_threshold',
914-
str(self.cellprob_threshold.value), '--stitch_threshold', str(self.stitch_threshold.value), '--min_size', str(self.min_size.value)]
915-
if self.cellpose_version.value in ['v2','v3']:
943+
if self.cellpose_version.value != 'omnipose':
944+
cmd += ['--cellprob_threshold', str(self.cellprob_threshold.value), '--min_size', str(self.min_size.value)]
945+
cmd += ['--anisotropy', str(anisotropy), '--flow_threshold', str(self.flow_threshold.value), '--stitch_threshold', str(self.stitch_threshold.value)]
946+
if self.cellpose_version.value in ['omnipose','v2','v3']:
916947
if self.invert.value:
917948
cmd += ['--invert']
918949
if self.remove_edge_masks.value:
@@ -1115,8 +1146,8 @@ def upgrade_settings(self, setting_values, variable_revision_number, module_name
11151146
variable_revision_number = 6
11161147
if variable_revision_number == 6:
11171148
new_setting_values = setting_values[0:2]
1118-
new_setting_values += ['v3', setting_values[2], CELLPOSE_DOCKERS['v3'][0], CELLPOSE_DOCKERS['v4'][0], setting_values[3]]
1119-
new_setting_values += [False, setting_values[4], MODEL_NAMES['v3'][0], MODEL_NAMES['v4'][0]]
1149+
new_setting_values += ['v3', CELLPOSE_DOCKERS['omnipose'][0], setting_values[2], CELLPOSE_DOCKERS['v3'][0], CELLPOSE_DOCKERS['v4'][0], setting_values[3]]
1150+
new_setting_values += [False, MODEL_NAMES['omnipose'][0], setting_values[4], MODEL_NAMES['v3'][0], MODEL_NAMES['v4'][0]]
11201151
new_setting_values += [setting_values[5:], False, DENOISER_NAMES[0], False, "Preprocessed"]
11211152
setting_values = new_setting_values
11221153
variable_revision_number = 7

dockerfiles/RunCellpose/Dockerfile renamed to dockerfiles/RunCellpose_Cellposev2/Dockerfile

File renamed without changes.

dockerfiles/RunCellpose/download_cellpose_models.py renamed to dockerfiles/RunCellpose_Cellposev2/download_cellpose_models.py

File renamed without changes.

‎dockerfiles/RunCellpose_Cellposev3/Dockerfile‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
FROM pytorch/pytorch:1.13.0-cuda11.6-cudnn8-runtime
2+
3+
RUN pip install numpy==1.26.4 cellpose==3.1.1.2
4+
5+
# Include if you wish the image to contain Cellpose pretrained models
6+
COPY download_cellpose_models.py /
7+
RUN python /download_cellpose_models.py

‎dockerfiles/RunCellpose_Cellposev3/download_cellpose_models.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import cellpose
2+
from cellpose.models import MODEL_NAMES
3+
4+
for model in MODEL_NAMES:
5+
for model_index in range(4):
6+
model_name = cellpose.models.model_path(model, model_index)
7+
if model in ("cyto", "nuclei", "cyto2"):
8+
size_model_name = cellpose.models.size_model_path(model)

‎dockerfiles/RunCellpose_Cellposev4/Dockerfile‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
FROM pytorch/pytorch:1.13.0-cuda11.6-cudnn8-runtime
2+
3+
RUN pip install numpy==1.26.4 cellpose==4.0.6
4+
5+
# Include if you wish the image to contain Cellpose pretrained models
6+
COPY download_cellpose_models.py /
7+
RUN python /download_cellpose_models.py

‎dockerfiles/RunCellpose_Cellposev4/download_cellpose_models.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
import cellpose
2+
from cellpose.models import MODEL_NAMES
3+
4+
for model in MODEL_NAMES:
5+
for model_index in range(4):
6+
model_name = cellpose.models.model_path(model, model_index)

0 commit comments

Comments
 (0)