forked from stardist/stardist
-
Notifications
You must be signed in to change notification settings - Fork 0
/
bioimageio_utils.py
428 lines (367 loc) · 17 KB
/
bioimageio_utils.py
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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
from pathlib import Path
from pkg_resources import get_distribution
from zipfile import ZipFile
import numpy as np
import tempfile
from csbdeep.utils import axes_check_and_normalize, normalize, _raise
DEEPIMAGEJ_MACRO = \
"""
//*******************************************************************
// Date: July-2021
// Credits: StarDist, DeepImageJ
// URL:
// https://github.com/stardist/stardist
// https://deepimagej.github.io/deepimagej
// This macro was adapted from
// https://github.com/deepimagej/imagej-macros/blob/648caa867f6ccb459649d4d3799efa1e2e0c5204/StarDist2D_Post-processing.ijm
// Please cite the respective contributions when using this code.
//*******************************************************************
// Macro to run StarDist postprocessing on 2D images.
// StarDist and deepImageJ plugins need to be installed.
// The macro assumes that the image to process is a stack in which
// the first channel corresponds to the object probability map
// and the remaining channels are the radial distances from each
// pixel to the object boundary.
//*******************************************************************
// Get the name of the image to call it
getDimensions(width, height, channels, slices, frames);
name=getTitle();
probThresh={probThresh};
nmsThresh={nmsThresh};
// Isolate the detection probability scores
run("Make Substack...", "channels=1");
rename("scores");
// Isolate the oriented distances
run("Fire");
selectWindow(name);
run("Delete Slice", "delete=channel");
selectWindow(name);
run("Properties...", "channels=" + maxOf(channels, slices) - 1 + " slices=1 frames=1 pixel_width=1.0000 pixel_height=1.0000 voxel_depth=1.0000");
rename("distances");
run("royal");
// Run StarDist plugin
run("Command From Macro", "command=[de.csbdresden.stardist.StarDist2DNMS], args=['prob':'scores', 'dist':'distances', 'probThresh':'" + probThresh + "', 'nmsThresh':'" + nmsThresh + "', 'outputType':'Both', 'excludeBoundary':'2', 'roiPosition':'Stack', 'verbose':'false'], process=[false]");
"""
def _import(error=True):
try:
from importlib_metadata import metadata
from bioimageio.core.build_spec import build_model
import xarray as xr
import bioimageio.core
except ImportError:
if error:
raise RuntimeError(
"Required libraries are missing for bioimage.io model export.\n"
"Please install StarDist as follows: pip install 'stardist[bioimageio]'\n"
"(You do not need to uninstall StarDist first.)"
)
else:
return None
return metadata, build_model, bioimageio.core, xr
def _create_stardist_dependencies(outdir):
pkg_info = get_distribution("stardist")
reqs = ("tensorflow",) + tuple(map(str, pkg_info.requires()))
path = outdir / "requirements.txt"
with open(path, "w") as f:
f.write("\n".join(reqs))
return f"pip:{path}"
def _create_stardist_doc(outdir):
doc_path = outdir / "README.md"
text = (
"# StarDist Model\n"
"This is a model for object detection with star-convex shapes.\n"
"Please see the [StarDist repository](https://github.com/stardist/stardist) for details."
)
with open(doc_path, "w") as f:
f.write(text)
return doc_path
def _get_stardist_metadata(outdir, model):
metadata, *_ = _import()
package_data = metadata("stardist")
doi_2d = "https://doi.org/10.1007/978-3-030-00934-2_30"
doi_3d = "https://doi.org/10.1109/WACV45572.2020.9093435"
data = dict(
description=package_data["Summary"],
authors=list(dict(name=name.strip()) for name in package_data["Author"].split(",")),
git_repo=package_data["Home-Page"],
license=package_data["License"],
dependencies=_create_stardist_dependencies(outdir),
cite=[{"text": "Cell Detection with Star-Convex Polygons", "doi": doi_2d},
{"text": "Star-convex Polyhedra for 3D Object Detection and Segmentation in Microscopy", "doi": doi_3d}],
tags=[
'fluorescence-light-microscopy', 'whole-slide-imaging', 'other', # modality
f'{model.config.n_dim}d', # dims
'cells', 'nuclei', # content
'tensorflow', # framework
'fiji', # software
'unet', # network
'instance-segmentation', 'object-detection', # task
'stardist',
],
covers=["https://raw.githubusercontent.com/stardist/stardist/master/images/stardist_logo.jpg"],
documentation=_create_stardist_doc(outdir),
)
return data
def _predict_tf(model_path, test_input):
import tensorflow as tf
from csbdeep.utils.tf import IS_TF_1
# need to unzip the model assets
model_assets = model_path.parent / "tf_model"
with ZipFile(model_path, "r") as f:
f.extractall(model_assets)
if IS_TF_1:
# make a new graph, i.e. don't use the global default graph
with tf.Graph().as_default():
with tf.Session() as sess:
tf_model = tf.saved_model.load_v2(str(model_assets))
x = tf.convert_to_tensor(test_input, dtype=tf.float32)
model = tf_model.signatures["serving_default"]
y = model(x)
sess.run(tf.global_variables_initializer())
output = sess.run(y["output"])
else:
tf_model = tf.saved_model.load(str(model_assets))
x = tf.convert_to_tensor(test_input, dtype=tf.float32)
model = tf_model.signatures["serving_default"]
y = model(x)
output = y["output"].numpy()
return output
def _get_weights_and_model_metadata(outdir, model, test_input, test_input_axes, test_input_norm_axes, mode, min_percentile, max_percentile):
# get the path to the exported model assets (saved in outdir)
if mode == "keras_hdf5":
raise NotImplementedError("Export to keras format is not supported yet")
elif mode == "tensorflow_saved_model_bundle":
assets_uri = outdir / "TF_SavedModel.zip"
model_csbdeep = model.export_TF(assets_uri, single_output=True, upsample_grid=True)
else:
raise ValueError(f"Unsupported mode: {mode}")
# to force "inputs.data_type: float32" in the spec (bonus: disables normalization warning in model._predict_setup)
test_input = test_input.astype(np.float32)
# convert test_input to axes_net semantics and shape, also resize if necessary (to adhere to axes_net_div_by)
test_input, axes_img, axes_net, axes_net_div_by, *_ = model._predict_setup(
img=test_input,
axes=test_input_axes,
normalizer=None,
n_tiles=None,
show_tile_progress=False,
predict_kwargs={},
)
# normalization axes string and numeric indices
axes_norm = set(axes_net).intersection(set(axes_check_and_normalize(test_input_norm_axes, disallowed='S')))
axes_norm = "".join(a for a in axes_net if a in axes_norm) # preserve order of axes_net
axes_norm_num = tuple(axes_net.index(a) for a in axes_norm)
# normalize input image
test_input_norm = normalize(test_input, pmin=min_percentile, pmax=max_percentile, axis=axes_norm_num)
net_axes_in = axes_net.lower()
net_axes_out = axes_check_and_normalize(model._axes_out).lower()
ndim_tensor = len(net_axes_out) + 1
input_min_shape = list(axes_net_div_by)
input_min_shape[axes_net.index('C')] = model.config.n_channel_in
input_step = list(axes_net_div_by)
input_step[axes_net.index('C')] = 0
# add the batch axis to shape and step
input_min_shape = [1] + input_min_shape
input_step = [0] + input_step
# the axes strings in bioimageio convention
input_axes = "b" + net_axes_in.lower()
output_axes = "b" + net_axes_out.lower()
if mode == "keras_hdf5":
output_names = ("prob", "dist") + (("class_prob",) if model._is_multiclass() else ())
output_n_channels = (1, model.config.n_rays,) + ((1,) if model._is_multiclass() else ())
# the output shape is computed from the input shape using
# output_shape[i] = output_scale[i] * input_shape[i] + 2 * output_offset[i]
output_scale = [1]+list(1/g for g in model.config.grid) + [0]
output_offset = [0]*(ndim_tensor)
elif mode == "tensorflow_saved_model_bundle":
if model._is_multiclass():
raise NotImplementedError("Tensorflow SavedModel not supported for multiclass models yet")
# regarding input/output names: https://github.com/CSBDeep/CSBDeep/blob/b0d2f5f344ebe65a9b4c3007f4567fe74268c813/csbdeep/utils/tf.py#L193-L194
input_names = ["input"]
output_names = ["output"]
output_n_channels = (1 + model.config.n_rays,)
# the output shape is computed from the input shape using
# output_shape[i] = output_scale[i] * input_shape[i] + 2 * output_offset[i]
# same shape as input except for the channel dimension
output_scale = [1]*(ndim_tensor)
output_scale[output_axes.index("c")] = 0
# no offset, except for the input axes, where it is output channel / 2
output_offset = [0.0]*(ndim_tensor)
output_offset[output_axes.index("c")] = output_n_channels[0] / 2.0
assert all(s in (0, 1) for s in output_scale), "halo computation assumption violated"
halo = model._axes_tile_overlap(output_axes.replace('b', 's'))
halo = [int(np.ceil(v/8)*8) for v in halo] # optional: round up to be divisible by 8
# the output shape needs to be valid after cropping the halo, so we add the halo to the input min shape
input_min_shape = [ms + 2 * ha for ms, ha in zip(input_min_shape, halo)]
# make sure the input min shape is still divisible by the min axis divisor
input_min_shape = input_min_shape[:1] + [ms + (-ms % div_by) for ms, div_by in zip(input_min_shape[1:], axes_net_div_by)]
assert all(ms % div_by == 0 for ms, div_by in zip(input_min_shape[1:], axes_net_div_by))
metadata, *_ = _import()
package_data = metadata("stardist")
is_2D = model.config.n_dim == 2
weights_file = outdir / "stardist_weights.h5"
model.keras_model.save_weights(str(weights_file))
config = dict(
stardist=dict(
python_version=package_data["Version"],
thresholds=dict(model.thresholds._asdict()),
weights=weights_file.name,
config=vars(model.config),
)
)
if is_2D:
macro_file = outdir / "stardist_postprocessing.ijm"
with open(str(macro_file), 'w', encoding='utf-8') as f:
f.write(DEEPIMAGEJ_MACRO.format(probThresh=model.thresholds.prob, nmsThresh=model.thresholds.nms))
config['stardist'].update(postprocessing_macro=macro_file.name)
n_inputs = len(input_names)
assert n_inputs == 1
input_config = dict(
input_names=input_names,
input_min_shape=[input_min_shape],
input_step=[input_step],
input_axes=[input_axes],
input_data_range=[["-inf", "inf"]],
preprocessing=[[dict(
name="scale_range",
kwargs=dict(
mode="per_sample",
axes=axes_norm.lower(),
min_percentile=min_percentile,
max_percentile=max_percentile,
))]]
)
n_outputs = len(output_names)
output_config = dict(
output_names=output_names,
output_data_range=[["-inf", "inf"]] * n_outputs,
output_axes=[output_axes] * n_outputs,
output_reference=[input_names[0]] * n_outputs,
output_scale=[output_scale] * n_outputs,
output_offset=[output_offset] * n_outputs,
halo=[halo] * n_outputs
)
in_path = outdir / "test_input.npy"
np.save(in_path, test_input[np.newaxis])
if mode == "tensorflow_saved_model_bundle":
test_outputs = _predict_tf(assets_uri, test_input_norm[np.newaxis])
else:
test_outputs = model.predict(test_input_norm)
# out_paths = []
# for i, out in enumerate(test_outputs):
# p = outdir / f"test_output{i}.npy"
# np.save(p, out)
# out_paths.append(p)
assert n_outputs == 1
out_paths = [outdir / "test_output.npy"]
np.save(out_paths[0], test_outputs)
from tensorflow import __version__ as tf_version
data = dict(weight_uri=assets_uri, test_inputs=[in_path], test_outputs=out_paths,
config=config, tensorflow_version=tf_version)
data.update(input_config)
data.update(output_config)
_files = [str(weights_file)]
if is_2D:
_files.append(str(macro_file))
data.update(attachments=dict(files=_files))
return data
def export_bioimageio(
model,
outpath,
test_input,
test_input_axes=None,
test_input_norm_axes='ZYX',
name="bioimageio_model",
mode="tensorflow_saved_model_bundle",
min_percentile=1.0,
max_percentile=99.8,
overwrite_spec_kwargs=None,
):
"""Export stardist model into bioimage.io format, https://github.com/bioimage-io/spec-bioimage-io.
Parameters
----------
model: StarDist2D, StarDist3D
the model to convert
outpath: str, Path
where to save the model
test_input: np.ndarray
input image for generating test data
test_input_axes: str or None
the axes of the test input, for example 'YX' for a 2d image or 'ZYX' for a 3d volume
using None assumes that axes of test_input are the same as those of model
test_input_norm_axes: str
the axes of the test input which will be jointly normalized, for example 'ZYX' for all spatial dimensions ('Z' ignored for 2D input)
use 'ZYXC' to also jointly normalize channels (e.g. for RGB input images)
name: str
the name of this model (default: "bioimageio_model")
mode: str
the export type for this model (default: "tensorflow_saved_model_bundle")
min_percentile: float
min percentile to be used for image normalization (default: 1.0)
max_percentile: float
max percentile to be used for image normalization (default: 99.8)
overwrite_spec_kwargs: dict or None
spec keywords that should be overloaded (default: None)
"""
_, build_model, *_ = _import()
from .models import StarDist2D, StarDist3D
isinstance(model, (StarDist2D, StarDist3D)) or _raise(ValueError("not a valid model"))
0 <= min_percentile < max_percentile <= 100 or _raise(ValueError("invalid percentile values"))
outpath = Path(outpath)
if outpath.suffix == "":
outdir = outpath
zip_path = outdir / f"{name}.zip"
elif outpath.suffix == ".zip":
outdir = outpath.parent
zip_path = outpath
else:
raise ValueError(f"outpath has to be a folder or zip file, got {outpath}")
outdir.mkdir(exist_ok=True, parents=True)
with tempfile.TemporaryDirectory() as _tmp_dir:
tmp_dir = Path(_tmp_dir)
kwargs = _get_stardist_metadata(tmp_dir, model)
model_kwargs = _get_weights_and_model_metadata(tmp_dir, model, test_input, test_input_axes, test_input_norm_axes, mode,
min_percentile=min_percentile, max_percentile=max_percentile)
kwargs.update(model_kwargs)
if overwrite_spec_kwargs is not None:
kwargs.update(overwrite_spec_kwargs)
build_model(name=name, output_path=zip_path, add_deepimagej_config=(model.config.n_dim==2), root=tmp_dir, **kwargs)
print(f"\nbioimage.io model exported to '{zip_path}'")
def import_bioimageio(source, outpath):
"""Import stardist model from bioimage.io format, https://github.com/bioimage-io/spec-bioimage-io.
Load a model in bioimage.io format from the given `source` (e.g. path to zip file, URL)
and convert it to a regular stardist model, which will be saved in the folder `outpath`.
Parameters
----------
source: str, Path
bioimage.io resource (e.g. path, URL)
outpath: str, Path
folder to save the stardist model (must not exist previously)
Returns
-------
StarDist2D or StarDist3D
stardist model loaded from `outpath`
"""
import shutil
from csbdeep.utils import save_json
from .models import StarDist2D, StarDist3D
*_, bioimageio_core, _ = _import()
biomodel = bioimageio_core.load_resource_description(source)
outpath = Path(outpath)
not outpath.exists() or _raise(FileExistsError(f"'{outpath}' already exists"))
'stardist' in biomodel.config or _raise(RuntimeError("bioimage.io model not compatible"))
config = biomodel.config['stardist']['config']
thresholds = biomodel.config['stardist']['thresholds']
weights = biomodel.config['stardist']['weights']
weights_file = None
for f in biomodel.attachments.files:
if f.name == weights and f.exists():
weights_file = f
break
weights_file is not None or _raise(FileNotFoundError(f"couldn't find weights file '{weights}'"))
outpath.mkdir(parents=True, exist_ok=True)
save_json(config, str(outpath / 'config.json'))
save_json(thresholds, str(outpath / 'thresholds.json'))
shutil.copy(str(weights_file), str(outpath / weights))
model_class = (StarDist2D if config['n_dim'] == 2 else StarDist3D)
model = model_class(None, outpath.name, basedir=str(outpath.parent))
return model