senoquant 1.0.0b1__py3-none-any.whl

senoquant/tabs/segmentation/stardist_onnx_utils/_stardist/bioimageio_utils.py

@@ -0,0 +1,494 @@
+from pathlib import Path
+try:
+    from importlib.metadata import requires
+except ImportError:
+    from importlib_metadata import requires
+from zipfile import ZipFile
+import numpy as np
+import tempfile
+from packaging.version import Version
+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 # type: ignore
+        import xarray as xr
+        import bioimageio.core # type: ignore
+    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):
+    from ruamel.yaml import YAML
+    from packaging.requirements import Requirement
+    from tensorflow import __version__ as tf_version
+    from . import __version__ as stardist_version
+    # dependencies that start with the name "bioimageio" will be added as conda dependencies
+    reqs_conda = []
+    for req_str in requires("stardist"):
+        req = Requirement(req_str)
+        if (
+            req.marker is not None
+            # only include requirements that are for the "bioimageio" extra
+            and req.marker.evaluate({'extra': 'bioimageio'})
+            and not req.marker.evaluate({'extra': ''})
+            # and package name starts with "bioimageio"
+            and req.name.startswith('bioimageio')
+        ):
+            # https://packaging.pypa.io/en/stable/requirements.html
+            reqs_conda.append(f"{req.name}{req.specifier}")
+    # only stardist and tensorflow as pip dependencies
+    v_tf = Version(tf_version)
+    reqs_pip = (f"stardist>={stardist_version}", f"tensorflow>={v_tf.major}.{v_tf.minor},<{v_tf.major+1}")
+    # conda environment
+    env = dict(
+        name = 'stardist',
+        channels = ['defaults', 'conda-forge'],
+        dependencies = [
+            ('python>=3.7,<3.8' if v_tf.major == 1 else 'python>=3.7'),
+            *reqs_conda,
+            'pip', {'pip': reqs_pip},
+        ],
+    )
+    yaml = YAML(typ='safe')
+    path = outdir / "environment.yaml"
+    with open(path, "w") as f:
+        yaml.dump(env, f)
+    return f"conda:{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, generate_default_deps):
+    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"
+    authors = {
+        'Martin Weigert': dict(name='Martin Weigert', github_user='maweigert'),
+        'Uwe Schmidt': dict(name='Uwe Schmidt', github_user='uschmidt83'),
+    }
+    data = dict(
+        description=package_data["Summary"],
+        authors=list(authors.get(name.strip(),dict(name=name.strip())) for name in package_data["Author"].split(",")),
+        git_repo=package_data["Home-Page"],
+        license=package_data["License"],
+        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/main/images/stardist_logo.jpg"],
+        documentation=_create_stardist_doc(outdir),
+    )
+    if generate_default_deps:  # only if requested, as not required for bioimage.io
+        data['dependencies'] = _create_stardist_dependencies(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=None,
+    mode="tensorflow_saved_model_bundle",
+    min_percentile=1.0,
+    max_percentile=99.8,
+    overwrite_spec_kwargs=None,
+    generate_default_deps=False,
+):
+    """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: None)
+        if None, uses the (folder) name of the model (i.e. `model.name`)
+    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)
+    generate_default_deps: bool
+        not required for bioimage.io, i.e. StarDist models don't need a dependencies field in rdf.yaml (default: False)
+        if True, generate an environment.yaml file recording the python, bioimageio.core, stardist and tensorflow requirements
+        from which a conda environment can be recreated to run this export
+    """
+    _, 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"))
+
+    if name is None:
+        name = model.name
+    name = str(name)
+
+    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, generate_default_deps)
+        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 with name '{name}' 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, uuid
+    from csbdeep.utils import save_json
+    from .models import StarDist2D, StarDist3D
+    *_, bioimageio_core, _ = _import()
+
+    outpath = Path(outpath)
+    not outpath.exists() or _raise(FileExistsError(f"'{outpath}' already exists"))
+
+    with tempfile.TemporaryDirectory() as _tmp_dir:
+        tmp_dir = Path(_tmp_dir)
+        # download the full model content to a temporary folder
+        zip_path = tmp_dir / f"{str(uuid.uuid4())}.zip"
+        bioimageio_core.export_resource_package(source, output_path=zip_path)
+        with ZipFile(zip_path, "r") as zip_ref:
+            zip_ref.extractall(tmp_dir)
+        zip_path.unlink()
+        rdf_path = tmp_dir / "rdf.yaml"
+        biomodel = bioimageio_core.load_resource_description(rdf_path)
+
+        # read the stardist specific content
+        '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']
+
+        # make sure that the keras weights are in the attachments
+        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}'"))
+
+        # save the config and threshold to json, and weights to hdf5 to enable loading as stardist model
+        # copy bioimageio files to separate sub-folder
+        outpath.mkdir(parents=True)
+        save_json(config, str(outpath / 'config.json'))
+        save_json(thresholds, str(outpath / 'thresholds.json'))
+        shutil.copy(str(weights_file), str(outpath / "weights_bioimageio.h5"))
+        shutil.copytree(str(tmp_dir), str(outpath / "bioimageio"))
+
+    model_class = (StarDist2D if config['n_dim'] == 2 else StarDist3D)
+    model = model_class(None, outpath.name, basedir=str(outpath.parent))
+
+    return model

senoquant/tabs/segmentation/stardist_onnx_utils/_stardist/data/__init__.py

@@ -0,0 +1,39 @@
+def abspath(path):
+    import os
+    base_path = os.path.abspath(os.path.dirname(__file__))
+    return os.path.join(base_path, path)
+
+
+def test_image_nuclei_2d(return_mask=False):
+    """ Fluorescence microscopy image and mask from the 2018 kaggle DSB challenge
+
+    Caicedo et al. "Nucleus segmentation across imaging experiments: the 2018 Data Science Bowl." Nature methods 16.12
+    """
+    from tifffile import imread
+    img = imread(abspath("images/img2d.tif"))
+    mask = imread(abspath("images/mask2d.tif"))
+    if return_mask:
+        return img, mask
+    else:
+        return img
+
+
+def test_image_he_2d():
+    """ H&E stained RGB example image from the Cancer Imaging Archive
+    https://www.cancerimagingarchive.net
+    """
+    from imageio import imread
+    img = imread(abspath("images/histo.jpg"))
+    return img
+
+
+def test_image_nuclei_3d(return_mask=False):
+    """ synthetic nuclei 
+    """
+    from tifffile import imread
+    img = imread(abspath("images/img3d.tif"))
+    mask = imread(abspath("images/mask3d.tif"))
+    if return_mask:
+        return img, mask
+    else:
+        return img

senoquant/tabs/segmentation/stardist_onnx_utils/_stardist/geometry/__init__.py

@@ -0,0 +1,10 @@
+from __future__ import absolute_import, print_function
+
+# TODO: rethink naming for 2D/3D functions
+
+from .geom2d import star_dist, relabel_image_stardist, ray_angles, dist_to_coord, polygons_to_label, polygons_to_label_coord
+from .geom3d import star_dist3D, polyhedron_to_label, relabel_image_stardist3D, dist_to_coord3D, export_to_obj_file3D
+
+from .geom2d import _dist_to_coord_old, _polygons_to_label_old
+
+#, dist_to_volume, dist_to_centroid