Upload 40 files

README.md

@@ -1,13 +1,109 @@
 ---
-title: Lewislou Cell Seg Sribd
-emoji: ⚡
-colorFrom: blue
-colorTo: blue
-sdk: gradio
-sdk_version: 3.38.0
-app_file: app.py
-pinned: false
 license: apache-2.0
+language:
+- en
+metrics:
+- f1
+tags:
+- cell segmentation
+- stardist
+- hover-net
+library_name: transformers
+pipeline_tag: image-segmentation
+datasets:
+- Lewislou/cell_samples
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# Model Card for cell-seg-sribd
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+This repository provides the solution of team Sribd-med for NeurIPS-CellSeg Challenge. The details of our method are described in our paper [Multi-stream Cell Segmentation with Low-level Cues for Multi-modality Images]. Some parts of the codes are from the baseline codes of the NeurIPS-CellSeg-Baseline repository,
+
+You can reproduce our method as follows step by step:
+
+
+### How to Get Started with the Model
+
+Install requirements by python -m pip install -r requirements.txt
+
+## Training Details
+
+### Training Data
+
+The competition training and tuning data can be downloaded from https://neurips22-cellseg.grand-challenge.org/dataset/ Besides, you can download three publiced data from the following link: Cellpose: https://www.cellpose.org/dataset  Omnipose: http://www.cellpose.org/dataset_omnipose Sartorius: https://www.kaggle.com/competitions/sartorius-cell-instance-segmentation/overview 
+
+## Environments and Requirements:
+Install requirements by
+
+```shell
+python -m pip install -r requirements.txt
+```
+
+### How to use
+
+Here is how to use this model:
+
+```python
+
+
+from skimage import io, segmentation, morphology, measure, exposure
+from sribd_cellseg_models import MultiStreamCellSegModel,ModelConfig
+import numpy as np
+import tifffile as tif
+import requests
+import torch
+from PIL import Image
+from overlay import visualize_instances_map
+import cv2
+img_name = 'test_images/cell_00551.tiff'
+
+def normalize_channel(img, lower=1, upper=99):
+    non_zero_vals = img[np.nonzero(img)]
+    percentiles = np.percentile(non_zero_vals, [lower, upper])
+    if percentiles[1] - percentiles[0] > 0.001:
+        img_norm = exposure.rescale_intensity(img, in_range=(percentiles[0], percentiles[1]), out_range='uint8')
+    else:
+        img_norm = img
+    return img_norm.astype(np.uint8)
+if img_name.endswith('.tif') or img_name.endswith('.tiff'):
+    img_data = tif.imread(img_name)
+else:
+    img_data = io.imread(img_name)
+        # normalize image data
+if len(img_data.shape) == 2:
+    img_data = np.repeat(np.expand_dims(img_data, axis=-1), 3, axis=-1)
+elif len(img_data.shape) == 3 and img_data.shape[-1] > 3:
+    img_data = img_data[:,:, :3]
+else:
+    pass
+pre_img_data = np.zeros(img_data.shape, dtype=np.uint8)
+for i in range(3):
+    img_channel_i = img_data[:,:,i]
+    if len(img_channel_i[np.nonzero(img_channel_i)])>0:
+        pre_img_data[:,:,i] = normalize_channel(img_channel_i, lower=1, upper=99)
+#dummy_input = np.zeros((512,512,3)).astype(np.uint8)
+my_model = MultiStreamCellSegModel.from_pretrained("Lewislou/cellseg_sribd")
+checkpoints = torch.load('model.pt')
+my_model.__init__(ModelConfig())
+my_model.load_checkpoints(checkpoints)
+with torch.no_grad():
+    output = my_model(pre_img_data)
+overlay = visualize_instances_map(pre_img_data,star_label)
+cv2.imwrite('prediction.png', cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
+
+
+```
+
+## Citation
+If any part of this code is used, please acknowledge it appropriately and cite the paper:
+```bibtex
+@misc{
+lou2022multistream,
+title={Multi-stream Cell Segmentation with Low-level Cues for Multi-modality Images},
+author={WEI LOU and Xinyi Yu and Chenyu Liu and Xiang Wan and Guanbin Li and Siqi Liu and Haofeng Li},
+year={2022},
+url={https://openreview.net/forum?id=G24BybwKe9}
+}
+```

classifiers.py

@@ -0,0 +1,261 @@
+from functools import partial
+from typing import Any, Callable, List, Optional, Type, Union
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
+    """3x3 convolution with padding"""
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=dilation,
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
+
+
+def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion: int = 1
+
+    def __init__(
+        self,
+        inplanes: int,
+        planes: int,
+        stride: int = 1,
+        downsample: Optional[nn.Module] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError("BasicBlock only supports groups=1 and base_width=64")
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+    expansion: int = 4
+
+    def __init__(
+        self,
+        inplanes: int,
+        planes: int,
+        stride: int = 1,
+        downsample: Optional[nn.Module] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.0)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+class ResNet(nn.Module):
+    def __init__(
+        self,
+        block: Type[Union[BasicBlock, Bottleneck]],
+        layers: List[int],
+        num_classes: int = 1000,
+        zero_init_residual: bool = False,
+        groups: int = 1,
+        width_per_group: int = 64,
+        replace_stride_with_dilation: Optional[List[bool]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super().__init__()
+        # _log_api_usage_once(self)
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError(
+                "replace_stride_with_dilation should be None "
+                f"or a 3-element tuple, got {replace_stride_with_dilation}"
+            )
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2])
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck) and m.bn3.weight is not None:
+                    nn.init.constant_(m.bn3.weight, 0)  # type: ignore[arg-type]
+                elif isinstance(m, BasicBlock) and m.bn2.weight is not None:
+                    nn.init.constant_(m.bn2.weight, 0)  # type: ignore[arg-type]
+
+    def _make_layer(
+        self,
+        block: Type[Union[BasicBlock, Bottleneck]],
+        planes: int,
+        blocks: int,
+        stride: int = 1,
+        dilate: bool = False,
+    ) -> nn.Sequential:
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer
+            )
+        )
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(
+                    self.inplanes,
+                    planes,
+                    groups=self.groups,
+                    base_width=self.base_width,
+                    dilation=self.dilation,
+                    norm_layer=norm_layer,
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        # See note [TorchScript super()]
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+def resnet18(weights=None):
+    # weights: path 
+    model = ResNet(BasicBlock, [2, 2, 2, 2], num_classes=4)
+    if weights is not None:
+        model.load_state_dict(torch.load(weights))
+    return model
+
+def resnet10():
+    return ResNet(BasicBlock, [1, 1, 1, 1], num_classes=4)

config.json

@@ -0,0 +1,25 @@
+{
+  "architectures": [
+    "MultiStreamCellSegModel"
+  ],
+  "block_size": 2048,
+  "context": 128,
+  "device": "cpu",
+  "input_channels": 3,
+  "ksize": 15,
+  "min_overlap": 128,
+  "model_type": "cell_sribd",
+  "n_rays": 32,
+  "np_thres": 0.6,
+  "num_classes": 4,
+  "obj_size_thres": 100,
+  "overall_thres": 0.4,
+  "overlap": 0.5,
+  "roi_size": [
+    512,
+    512
+  ],
+  "sw_batch_size": 4,
+  "torch_dtype": "float32",
+  "transformers_version": "4.27.1"
+}

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:460f2c3a9168220ef03404df983b923c7f59d6db873536cd44d9e4b7e4354f6c
+size 135

models/__init__.py

@@ -0,0 +1,10 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Sun Mar 20 14:23:55 2022
+
+@author: jma
+"""
+
+#from .unetr2d import UNETR2D
+#from .swin_unetr import SwinUNETR

models/convnext.py

@@ -0,0 +1,220 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from functools import partial
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import trunc_normal_, DropPath
+from timm.models.registry import register_model
+from monai.networks.layers.factories import Act, Conv, Pad, Pool
+from monai.networks.layers.utils import get_norm_layer
+from monai.utils.module import look_up_option
+from typing import List, NamedTuple, Optional, Tuple, Type, Union
+class Block(nn.Module):
+    r""" ConvNeXt Block. There are two equivalent implementations:
+    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
+    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
+    We use (2) as we find it slightly faster in PyTorch
+    
+    Args:
+        dim (int): Number of input channels.
+        drop_path (float): Stochastic depth rate. Default: 0.0
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+    """
+    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
+        self.norm = LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), 
+                                    requires_grad=True) if layer_scale_init_value > 0 else None
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + self.drop_path(x)
+        return x
+
+class ConvNeXt(nn.Module):
+    r""" ConvNeXt
+        A PyTorch impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+    def __init__(self, in_chans=3, num_classes=21841, 
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], drop_path_rate=0., 
+                 layer_scale_init_value=1e-6, head_init_scale=1., out_indices=[0, 1, 2, 3],
+                 ):
+        super().__init__()
+        # conv_type: Type[Union[nn.Conv1d, nn.Conv2d, nn.Conv3d]] = Conv["conv", 2]
+        # self._conv_stem = conv_type(self.in_channels, self.in_channels, kernel_size=3, stride=stride, bias=False)
+        # self._conv_stem_padding = _make_same_padder(self._conv_stem, current_image_size)
+       
+        self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers
+        stem = nn.Sequential(
+            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+            LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = nn.Sequential(
+                    LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
+                    nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates=[x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] 
+        cur = 0
+        for i in range(4):
+            stage = nn.Sequential(
+                *[Block(dim=dims[i], drop_path=dp_rates[cur + j], 
+                layer_scale_init_value=layer_scale_init_value) for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+
+        self.out_indices = out_indices
+
+        norm_layer = partial(LayerNorm, eps=1e-6, data_format="channels_first")
+        for i_layer in range(4):
+            layer = norm_layer(dims[i_layer])
+            layer_name = f'norm{i_layer}'
+            self.add_module(layer_name, layer)
+        self.apply(self._init_weights)
+    
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward_features(self, x):
+        outs = []
+
+        for i in range(4):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+            if i in self.out_indices:
+                norm_layer = getattr(self, f'norm{i}')
+                x_out = norm_layer(x)
+         
+                outs.append(x_out)
+
+        return tuple(outs)
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        
+        return x
+
+class LayerNorm(nn.Module):
+    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
+    with shape (batch_size, channels, height, width).
+    """
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError 
+        self.normalized_shape = (normalized_shape, )
+    
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x
+
+
+model_urls = {
+    "convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
+    "convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
+    "convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
+    "convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
+    "convnext_tiny_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth",
+    "convnext_small_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth",
+    "convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
+    "convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
+    "convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
+}
+
+@register_model
+def convnext_tiny(pretrained=False,in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_tiny_22k'] if in_22k else model_urls['convnext_tiny_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
+        model.load_state_dict(checkpoint["model"])
+    return model
+
+@register_model
+def convnext_small(pretrained=False,in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_small_22k'] if in_22k else model_urls['convnext_small_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
+        model.load_state_dict(checkpoint["model"], strict=False)
+    return model
+
+@register_model
+def convnext_base(pretrained=False, in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_base_22k'] if in_22k else model_urls['convnext_base_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
+        model.load_state_dict(checkpoint["model"], strict=False)
+    return model
+
+@register_model
+def convnext_large(pretrained=False, in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_large_22k'] if in_22k else model_urls['convnext_large_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
+        model.load_state_dict(checkpoint["model"])
+    return model
+
+@register_model
+def convnext_xlarge(pretrained=False, in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
+    if pretrained:
+        assert in_22k, "only ImageNet-22K pre-trained ConvNeXt-XL is available; please set in_22k=True"
+        url = model_urls['convnext_xlarge_22k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
+        model.load_state_dict(checkpoint["model"])
+    return model

models/flexible_unet.py

@@ -0,0 +1,312 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Optional, Sequence, Tuple, Union
+
+import torch
+from torch import nn
+
+from monai.networks.blocks import UpSample
+from monai.networks.layers.factories import Conv
+from monai.networks.layers.utils import get_act_layer
+from monai.networks.nets import EfficientNetBNFeatures
+from monai.networks.nets.basic_unet import UpCat
+from monai.utils import InterpolateMode
+
+__all__ = ["FlexibleUNet"]
+
+encoder_feature_channel = {
+    "efficientnet-b0": (16, 24, 40, 112, 320),
+    "efficientnet-b1": (16, 24, 40, 112, 320),
+    "efficientnet-b2": (16, 24, 48, 120, 352),
+    "efficientnet-b3": (24, 32, 48, 136, 384),
+    "efficientnet-b4": (24, 32, 56, 160, 448),
+    "efficientnet-b5": (24, 40, 64, 176, 512),
+    "efficientnet-b6": (32, 40, 72, 200, 576),
+    "efficientnet-b7": (32, 48, 80, 224, 640),
+    "efficientnet-b8": (32, 56, 88, 248, 704),
+    "efficientnet-l2": (72, 104, 176, 480, 1376),
+}
+
+
+def _get_encoder_channels_by_backbone(backbone: str, in_channels: int = 3) -> tuple:
+    """
+    Get the encoder output channels by given backbone name.
+
+    Args:
+        backbone: name of backbone to generate features, can be from [efficientnet-b0, ..., efficientnet-b7].
+        in_channels: channel of input tensor, default to 3.
+
+    Returns:
+        A tuple of output feature map channels' length .
+    """
+    encoder_channel_tuple = encoder_feature_channel[backbone]
+    encoder_channel_list = [in_channels] + list(encoder_channel_tuple)
+    encoder_channel = tuple(encoder_channel_list)
+    return encoder_channel
+
+
+class UNetDecoder(nn.Module):
+    """
+    UNet Decoder.
+    This class refers to `segmentation_models.pytorch
+    <https://github.com/qubvel/segmentation_models.pytorch>`_.
+
+    Args:
+        spatial_dims: number of spatial dimensions.
+        encoder_channels: number of output channels for all feature maps in encoder.
+            `len(encoder_channels)` should be no less than 2.
+        decoder_channels: number of output channels for all feature maps in decoder.
+            `len(decoder_channels)` should equal to `len(encoder_channels) - 1`.
+        act: activation type and arguments.
+        norm: feature normalization type and arguments.
+        dropout: dropout ratio.
+        bias: whether to have a bias term in convolution blocks in this decoder.
+        upsample: upsampling mode, available options are
+            ``"deconv"``, ``"pixelshuffle"``, ``"nontrainable"``.
+        pre_conv: a conv block applied before upsampling.
+            Only used in the "nontrainable" or "pixelshuffle" mode.
+        interp_mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``}
+            Only used in the "nontrainable" mode.
+        align_corners: set the align_corners parameter for upsample. Defaults to True.
+            Only used in the "nontrainable" mode.
+        is_pad: whether to pad upsampling features to fit the encoder spatial dims.
+
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int,
+        encoder_channels: Sequence[int],
+        decoder_channels: Sequence[int],
+        act: Union[str, tuple],
+        norm: Union[str, tuple],
+        dropout: Union[float, tuple],
+        bias: bool,
+        upsample: str,
+        pre_conv: Optional[str],
+        interp_mode: str,
+        align_corners: Optional[bool],
+        is_pad: bool,
+    ):
+
+        super().__init__()
+        if len(encoder_channels) < 2:
+            raise ValueError("the length of `encoder_channels` should be no less than 2.")
+        if len(decoder_channels) != len(encoder_channels) - 1:
+            raise ValueError("`len(decoder_channels)` should equal to `len(encoder_channels) - 1`.")
+
+        in_channels = [encoder_channels[-1]] + list(decoder_channels[:-1])
+        skip_channels = list(encoder_channels[1:-1][::-1]) + [0]
+        halves = [True] * (len(skip_channels) - 1)
+        halves.append(False)
+        blocks = []
+        for in_chn, skip_chn, out_chn, halve in zip(in_channels, skip_channels, decoder_channels, halves):
+            blocks.append(
+                UpCat(
+                    spatial_dims=spatial_dims,
+                    in_chns=in_chn,
+                    cat_chns=skip_chn,
+                    out_chns=out_chn,
+                    act=act,
+                    norm=norm,
+                    dropout=dropout,
+                    bias=bias,
+                    upsample=upsample,
+                    pre_conv=pre_conv,
+                    interp_mode=interp_mode,
+                    align_corners=align_corners,
+                    halves=halve,
+                    is_pad=is_pad,
+                )
+            )
+        self.blocks = nn.ModuleList(blocks)
+
+    def forward(self, features: List[torch.Tensor], skip_connect: int = 4):
+        skips = features[:-1][::-1]
+        features = features[1:][::-1]
+
+        x = features[0]
+        for i, block in enumerate(self.blocks):
+            if i < skip_connect:
+                skip = skips[i]
+            else:
+                skip = None
+            x = block(x, skip)
+
+        return x
+
+
+class SegmentationHead(nn.Sequential):
+    """
+    Segmentation head.
+    This class refers to `segmentation_models.pytorch
+    <https://github.com/qubvel/segmentation_models.pytorch>`_.
+
+    Args:
+        spatial_dims: number of spatial dimensions.
+        in_channels: number of input channels for the block.
+        out_channels: number of output channels for the block.
+        kernel_size: kernel size for the conv layer.
+        act: activation type and arguments.
+        scale_factor: multiplier for spatial size. Has to match input size if it is a tuple.
+
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        act: Optional[Union[Tuple, str]] = None,
+        scale_factor: float = 1.0,
+    ):
+
+        conv_layer = Conv[Conv.CONV, spatial_dims](
+            in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=kernel_size // 2
+        )
+        up_layer: nn.Module = nn.Identity()
+        if scale_factor > 1.0:
+            up_layer = UpSample(
+                spatial_dims=spatial_dims,
+                scale_factor=scale_factor,
+                mode="nontrainable",
+                pre_conv=None,
+                interp_mode=InterpolateMode.LINEAR,
+            )
+        if act is not None:
+            act_layer = get_act_layer(act)
+        else:
+            act_layer = nn.Identity()
+        super().__init__(conv_layer, up_layer, act_layer)
+
+
+class FlexibleUNet(nn.Module):
+    """
+    A flexible implementation of UNet-like encoder-decoder architecture.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        backbone: str,
+        pretrained: bool = False,
+        decoder_channels: Tuple = (256, 128, 64, 32, 16),
+        spatial_dims: int = 2,
+        norm: Union[str, tuple] = ("batch", {"eps": 1e-3, "momentum": 0.1}),
+        act: Union[str, tuple] = ("relu", {"inplace": True}),
+        dropout: Union[float, tuple] = 0.0,
+        decoder_bias: bool = False,
+        upsample: str = "nontrainable",
+        interp_mode: str = "nearest",
+        is_pad: bool = True,
+    ) -> None:
+        """
+        A flexible implement of UNet, in which the backbone/encoder can be replaced with
+        any efficient network. Currently the input must have a 2 or 3 spatial dimension
+        and the spatial size of each dimension must be a multiple of 32 if is pad parameter
+        is False
+
+        Args:
+            in_channels: number of input channels.
+            out_channels: number of output channels.
+            backbone: name of backbones to initialize, only support efficientnet right now,
+                can be from [efficientnet-b0,..., efficientnet-b8, efficientnet-l2].
+            pretrained: whether to initialize pretrained ImageNet weights, only available
+                for spatial_dims=2 and batch norm is used, default to False.
+            decoder_channels: number of output channels for all feature maps in decoder.
+                `len(decoder_channels)` should equal to `len(encoder_channels) - 1`,default
+                to (256, 128, 64, 32, 16).
+            spatial_dims: number of spatial dimensions, default to 2.
+            norm: normalization type and arguments, default to ("batch", {"eps": 1e-3,
+                "momentum": 0.1}).
+            act: activation type and arguments, default to ("relu", {"inplace": True}).
+            dropout: dropout ratio, default to 0.0.
+            decoder_bias: whether to have a bias term in decoder's convolution blocks.
+            upsample: upsampling mode, available options are``"deconv"``, ``"pixelshuffle"``,
+                ``"nontrainable"``.
+            interp_mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``}
+                Only used in the "nontrainable" mode.
+            is_pad: whether to pad upsampling features to fit features from encoder. Default to True.
+                If this parameter is set to "True", the spatial dim of network input can be arbitary
+                size, which is not supported by TensorRT. Otherwise, it must be a multiple of 32.
+        """
+        super().__init__()
+
+        if backbone not in encoder_feature_channel:
+            raise ValueError(f"invalid model_name {backbone} found, must be one of {encoder_feature_channel.keys()}.")
+
+        if spatial_dims not in (2, 3):
+            raise ValueError("spatial_dims can only be 2 or 3.")
+
+        adv_prop = "ap" in backbone
+
+        self.backbone = backbone
+        self.spatial_dims = spatial_dims
+        model_name = backbone
+        encoder_channels = _get_encoder_channels_by_backbone(backbone, in_channels)
+        self.encoder = EfficientNetBNFeatures(
+            model_name=model_name,
+            pretrained=pretrained,
+            in_channels=in_channels,
+            spatial_dims=spatial_dims,
+            norm=norm,
+            adv_prop=adv_prop,
+        )
+        self.decoder = UNetDecoder(
+            spatial_dims=spatial_dims,
+            encoder_channels=encoder_channels,
+            decoder_channels=decoder_channels,
+            act=act,
+            norm=norm,
+            dropout=dropout,
+            bias=decoder_bias,
+            upsample=upsample,
+            interp_mode=interp_mode,
+            pre_conv=None,
+            align_corners=None,
+            is_pad=is_pad,
+        )
+        self.dist_head = SegmentationHead(
+            spatial_dims=spatial_dims,
+            in_channels=decoder_channels[-1],
+            out_channels=32,
+            kernel_size=1,
+            act='relu',
+        )
+        self.prob_head = SegmentationHead(
+            spatial_dims=spatial_dims,
+            in_channels=decoder_channels[-1],
+            out_channels=1,
+            kernel_size=1,
+            act='sigmoid',
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        """
+        Do a typical encoder-decoder-header inference.
+
+        Args:
+            inputs: input should have spatially N dimensions ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N])``,
+                N is defined by `dimensions`.
+
+        Returns:
+            A torch Tensor of "raw" predictions in shape ``(Batch, out_channels, dim_0[, dim_1, ..., dim_N])``.
+
+        """
+        x = inputs
+        enc_out = self.encoder(x)
+        decoder_out = self.decoder(enc_out)
+        dist = self.dist_head(decoder_out)
+        prob = self.prob_head(decoder_out)
+        return dist,prob

models/flexible_unet_convnext.py

@@ -0,0 +1,447 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Optional, Sequence, Tuple, Union
+
+import torch
+from torch import nn
+from . import convnext
+from monai.networks.blocks import UpSample
+from monai.networks.layers.factories import Conv
+from monai.networks.layers.utils import get_act_layer
+from monai.networks.nets import EfficientNetBNFeatures
+from monai.networks.nets.basic_unet import UpCat
+from monai.utils import InterpolateMode
+
+__all__ = ["FlexibleUNet"]
+
+encoder_feature_channel = {
+    "efficientnet-b0": (16, 24, 40, 112, 320),
+    "efficientnet-b1": (16, 24, 40, 112, 320),
+    "efficientnet-b2": (16, 24, 48, 120, 352),
+    "efficientnet-b3": (24, 32, 48, 136, 384),
+    "efficientnet-b4": (24, 32, 56, 160, 448),
+    "efficientnet-b5": (24, 40, 64, 176, 512),
+    "efficientnet-b6": (32, 40, 72, 200, 576),
+    "efficientnet-b7": (32, 48, 80, 224, 640),
+    "efficientnet-b8": (32, 56, 88, 248, 704),
+    "efficientnet-l2": (72, 104, 176, 480, 1376),
+    "convnext_small": (96, 192, 384, 768),
+    "convnext_base": (128, 256, 512, 1024),
+    "van_b2": (64, 128, 320, 512),
+    "van_b1": (64, 128, 320, 512),
+}
+
+
+def _get_encoder_channels_by_backbone(backbone: str, in_channels: int = 3) -> tuple:
+    """
+    Get the encoder output channels by given backbone name.
+
+    Args:
+        backbone: name of backbone to generate features, can be from [efficientnet-b0, ..., efficientnet-b7].
+        in_channels: channel of input tensor, default to 3.
+
+    Returns:
+        A tuple of output feature map channels' length .
+    """
+    encoder_channel_tuple = encoder_feature_channel[backbone]
+    encoder_channel_list = [in_channels] + list(encoder_channel_tuple)
+    encoder_channel = tuple(encoder_channel_list)
+    return encoder_channel
+
+
+class UNetDecoder(nn.Module):
+    """
+    UNet Decoder.
+    This class refers to `segmentation_models.pytorch
+    <https://github.com/qubvel/segmentation_models.pytorch>`_.
+
+    Args:
+        spatial_dims: number of spatial dimensions.
+        encoder_channels: number of output channels for all feature maps in encoder.
+            `len(encoder_channels)` should be no less than 2.
+        decoder_channels: number of output channels for all feature maps in decoder.
+            `len(decoder_channels)` should equal to `len(encoder_channels) - 1`.
+        act: activation type and arguments.
+        norm: feature normalization type and arguments.
+        dropout: dropout ratio.
+        bias: whether to have a bias term in convolution blocks in this decoder.
+        upsample: upsampling mode, available options are
+            ``"deconv"``, ``"pixelshuffle"``, ``"nontrainable"``.
+        pre_conv: a conv block applied before upsampling.
+            Only used in the "nontrainable" or "pixelshuffle" mode.
+        interp_mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``}
+            Only used in the "nontrainable" mode.
+        align_corners: set the align_corners parameter for upsample. Defaults to True.
+            Only used in the "nontrainable" mode.
+        is_pad: whether to pad upsampling features to fit the encoder spatial dims.
+
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int,
+        encoder_channels: Sequence[int],
+        decoder_channels: Sequence[int],
+        act: Union[str, tuple],
+        norm: Union[str, tuple],
+        dropout: Union[float, tuple],
+        bias: bool,
+        upsample: str,
+        pre_conv: Optional[str],
+        interp_mode: str,
+        align_corners: Optional[bool],
+        is_pad: bool,
+    ):
+
+        super().__init__()
+        if len(encoder_channels) < 2:
+            raise ValueError("the length of `encoder_channels` should be no less than 2.")
+        if len(decoder_channels) != len(encoder_channels) - 1:
+            raise ValueError("`len(decoder_channels)` should equal to `len(encoder_channels) - 1`.")
+
+        in_channels = [encoder_channels[-1]] + list(decoder_channels[:-1])
+        skip_channels = list(encoder_channels[1:-1][::-1]) + [0]
+        halves = [True] * (len(skip_channels) - 1)
+        halves.append(False)
+        blocks = []
+        for in_chn, skip_chn, out_chn, halve in zip(in_channels, skip_channels, decoder_channels, halves):
+            blocks.append(
+                UpCat(
+                    spatial_dims=spatial_dims,
+                    in_chns=in_chn,
+                    cat_chns=skip_chn,
+                    out_chns=out_chn,
+                    act=act,
+                    norm=norm,
+                    dropout=dropout,
+                    bias=bias,
+                    upsample=upsample,
+                    pre_conv=pre_conv,
+                    interp_mode=interp_mode,
+                    align_corners=align_corners,
+                    halves=halve,
+                    is_pad=is_pad,
+                )
+            )
+        self.blocks = nn.ModuleList(blocks)
+
+    def forward(self, features: List[torch.Tensor], skip_connect: int = 3):
+        skips = features[:-1][::-1]
+        features = features[1:][::-1]
+
+        x = features[0]
+        for i, block in enumerate(self.blocks):
+            if i < skip_connect:
+                skip = skips[i]
+            else:
+                skip = None
+            x = block(x, skip)
+
+        return x
+
+
+class SegmentationHead(nn.Sequential):
+    """
+    Segmentation head.
+    This class refers to `segmentation_models.pytorch
+    <https://github.com/qubvel/segmentation_models.pytorch>`_.
+
+    Args:
+        spatial_dims: number of spatial dimensions.
+        in_channels: number of input channels for the block.
+        out_channels: number of output channels for the block.
+        kernel_size: kernel size for the conv layer.
+        act: activation type and arguments.
+        scale_factor: multiplier for spatial size. Has to match input size if it is a tuple.
+
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        act: Optional[Union[Tuple, str]] = None,
+        scale_factor: float = 1.0,
+    ):
+
+        conv_layer = Conv[Conv.CONV, spatial_dims](
+            in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=kernel_size // 2
+        )
+        up_layer: nn.Module = nn.Identity()
+        # if scale_factor > 1.0:
+        #     up_layer = UpSample(
+        #         in_channels=out_channels,
+        #         spatial_dims=spatial_dims,
+        #         scale_factor=scale_factor,
+        #         mode="deconv",
+        #         pre_conv=None,
+        #         interp_mode=InterpolateMode.LINEAR,
+        #     )
+        if scale_factor > 1.0:
+            up_layer = UpSample(
+                spatial_dims=spatial_dims,
+                scale_factor=scale_factor,
+                mode="nontrainable",
+                pre_conv=None,
+                interp_mode=InterpolateMode.LINEAR,
+            )
+        if act is not None:
+            act_layer = get_act_layer(act)
+        else:
+            act_layer = nn.Identity()
+        super().__init__(conv_layer, up_layer, act_layer)
+
+
+class FlexibleUNet_star(nn.Module):
+    """
+    A flexible implementation of UNet-like encoder-decoder architecture.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        backbone: str,
+        pretrained: bool = False,
+        decoder_channels: Tuple = (256, 128, 64, 32),
+        #decoder_channels: Tuple = (1024, 512, 256, 128),
+        spatial_dims: int = 2,
+        norm: Union[str, tuple] = ("batch", {"eps": 1e-3, "momentum": 0.1}),
+        act: Union[str, tuple] = ("relu", {"inplace": True}),
+        dropout: Union[float, tuple] = 0.0,
+        decoder_bias: bool = False,
+        upsample: str = "nontrainable",
+        interp_mode: str = "nearest",
+        is_pad: bool = True,
+        n_rays: int = 32,
+        prob_out_channels: int = 1,
+    ) -> None:
+        """
+        A flexible implement of UNet, in which the backbone/encoder can be replaced with
+        any efficient network. Currently the input must have a 2 or 3 spatial dimension
+        and the spatial size of each dimension must be a multiple of 32 if is pad parameter
+        is False
+
+        Args:
+            in_channels: number of input channels.
+            out_channels: number of output channels.
+            backbone: name of backbones to initialize, only support efficientnet right now,
+                can be from [efficientnet-b0,..., efficientnet-b8, efficientnet-l2].
+            pretrained: whether to initialize pretrained ImageNet weights, only available
+                for spatial_dims=2 and batch norm is used, default to False.
+            decoder_channels: number of output channels for all feature maps in decoder.
+                `len(decoder_channels)` should equal to `len(encoder_channels) - 1`,default
+                to (256, 128, 64, 32, 16).
+            spatial_dims: number of spatial dimensions, default to 2.
+            norm: normalization type and arguments, default to ("batch", {"eps": 1e-3,
+                "momentum": 0.1}).
+            act: activation type and arguments, default to ("relu", {"inplace": True}).
+            dropout: dropout ratio, default to 0.0.
+            decoder_bias: whether to have a bias term in decoder's convolution blocks.
+            upsample: upsampling mode, available options are``"deconv"``, ``"pixelshuffle"``,
+                ``"nontrainable"``.
+            interp_mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``}
+                Only used in the "nontrainable" mode.
+            is_pad: whether to pad upsampling features to fit features from encoder. Default to True.
+                If this parameter is set to "True", the spatial dim of network input can be arbitary
+                size, which is not supported by TensorRT. Otherwise, it must be a multiple of 32.
+        """
+        super().__init__()
+
+        if backbone not in encoder_feature_channel:
+            raise ValueError(f"invalid model_name {backbone} found, must be one of {encoder_feature_channel.keys()}.")
+
+        if spatial_dims not in (2, 3):
+            raise ValueError("spatial_dims can only be 2 or 3.")
+
+        adv_prop = "ap" in backbone
+
+        self.backbone = backbone
+        self.spatial_dims = spatial_dims
+        model_name = backbone
+        encoder_channels = _get_encoder_channels_by_backbone(backbone, in_channels)
+
+        self.encoder = convnext.convnext_small(pretrained=False,in_22k=True)
+
+        self.decoder = UNetDecoder(
+            spatial_dims=spatial_dims,
+            encoder_channels=encoder_channels,
+            decoder_channels=decoder_channels,
+            act=act,
+            norm=norm,
+            dropout=dropout,
+            bias=decoder_bias,
+            upsample=upsample,
+            interp_mode=interp_mode,
+            pre_conv=None,
+            align_corners=None,
+            is_pad=is_pad,
+        )
+        self.dist_head = SegmentationHead(
+            spatial_dims=spatial_dims,
+            in_channels=decoder_channels[-1],
+            out_channels=n_rays,
+            kernel_size=1,
+            act='relu',
+            scale_factor = 2,
+        )
+        self.prob_head = SegmentationHead(
+            spatial_dims=spatial_dims,
+            in_channels=decoder_channels[-1],
+            out_channels=prob_out_channels,
+            kernel_size=1,
+            act='sigmoid',
+            scale_factor = 2,
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        """
+        Do a typical encoder-decoder-header inference.
+
+        Args:
+            inputs: input should have spatially N dimensions ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N])``,
+                N is defined by `dimensions`.
+
+        Returns:
+            A torch Tensor of "raw" predictions in shape ``(Batch, out_channels, dim_0[, dim_1, ..., dim_N])``.
+
+        """
+        x = inputs
+        enc_out = self.encoder(x)
+        decoder_out = self.decoder(enc_out)
+
+        dist = self.dist_head(decoder_out)
+        prob = self.prob_head(decoder_out)
+
+        return dist,prob
+        
+
+
+class FlexibleUNet_hv(nn.Module):
+    """
+    A flexible implementation of UNet-like encoder-decoder architecture.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        backbone: str,
+        pretrained: bool = False,
+        decoder_channels: Tuple = (1024, 512, 256, 128),
+        spatial_dims: int = 2,
+        norm: Union[str, tuple] = ("batch", {"eps": 1e-3, "momentum": 0.1}),
+        act: Union[str, tuple] = ("relu", {"inplace": True}),
+        dropout: Union[float, tuple] = 0.0,
+        decoder_bias: bool = False,
+        upsample: str = "nontrainable",
+        interp_mode: str = "nearest",
+        is_pad: bool = True,
+        n_rays: int = 32,
+        prob_out_channels: int = 1,
+    ) -> None:
+        """
+        A flexible implement of UNet, in which the backbone/encoder can be replaced with
+        any efficient network. Currently the input must have a 2 or 3 spatial dimension
+        and the spatial size of each dimension must be a multiple of 32 if is pad parameter
+        is False
+
+        Args:
+            in_channels: number of input channels.
+            out_channels: number of output channels.
+            backbone: name of backbones to initialize, only support efficientnet right now,
+                can be from [efficientnet-b0,..., efficientnet-b8, efficientnet-l2].
+            pretrained: whether to initialize pretrained ImageNet weights, only available
+                for spatial_dims=2 and batch norm is used, default to False.
+            decoder_channels: number of output channels for all feature maps in decoder.
+                `len(decoder_channels)` should equal to `len(encoder_channels) - 1`,default
+                to (256, 128, 64, 32, 16).
+            spatial_dims: number of spatial dimensions, default to 2.
+            norm: normalization type and arguments, default to ("batch", {"eps": 1e-3,
+                "momentum": 0.1}).
+            act: activation type and arguments, default to ("relu", {"inplace": True}).
+            dropout: dropout ratio, default to 0.0.
+            decoder_bias: whether to have a bias term in decoder's convolution blocks.
+            upsample: upsampling mode, available options are``"deconv"``, ``"pixelshuffle"``,
+                ``"nontrainable"``.
+            interp_mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``}
+                Only used in the "nontrainable" mode.
+            is_pad: whether to pad upsampling features to fit features from encoder. Default to True.
+                If this parameter is set to "True", the spatial dim of network input can be arbitary
+                size, which is not supported by TensorRT. Otherwise, it must be a multiple of 32.
+        """
+        super().__init__()
+
+        if backbone not in encoder_feature_channel:
+            raise ValueError(f"invalid model_name {backbone} found, must be one of {encoder_feature_channel.keys()}.")
+
+        if spatial_dims not in (2, 3):
+            raise ValueError("spatial_dims can only be 2 or 3.")
+
+        adv_prop = "ap" in backbone
+
+        self.backbone = backbone
+        self.spatial_dims = spatial_dims
+        model_name = backbone
+        encoder_channels = _get_encoder_channels_by_backbone(backbone, in_channels)
+        self.encoder = convnext.convnext_small(pretrained=False,in_22k=True)
+        self.decoder = UNetDecoder(
+            spatial_dims=spatial_dims,
+            encoder_channels=encoder_channels,
+            decoder_channels=decoder_channels,
+            act=act,
+            norm=norm,
+            dropout=dropout,
+            bias=decoder_bias,
+            upsample=upsample,
+            interp_mode=interp_mode,
+            pre_conv=None,
+            align_corners=None,
+            is_pad=is_pad,
+        )
+        self.dist_head = SegmentationHead(
+            spatial_dims=spatial_dims,
+            in_channels=decoder_channels[-1],
+            out_channels=n_rays,
+            kernel_size=1,
+            act=None,
+            scale_factor = 2,
+        )
+        self.prob_head = SegmentationHead(
+            spatial_dims=spatial_dims,
+            in_channels=decoder_channels[-1],
+            out_channels=prob_out_channels,
+            kernel_size=1,
+            act='sigmoid',
+            scale_factor = 2,
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        """
+        Do a typical encoder-decoder-header inference.
+
+        Args:
+            inputs: input should have spatially N dimensions ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N])``,
+                N is defined by `dimensions`.
+
+        Returns:
+            A torch Tensor of "raw" predictions in shape ``(Batch, out_channels, dim_0[, dim_1, ..., dim_N])``.
+
+        """
+        x = inputs
+        enc_out = self.encoder(x)
+        decoder_out = self.decoder(enc_out)
+        dist = self.dist_head(decoder_out)
+        prob = self.prob_head(decoder_out)
+        return dist,prob

overlay.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+###overlay
+import cv2
+import math
+import random
+import colorsys
+import numpy as np
+import itertools
+import matplotlib.pyplot as plt
+from matplotlib import cm
+import os
+import scipy.io as io
+def get_bounding_box(img):
+    """Get bounding box coordinate information."""
+    rows = np.any(img, axis=1)
+    cols = np.any(img, axis=0)
+    rmin, rmax = np.where(rows)[0][[0, -1]]
+    cmin, cmax = np.where(cols)[0][[0, -1]]
+    # due to python indexing, need to add 1 to max
+    # else accessing will be 1px in the box, not out
+    rmax += 1
+    cmax += 1
+    return [rmin, rmax, cmin, cmax]
+####
+def colorize(ch, vmin, vmax):
+    """Will clamp value value outside the provided range to vmax and vmin."""
+    cmap = plt.get_cmap("jet")
+    ch = np.squeeze(ch.astype("float32"))
+    vmin = vmin if vmin is not None else ch.min()
+    vmax = vmax if vmax is not None else ch.max()
+    ch[ch > vmax] = vmax  # clamp value
+    ch[ch < vmin] = vmin
+    ch = (ch - vmin) / (vmax - vmin + 1.0e-16)
+    # take RGB from RGBA heat map
+    ch_cmap = (cmap(ch)[..., :3] * 255).astype("uint8")
+    return ch_cmap
+
+
+####
+def random_colors(N, bright=True):
+    """Generate random colors.
+    
+    To get visually distinct colors, generate them in HSV space then
+    convert to RGB.
+    """
+    brightness = 1.0 if bright else 0.7
+    hsv = [(i / N, 1, brightness) for i in range(N)]
+    colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv))
+    random.shuffle(colors)
+    return colors
+
+
+####
+def visualize_instances_map(
+    input_image, inst_map, type_map=None, type_colour=None, line_thickness=2
+):
+    """Overlays segmentation results on image as contours.
+
+    Args:
+        input_image: input image
+        inst_map: instance mask with unique value for every object
+        type_map: type mask with unique value for every class
+        type_colour: a dict of {type : colour} , `type` is from 0-N
+                     and `colour` is a tuple of (R, G, B)
+        line_thickness: line thickness of contours
+
+    Returns:
+        overlay: output image with segmentation overlay as contours
+    """
+    overlay = np.copy((input_image).astype(np.uint8))
+
+    inst_list = list(np.unique(inst_map))  # get list of instances
+    inst_list.remove(0)  # remove background
+
+    inst_rng_colors = random_colors(len(inst_list))
+    inst_rng_colors = np.array(inst_rng_colors) * 255
+    inst_rng_colors = inst_rng_colors.astype(np.uint8)
+
+    for inst_idx, inst_id in enumerate(inst_list):
+        inst_map_mask = np.array(inst_map == inst_id, np.uint8)  # get single object
+        y1, y2, x1, x2 = get_bounding_box(inst_map_mask)
+        y1 = y1 - 2 if y1 - 2 >= 0 else y1
+        x1 = x1 - 2 if x1 - 2 >= 0 else x1
+        x2 = x2 + 2 if x2 + 2 <= inst_map.shape[1] - 1 else x2
+        y2 = y2 + 2 if y2 + 2 <= inst_map.shape[0] - 1 else y2
+        inst_map_crop = inst_map_mask[y1:y2, x1:x2]
+        contours_crop = cv2.findContours(
+            inst_map_crop, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+        )
+        # only has 1 instance per map, no need to check #contour detected by opencv
+        #print(contours_crop)
+        contours_crop = np.squeeze(
+            contours_crop[0][0].astype("int32")
+        )  # * opencv protocol format may break
+        
+        if len(contours_crop.shape) == 1:
+            contours_crop = contours_crop.reshape(1,-1)
+        #print(contours_crop.shape)
+        contours_crop += np.asarray([[x1, y1]])  # index correction
+        if type_map is not None:
+            type_map_crop = type_map[y1:y2, x1:x2]
+            type_id = np.unique(type_map_crop).max()  # non-zero
+            inst_colour = type_colour[type_id]
+        else:
+            inst_colour = (inst_rng_colors[inst_idx]).tolist()
+        cv2.drawContours(overlay, [contours_crop], -1, inst_colour, line_thickness)
+    return overlay
+
+
+# In[ ]:
+
+
+
+

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d978b42e9e63e949f0dcd3685be14146b6c5b5bfb48f703bfbc308b4ac190b64
+size 135

requirements.txt

@@ -0,0 +1,37 @@
+gputools==0.2.13
+h5py==3.7.0
+huggingface-hub==0.10.1
+imagecodecs
+imageio==2.22.2
+importlib-metadata==5.0.0
+kiwisolver==1.4.4
+llvmlite==0.39.1
+Mako==1.2.3
+Markdown==3.4.1
+MarkupSafe==2.1.1
+matplotlib==3.6.1
+mkl-fft==1.3.1
+mkl-service==2.4.0
+monai==1.0.0
+networkx==2.8.7
+numba==0.56.3
+numexpr
+numpy
+oauthlib==3.2.2
+opencv-python==4.6.0.66
+packaging
+pandas==1.4.4
+Pillow==9.2.0
+scikit-image==0.19.3
+scipy==1.9.2
+stardist==0.8.3
+tensorboard==2.10.1
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.1
+tifffile==2022.10.10
+timm==0.6.11
+torch==1.12.1
+torchaudio==0.12.1
+torchvision==0.13.1
+tqdm==4.64.1
+

sribd_cellseg_models.py

@@ -0,0 +1,100 @@
+
+import os
+join = os.path.join
+import argparse
+import numpy as np
+import torch
+import torch.nn as nn
+from collections import OrderedDict
+from torchvision import datasets, models, transforms
+from classifiers import resnet10, resnet18
+
+from utils_modify import sliding_window_inference,sliding_window_inference_large,__proc_np_hv
+from PIL import Image
+import torch.nn.functional as F
+from skimage import io, segmentation, morphology, measure, exposure
+import tifffile as tif
+from models.flexible_unet_convnext import FlexibleUNet_star,FlexibleUNet_hv 
+from transformers import PretrainedConfig
+from typing import List
+from transformers import PreTrainedModel
+from huggingface_hub import PyTorchModelHubMixin
+from torch import nn
+class ModelConfig(PretrainedConfig):
+    model_type = "cell_sribd"
+    def __init__(
+        self,
+        version = 1,
+        input_channels: int = 3,
+        roi_size: int = 512,
+        overlap: float = 0.5,
+        device: str = 'cpu',
+        **kwargs,
+    ):
+        
+        self.device = device
+        self.roi_size = (roi_size, roi_size)
+        self.input_channels = input_channels
+        self.overlap = overlap
+        self.np_thres, self.ksize, self.overall_thres, self.obj_size_thres = 0.6, 15, 0.4, 100
+        self.n_rays = 32
+        self.sw_batch_size = 4
+        self.num_classes= 4
+        self.block_size = 2048
+        self.min_overlap = 128
+        self.context = 128
+        super().__init__(**kwargs)
+        
+        
+class MultiStreamCellSegModel(PreTrainedModel):
+    config_class = ModelConfig
+    #print(config.input_channels)
+    def __init__(self, config):
+        super().__init__(config)
+        #print(config.input_channels)
+        self.config = config
+        self.cls_model = resnet18()
+        self.model0 = FlexibleUNet_star(in_channels=config.input_channels,out_channels=config.n_rays+1,backbone='convnext_small',pretrained=False,n_rays=config.n_rays,prob_out_channels=1,)
+        self.model1 = FlexibleUNet_star(in_channels=config.input_channels,out_channels=config.n_rays+1,backbone='convnext_small',pretrained=False,n_rays=config.n_rays,prob_out_channels=1,)
+        self.model2 = FlexibleUNet_star(in_channels=config.input_channels,out_channels=config.n_rays+1,backbone='convnext_small',pretrained=False,n_rays=config.n_rays,prob_out_channels=1,)
+        self.model3 = FlexibleUNet_hv(in_channels=config.input_channels,out_channels=2+2,backbone='convnext_small',pretrained=False,n_rays=2,prob_out_channels=2,)
+        self.preprocess=transforms.Compose([
+            transforms.Resize(size=256),
+            transforms.CenterCrop(size=224),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])])
+    def load_checkpoints(self,checkpoints):
+        self.cls_model.load_state_dict(checkpoints['cls_model'])
+        self.model0.load_state_dict(checkpoints['class1_model']['model_state_dict'])
+        self.model1.load_state_dict(checkpoints['class2_model']['model_state_dict'])
+        self.model2.load_state_dict(checkpoints['class3_model']['model_state_dict'])
+        self.model3.load_state_dict(checkpoints['class4_model'])
+        
+    def forward(self, pre_img_data):
+        inputs=self.preprocess(Image.fromarray(pre_img_data)).unsqueeze(0)
+        outputs = self.cls_model(inputs)
+        _, preds = torch.max(outputs, 1)    
+        label=preds[0].cpu().numpy()
+        test_npy01 = pre_img_data
+        if label in [0,1,2]:
+            if label == 0:
+                output_label = sliding_window_inference_large(test_npy01,self.config.block_size,self.config.min_overlap,self.config.context, self.config.roi_size,self.config.sw_batch_size,predictor=self.model0,device=self.config.device)
+            elif label == 1:
+                output_label = sliding_window_inference_large(test_npy01,self.config.block_size,self.config.min_overlap,self.config.context, self.config.roi_size,self.config.sw_batch_size,predictor=self.model1,device=self.config.device)
+            elif label == 2:
+                output_label = sliding_window_inference_large(test_npy01,self.config.block_size,self.config.min_overlap,self.config.context, self.config.roi_size,self.config.sw_batch_size,predictor=self.model2,device=self.config.device)
+        else:
+            test_tensor = torch.from_numpy(np.expand_dims(test_npy01, 0)).permute(0, 3, 1, 2).type(torch.FloatTensor)
+
+            output_hv, output_np = sliding_window_inference(test_tensor, self.config.roi, self.config.sw_batch_size, self.model3, overlap=self.config.overlap,device=self.config.device)
+            pred_dict = {'np': output_np, 'hv': output_hv}
+            pred_dict = OrderedDict(
+                    [[k, v.permute(0, 2, 3, 1).contiguous()] for k, v in pred_dict.items()]  # NHWC
+                )
+            pred_dict["np"] = F.softmax(pred_dict["np"], dim=-1)[..., 1:]
+            pred_output = torch.cat(list(pred_dict.values()), -1).cpu().numpy() # NHW3
+            pred_map = np.squeeze(pred_output) # HW3
+            pred_inst = __proc_np_hv(pred_map, self.config.np_thres, self.config.ksize, self.config.overall_thres, self.config.obj_size_thres)
+            raw_pred_shape = pred_inst.shape[:2]
+            output_label = pred_inst
+        return output_label       

stardist_pkg/__init__.py

@@ -0,0 +1,26 @@
+from __future__ import absolute_import, print_function
+
+import warnings
+def format_warning(message, category, filename, lineno, line=''):
+    import pathlib
+    return f"{pathlib.Path(filename).name} ({lineno}): {message}\n"
+warnings.formatwarning = format_warning
+del warnings
+
+from .version import __version__
+
+# TODO: which functions to expose here? all?
+from .nms import non_maximum_suppression
+from .utils import edt_prob, fill_label_holes, sample_points, calculate_extents, export_imagej_rois, gputools_available
+from .geometry import star_dist,   polygons_to_label,   relabel_image_stardist, ray_angles, dist_to_coord
+from .sample_patches import sample_patches
+from .bioimageio_utils import export_bioimageio, import_bioimageio
+
+def _py_deprecation(ver_python=(3,6), ver_stardist='0.9.0'):
+     import sys
+     from distutils.version import LooseVersion
+     if sys.version_info[:2] == ver_python and LooseVersion(__version__) < LooseVersion(ver_stardist):
+         print(f"You are using Python {ver_python[0]}.{ver_python[1]}, which will no longer be supported in StarDist {ver_stardist}.\n"
+               f"→ Please upgrade to Python {ver_python[0]}.{ver_python[1]+1} or later.", file=sys.stderr, flush=True)
+_py_deprecation()
+del _py_deprecation

stardist_pkg/big.py

@@ -0,0 +1,601 @@
+import numpy as np
+import warnings
+import math
+from tqdm import tqdm
+from skimage.measure import regionprops
+from skimage.draw import polygon
+from csbdeep.utils import _raise, axes_check_and_normalize, axes_dict
+from itertools import product
+
+
+
+
+OBJECT_KEYS = set(('prob', 'points', 'coord', 'dist', 'class_prob', 'class_id'))
+COORD_KEYS = set(('points', 'coord'))
+
+
+
+class Block:
+    """One-dimensional block as part of a chain.
+
+    There are no explicit start and end positions. Instead, each block is
+    aware of its predecessor and successor and derives such things (recursively)
+    based on its neighbors.
+
+    Blocks overlap with one another (at least min_overlap + 2*context) and
+    have a read region (the entire block) and a write region (ignoring context).
+    Given a query interval, Block.is_responsible will return true for only one
+    block of a chain (or raise an exception if the interval is larger than
+    min_overlap or even the entire block without context).
+
+    """
+    def __init__(self, size, min_overlap, context, pred):
+        self.size = int(size)
+        self.min_overlap = int(min_overlap)
+        self.context = int(context)
+        self.pred = pred
+        self.succ = None
+        assert 0 <= self.min_overlap + 2*self.context < self.size
+        self.stride = self.size - (self.min_overlap + 2*self.context)
+        self._start = 0
+        self._frozen = False
+
+    @property
+    def start(self):
+        return self._start if (self.frozen or self.at_begin) else self.pred.succ_start
+
+    @property
+    def end(self):
+        return self.start + self.size
+
+    @property
+    def succ_start(self):
+        return self.start + self.stride
+
+    def add_succ(self):
+        assert self.succ is None and not self.frozen
+        self.succ = Block(self.size, self.min_overlap, self.context, self)
+        return self.succ
+
+    def decrease_stride(self, amount):
+        amount = int(amount)
+        assert 0 <= amount < self.stride and not self.frozen
+        self.stride -= amount
+
+    def freeze(self):
+        """Call on first block to freeze entire chain (after construction is done)"""
+        assert not self.frozen and (self.at_begin or self.pred.frozen)
+        self._start = self.start
+        self._frozen = True
+        if not self.at_end:
+            self.succ.freeze()
+
+    @property
+    def slice_read(self):
+        return slice(self.start, self.end)
+
+    @property
+    def slice_crop_context(self):
+        """Crop context relative to read region"""
+        return slice(self.context_start, self.size - self.context_end)
+
+    @property
+    def slice_write(self):
+        return slice(self.start + self.context_start, self.end - self.context_end)
+
+    def is_responsible(self, bbox):
+        """Responsibility for query interval bbox, which is assumed to be smaller than min_overlap.
+
+        If the assumption is met, only one block of a chain will return true.
+        If violated, one or more blocks of a chain may raise a NotFullyVisible exception.
+        The exception will have an argument that is
+            False if bbox is larger than min_overlap, and
+            True if bbox is even larger than the entire block without context.
+
+        bbox: (int,int)
+            1D bounding box interval with coordinates relative to size without context
+
+        """
+        bmin, bmax = bbox
+
+        r_start = 0 if self.at_begin else (self.pred.overlap - self.pred.context_end - self.context_start)
+        r_end = self.size - self.context_start - self.context_end
+        assert 0 <= bmin < bmax <= r_end
+
+        # assert not (bmin == 0 and bmax >= r_start and not self.at_begin), [(r_start,r_end), bbox, self]
+
+        if bmin == 0 and bmax >= r_start:
+            if bmax == r_end:
+                # object spans the entire block, i.e. is probably larger than size (minus the context)
+                raise NotFullyVisible(True)
+            if not self.at_begin:
+                # object spans the entire overlap region, i.e. is only partially visible here and also by the predecessor block
+                raise NotFullyVisible(False)
+
+        # object ends before responsible region start
+        if bmax < r_start: return False
+        # object touches the end of the responsible region (only take if at end)
+        if bmax == r_end and not self.at_end: return False
+        return True
+
+    # ------------------------
+
+    @property
+    def frozen(self):
+        return self._frozen
+
+    @property
+    def at_begin(self):
+        return self.pred is None
+
+    @property
+    def at_end(self):
+        return self.succ is None
+
+    @property
+    def overlap(self):
+        return self.size - self.stride
+
+    @property
+    def context_start(self):
+        return 0 if self.at_begin else self.context
+
+    @property
+    def context_end(self):
+        return 0 if self.at_end else self.context
+
+    def __repr__(self):
+        shared  = f'{self.start:03}:{self.end:03}'
+        shared += f', size={self.context_start}-{self.size-self.context_start-self.context_end}-{self.context_end}'
+        if self.at_end:
+            return f'{self.__class__.__name__}({shared})'
+        else:
+            return f'{self.__class__.__name__}({shared}, overlap={self.overlap}/{self.overlap-self.context_start-self.context_end})'
+
+    @property
+    def chain(self):
+        blocks = [self]
+        while not blocks[-1].at_end:
+            blocks.append(blocks[-1].succ)
+        return blocks
+
+    def __iter__(self):
+        return iter(self.chain)
+
+    # ------------------------
+
+    @staticmethod
+    def cover(size, block_size, min_overlap, context, grid=1, verbose=True):
+        """Return chain of grid-aligned blocks to cover the interval [0,size].
+
+        Parameters block_size, min_overlap, and context will be used
+        for all blocks of the chain. Only the size of the last block
+        may differ.
+
+        Except for the last block, start and end positions of all blocks will
+        be multiples of grid. To that end, the provided block parameters may
+        be increased to achieve that.
+
+        Note that parameters must be chosen such that the write regions of only
+        neighboring blocks are overlapping.
+
+        """
+        assert 0 <= min_overlap+2*context < block_size <= size
+        assert 0 < grid <= block_size
+        block_size = _grid_divisible(grid, block_size, name='block_size', verbose=verbose)
+        min_overlap = _grid_divisible(grid, min_overlap, name='min_overlap', verbose=verbose)
+        context = _grid_divisible(grid, context, name='context', verbose=verbose)
+
+        # allow size not to be divisible by grid
+        size_orig = size
+        size = _grid_divisible(grid, size, name='size', verbose=False)
+
+        # divide all sizes by grid
+        assert all(v % grid == 0 for v in (size, block_size, min_overlap, context))
+        size //= grid
+        block_size //= grid
+        min_overlap //= grid
+        context //= grid
+
+        # compute cover in grid-multiples
+        t = first = Block(block_size, min_overlap, context, None)
+        while t.end < size:
+            t = t.add_succ()
+        last = t
+
+        # [print(t) for t in first]
+
+        # move blocks around to make it fit
+        excess = last.end - size
+        t = first
+        while excess > 0:
+            t.decrease_stride(1)
+            excess -= 1
+            t = t.succ
+            if (t == last): t = first
+
+        # make a copy of the cover and multiply sizes by grid
+        if grid > 1:
+            size *= grid
+            block_size *= grid
+            min_overlap *= grid
+            context *= grid
+            #
+            _t = _first = first
+            t = first = Block(block_size, min_overlap, context, None)
+            t.stride = _t.stride*grid
+            while not _t.at_end:
+                _t = _t.succ
+                t = t.add_succ()
+                t.stride = _t.stride*grid
+            last = t
+
+            # change size of last block
+            # will be padded internally to the same size
+            # as the others by model.predict_instances
+            size_delta = size - size_orig
+            last.size -= size_delta
+            assert 0 <= size_delta < grid
+
+        # for efficiency (to not determine starts recursively from now on)
+        first.freeze()
+
+        blocks = first.chain
+
+        # sanity checks
+        assert first.start == 0 and last.end == size_orig
+        assert all(t.overlap-2*context >= min_overlap for t in blocks if t != last)
+        assert all(t.start % grid == 0 and t.end % grid == 0 for t in blocks if t != last)
+        # print(); [print(t) for t in first]
+
+        # only neighboring blocks should be overlapping
+        if len(blocks) >= 3:
+            for t in blocks[:-2]:
+                assert t.slice_write.stop <= t.succ.succ.slice_write.start
+
+        return blocks
+
+
+
+class BlockND:
+    """N-dimensional block.
+
+    Each BlockND simply consists of a 1-dimensional Block per axis and also
+    has an id (which should be unique). The n-dimensional region represented
+    by each BlockND is the intersection of all 1D Blocks per axis.
+
+    Also see `Block`.
+
+    """
+    def __init__(self, id, blocks, axes):
+        self.id = id
+        self.blocks = tuple(blocks)
+        self.axes = axes_check_and_normalize(axes, length=len(self.blocks))
+        self.axis_to_block = dict(zip(self.axes,self.blocks))
+
+    def blocks_for_axes(self, axes=None):
+        axes = self.axes if axes is None else axes_check_and_normalize(axes)
+        return tuple(self.axis_to_block[a] for a in axes)
+
+    def slice_read(self, axes=None):
+        return tuple(t.slice_read for t in self.blocks_for_axes(axes))
+
+    def slice_crop_context(self, axes=None):
+        return tuple(t.slice_crop_context for t in self.blocks_for_axes(axes))
+
+    def slice_write(self, axes=None):
+        return tuple(t.slice_write for t in self.blocks_for_axes(axes))
+
+    def read(self, x, axes=None):
+        """Read block "read region" from x (numpy.ndarray or similar)"""
+        return x[self.slice_read(axes)]
+
+    def crop_context(self, labels, axes=None):
+        return labels[self.slice_crop_context(axes)]
+
+    def write(self, x, labels, axes=None):
+        """Write (only entries > 0 of) labels to block "write region" of x (numpy.ndarray or similar)"""
+        s = self.slice_write(axes)
+        mask = labels > 0
+        # x[s][mask] = labels[mask] # doesn't work with zarr
+        region = x[s]
+        region[mask] = labels[mask]
+        x[s] = region
+
+    def is_responsible(self, slices, axes=None):
+        return all(t.is_responsible((s.start,s.stop)) for t,s in zip(self.blocks_for_axes(axes),slices))
+
+    def __repr__(self):
+        slices =  ','.join(f'{a}={t.start:03}:{t.end:03}' for t,a in zip(self.blocks,self.axes))
+        return f'{self.__class__.__name__}({self.id}|{slices})'
+
+    def __iter__(self):
+        return iter(self.blocks)
+
+    # ------------------------
+
+    def filter_objects(self, labels, polys, axes=None):
+        """Filter out objects that block is not responsible for.
+
+        Given label image 'labels' and dictionary 'polys' of polygon/polyhedron objects,
+        only retain those objects that this block is responsible for.
+
+        This function will return a pair (labels, polys) of the modified label image and dictionary.
+        It will raise a RuntimeError if an object is found in the overlap area
+        of neighboring blocks that violates the assumption to be smaller than 'min_overlap'.
+
+        If parameter 'polys' is None, only the filtered label image will be returned.
+
+        Notes
+        -----
+        - Important: It is assumed that the object label ids in 'labels' and
+          the entries in 'polys' are sorted in the same way.
+        - Does not modify 'labels' and 'polys', but returns modified copies.
+
+        Example
+        -------
+        >>> labels, polys = model.predict_instances(block.read(img))
+        >>> labels = block.crop_context(labels)
+        >>> labels, polys = block.filter_objects(labels, polys)
+
+        """
+        # TODO: option to update labels in-place
+        assert np.issubdtype(labels.dtype, np.integer)
+        ndim = len(self.blocks_for_axes(axes))
+        assert ndim in (2,3)
+        assert labels.ndim == ndim and labels.shape == tuple(s.stop-s.start for s in self.slice_crop_context(axes))
+
+        labels_filtered = np.zeros_like(labels)
+        # problem_ids = []
+        for r in regionprops(labels):
+            slices = tuple(slice(r.bbox[i],r.bbox[i+labels.ndim]) for i in range(labels.ndim))
+            try:
+                if self.is_responsible(slices, axes):
+                    labels_filtered[slices][r.image] = r.label
+            except NotFullyVisible as e:
+                # shape_block_write = tuple(s.stop-s.start for s in self.slice_write(axes))
+                shape_object = tuple(s.stop-s.start for s in slices)
+                shape_min_overlap = tuple(t.min_overlap for t in self.blocks_for_axes(axes))
+                raise RuntimeError(f"Found object of shape {shape_object}, which violates the assumption of being smaller than 'min_overlap' {shape_min_overlap}. Increase 'min_overlap' to avoid this problem.")
+
+                # if e.args[0]: # object larger than block write region
+                #     assert any(o >= b for o,b in zip(shape_object,shape_block_write))
+                #     # problem, since this object will probably be saved by another block too
+                #     raise RuntimeError(f"Found object of shape {shape_object}, larger than an entire block's write region of shape {shape_block_write}. Increase 'block_size' to avoid this problem.")
+                #     # print("found object larger than 'block_size'")
+                # else:
+                #     assert any(o >= b for o,b in zip(shape_object,shape_min_overlap))
+                #     # print("found object larger than 'min_overlap'")
+
+                # # keep object, because will be dealt with later, i.e.
+                # # render the poly again into the label image, but this is not
+                # # ideal since the assumption is that the object outside that
+                # # region is not reliable because it's in the context
+                # labels_filtered[slices][r.image] = r.label
+                # problem_ids.append(r.label)
+
+        if polys is None:
+            # assert len(problem_ids) == 0
+            return labels_filtered
+        else:
+            # it is assumed that ids in 'labels' map to entries in 'polys'
+            assert isinstance(polys,dict) and any(k in polys for k in COORD_KEYS)
+            filtered_labels = np.unique(labels_filtered)
+            filtered_ind = [i-1 for i in filtered_labels if i > 0]
+            polys_out = {k: (v[filtered_ind] if k in OBJECT_KEYS else v) for k,v in polys.items()}
+            for k in COORD_KEYS:
+                if k in polys_out.keys():
+                    polys_out[k] = self.translate_coordinates(polys_out[k], axes=axes)
+
+        return labels_filtered, polys_out#, tuple(problem_ids)
+
+    def translate_coordinates(self, coordinates, axes=None):
+        """Translate local block coordinates (of read region) to global ones based on block position"""
+        ndim = len(self.blocks_for_axes(axes))
+        assert isinstance(coordinates, np.ndarray) and coordinates.ndim >= 2 and coordinates.shape[1] == ndim
+        start = [s.start for s in self.slice_read(axes)]
+        shape = tuple(1 if d!=1 else ndim for d in range(coordinates.ndim))
+        start = np.array(start).reshape(shape)
+        return coordinates + start
+
+    # ------------------------
+
+    @staticmethod
+    def cover(shape, axes, block_size, min_overlap, context, grid=1):
+        """Return grid-aligned n-dimensional blocks to cover region
+        of the given shape with axes semantics.
+
+        Parameters block_size, min_overlap, and context can be different per
+        dimension/axis (if provided as list) or the same (if provided as
+        scalar value).
+
+        Also see `Block.cover`.
+
+        """
+        shape = tuple(shape)
+        n = len(shape)
+        axes = axes_check_and_normalize(axes, length=n)
+        if np.isscalar(block_size):  block_size  = n*[block_size]
+        if np.isscalar(min_overlap): min_overlap = n*[min_overlap]
+        if np.isscalar(context):     context     = n*[context]
+        if np.isscalar(grid):        grid        = n*[grid]
+        assert n == len(block_size) == len(min_overlap) == len(context) == len(grid)
+
+        # compute cover for each dimension
+        cover_1d = [Block.cover(*args) for args in zip(shape, block_size, min_overlap, context, grid)]
+        # return cover as Cartesian product of 1-dimensional blocks
+        return tuple(BlockND(i,blocks,axes) for i,blocks in enumerate(product(*cover_1d)))
+
+
+
+class Polygon:
+
+    def __init__(self, coord, bbox=None, shape_max=None):
+        self.bbox = self.coords_bbox(coord, shape_max=shape_max) if bbox is None else bbox
+        self.coord = coord - np.array([r[0] for r in self.bbox]).reshape(2,1)
+        self.slice = tuple(slice(*r) for r in self.bbox)
+        self.shape = tuple(r[1]-r[0] for r in self.bbox)
+        rr,cc = polygon(*self.coord, self.shape)
+        self.mask = np.zeros(self.shape, bool)
+        self.mask[rr,cc] = True
+
+    @staticmethod
+    def coords_bbox(*coords, shape_max=None):
+        assert all(isinstance(c, np.ndarray) and c.ndim==2 and c.shape[0]==2 for c in coords)
+        if shape_max is None:
+            shape_max = (np.inf, np.inf)
+        coord = np.concatenate(coords, axis=1)
+        mins = np.maximum(0,         np.floor(np.min(coord,axis=1))).astype(int)
+        maxs = np.minimum(shape_max, np.ceil (np.max(coord,axis=1))).astype(int)
+        return tuple(zip(tuple(mins),tuple(maxs)))
+
+
+
+class Polyhedron:
+
+    def __init__(self, dist, origin, rays, bbox=None, shape_max=None):
+        self.bbox = self.coords_bbox((dist, origin), rays=rays, shape_max=shape_max) if bbox is None else bbox
+        self.slice = tuple(slice(*r) for r in self.bbox)
+        self.shape = tuple(r[1]-r[0] for r in self.bbox)
+        _origin = origin.reshape(1,3) - np.array([r[0] for r in self.bbox]).reshape(1,3)
+        self.mask = polyhedron_to_label(dist[np.newaxis], _origin, rays, shape=self.shape, verbose=False).astype(bool)
+
+    @staticmethod
+    def coords_bbox(*dist_origin, rays, shape_max=None):
+        dists, points = zip(*dist_origin)
+        assert all(isinstance(d, np.ndarray) and d.ndim==1 and len(d)==len(rays) for d in dists)
+        assert all(isinstance(p, np.ndarray) and p.ndim==1 and len(p)==3 for p in points)
+        dists, points, verts = np.stack(dists)[...,np.newaxis], np.stack(points)[:,np.newaxis], rays.vertices[np.newaxis]
+        coord = dists * verts + points
+        coord = np.concatenate(coord, axis=0)
+        if shape_max is None:
+            shape_max = (np.inf, np.inf, np.inf)
+        mins = np.maximum(0,         np.floor(np.min(coord,axis=0))).astype(int)
+        maxs = np.minimum(shape_max, np.ceil (np.max(coord,axis=0))).astype(int)
+        return tuple(zip(tuple(mins),tuple(maxs)))
+
+
+
+# def repaint_labels(output, labels, polys, show_progress=True):
+#     """Repaint object instances in correct order based on probability scores.
+
+#     Does modify 'output' and 'polys' in-place, but will only write sparsely to 'output' where needed.
+
+#     output: numpy.ndarray or similar
+#         Label image (integer-valued)
+#     labels: iterable of int
+#         List of integer label ids that occur in output
+#     polys: dict
+#         Dictionary of polygon/polyhedra properties.
+#         Assumption is that the label id (-1) corresponds to the index in the polys dict
+
+#     """
+#     assert output.ndim in (2,3)
+
+#     if show_progress:
+#         labels = tqdm(labels, leave=True)
+
+#     labels_eliminated = set()
+
+#     # TODO: inelegant to have so much duplicated code here
+#     if output.ndim == 2:
+#         coord = lambda i: polys['coord'][i-1]
+#         prob  = lambda i: polys['prob'][i-1]
+
+#         for i in labels:
+#             if i in labels_eliminated: continue
+#             poly_i = Polygon(coord(i), shape_max=output.shape)
+
+#             # find all labels that overlap with i (including i)
+#             overlapping = set(np.unique(output[poly_i.slice][poly_i.mask])) - {0}
+#             assert i in overlapping
+#             # compute bbox union to find area to crop/replace in large output label image
+#             bbox_union = Polygon.coords_bbox(*[coord(j) for j in overlapping], shape_max=output.shape)
+
+#             # crop out label i, including the region that include all overlapping labels
+#             poly_i = Polygon(coord(i), bbox=bbox_union)
+#             mask = poly_i.mask.copy()
+
+#             # remove pixels from mask that belong to labels with higher probability
+#             for j in [j for j in overlapping if prob(j) > prob(i)]:
+#                 mask[ Polygon(coord(j), bbox=bbox_union).mask ] = False
+
+#             crop = output[poly_i.slice]
+#             crop[crop==i] = 0 # delete all remnants of i in crop
+#             crop[mask]    = i # paint i where mask still active
+
+#             labels_remaining = set(np.unique(output[poly_i.slice][poly_i.mask])) - {0}
+#             labels_eliminated.update(overlapping - labels_remaining)
+#     else:
+
+#         dist = lambda i: polys['dist'][i-1]
+#         origin = lambda i: polys['points'][i-1]
+#         prob = lambda i: polys['prob'][i-1]
+#         rays = polys['rays']
+
+#         for i in labels:
+#             if i in labels_eliminated: continue
+#             poly_i = Polyhedron(dist(i), origin(i), rays, shape_max=output.shape)
+
+#             # find all labels that overlap with i (including i)
+#             overlapping = set(np.unique(output[poly_i.slice][poly_i.mask])) - {0}
+#             assert i in overlapping
+#             # compute bbox union to find area to crop/replace in large output label image
+#             bbox_union = Polyhedron.coords_bbox(*[(dist(j),origin(j)) for j in overlapping], rays=rays, shape_max=output.shape)
+
+#             # crop out label i, including the region that include all overlapping labels
+#             poly_i = Polyhedron(dist(i), origin(i), rays, bbox=bbox_union)
+#             mask = poly_i.mask.copy()
+
+#             # remove pixels from mask that belong to labels with higher probability
+#             for j in [j for j in overlapping if prob(j) > prob(i)]:
+#                 mask[ Polyhedron(dist(j), origin(j), rays, bbox=bbox_union).mask ] = False
+
+#             crop = output[poly_i.slice]
+#             crop[crop==i] = 0 # delete all remnants of i in crop
+#             crop[mask]    = i # paint i where mask still active
+
+#             labels_remaining = set(np.unique(output[poly_i.slice][poly_i.mask])) - {0}
+#             labels_eliminated.update(overlapping - labels_remaining)
+
+#     if len(labels_eliminated) > 0:
+#         ind = [i-1 for i in labels_eliminated]
+#         for k,v in polys.items():
+#             if k in OBJECT_KEYS:
+#                 polys[k] = np.delete(v, ind, axis=0)
+
+
+
+############
+
+
+
+def predict_big(model, *args, **kwargs):
+    from .models import StarDist2D, StarDist3D
+    if isinstance(model,(StarDist2D,StarDist3D)):
+        dst = model.__class__.__name__
+    else:
+        dst = '{StarDist2D, StarDist3D}'
+    raise RuntimeError(f"This function has moved to {dst}.predict_instances_big.")
+
+
+
+class NotFullyVisible(Exception):
+    pass
+
+
+
+def _grid_divisible(grid, size, name=None, verbose=True):
+    if size % grid == 0:
+        return size
+    _size = size
+    size = math.ceil(size / grid) * grid
+    if bool(verbose):
+        print(f"{verbose if isinstance(verbose,str) else ''}increasing '{'value' if name is None else name}' from {_size} to {size} to be evenly divisible by {grid} (grid)", flush=True)
+    assert size % grid == 0
+    return size
+
+
+
+# def render_polygons(polys, shape):
+#     return polygons_to_label_coord(polys['coord'], shape=shape)

stardist_pkg/bioimageio_utils.py

@@ -0,0 +1,472 @@
+from pathlib import Path
+from pkg_resources import get_distribution
+from zipfile import ZipFile
+import numpy as np
+import tempfile
+from distutils.version import LooseVersion
+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 tensorflow import __version__ as tf_version
+    from . import __version__ as stardist_version
+    pkg_info = get_distribution("stardist")
+    # dependencies that start with the name "bioimageio" will be added as conda dependencies
+    reqs_conda = [str(req) for req in pkg_info.requires(extras=['bioimageio']) if str(req).startswith('bioimageio')]
+    # only stardist and tensorflow as pip dependencies
+    tf_major, tf_minor = LooseVersion(tf_version).version[:2]
+    reqs_pip = (f"stardist>={stardist_version}", f"tensorflow>={tf_major}.{tf_minor},<{tf_major+1}")
+    # conda environment
+    env = dict(
+        name = 'stardist',
+        channels = ['defaults', 'conda-forge'],
+        dependencies = [
+            ('python>=3.7,<3.8' if 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):
+    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"],
+        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=None,
+    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: 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)
+    """
+    _, 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)
+        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

stardist_pkg/geometry/__init__.py

@@ -0,0 +1,9 @@
+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 .geom2d import _dist_to_coord_old, _polygons_to_label_old
+
+#, dist_to_volume, dist_to_centroid

stardist_pkg/geometry/geom2d.py

@@ -0,0 +1,212 @@
+from __future__ import print_function, unicode_literals, absolute_import, division
+import numpy as np
+import warnings
+
+from skimage.measure import regionprops
+from skimage.draw import polygon
+from csbdeep.utils import _raise
+
+from ..utils import path_absolute, _is_power_of_2, _normalize_grid
+from ..matching import _check_label_array
+from stardist.lib.stardist2d import c_star_dist
+
+
+
+def _ocl_star_dist(lbl, n_rays=32, grid=(1,1)):
+    from gputools import OCLProgram, OCLArray, OCLImage
+    (np.isscalar(n_rays) and 0 < int(n_rays)) or _raise(ValueError())
+    n_rays = int(n_rays)
+    # slicing with grid is done with tuple(slice(0, None, g) for g in grid)
+    res_shape = tuple((s-1)//g+1 for s, g in zip(lbl.shape, grid))
+    
+    src = OCLImage.from_array(lbl.astype(np.uint16,copy=False))
+    dst = OCLArray.empty(res_shape+(n_rays,), dtype=np.float32)
+    program = OCLProgram(path_absolute("kernels/stardist2d.cl"), build_options=['-D', 'N_RAYS=%d' % n_rays])
+    program.run_kernel('star_dist', res_shape[::-1], None, dst.data, src, np.int32(grid[0]),np.int32(grid[1]))
+    return dst.get()
+
+
+def _cpp_star_dist(lbl, n_rays=32, grid=(1,1)):
+    (np.isscalar(n_rays) and 0 < int(n_rays)) or _raise(ValueError())
+    return c_star_dist(lbl.astype(np.uint16,copy=False), np.int32(n_rays), np.int32(grid[0]),np.int32(grid[1]))
+
+
+def _py_star_dist(a, n_rays=32, grid=(1,1)):
+    (np.isscalar(n_rays) and 0 < int(n_rays)) or _raise(ValueError())
+    if grid != (1,1):
+        raise NotImplementedError(grid)
+    
+    n_rays = int(n_rays)
+    a = a.astype(np.uint16,copy=False)
+    dst = np.empty(a.shape+(n_rays,),np.float32)
+
+    for i in range(a.shape[0]):
+        for j in range(a.shape[1]):
+            value = a[i,j]
+            if value == 0:
+                dst[i,j] = 0
+            else:
+                st_rays = np.float32((2*np.pi) / n_rays)
+                for k in range(n_rays):
+                    phi = np.float32(k*st_rays)
+                    dy = np.cos(phi)
+                    dx = np.sin(phi)
+                    x, y = np.float32(0), np.float32(0)
+                    while True:
+                        x += dx
+                        y += dy
+                        ii = int(round(i+x))
+                        jj = int(round(j+y))
+                        if (ii < 0 or ii >= a.shape[0] or
+                            jj < 0 or jj >= a.shape[1] or
+                            value != a[ii,jj]):
+                            # small correction as we overshoot the boundary
+                            t_corr = 1-.5/max(np.abs(dx),np.abs(dy))
+                            x -= t_corr*dx
+                            y -= t_corr*dy
+                            dist = np.sqrt(x**2+y**2)
+                            dst[i,j,k] = dist
+                            break
+    return dst
+
+
+def star_dist(a, n_rays=32, grid=(1,1), mode='cpp'):
+    """'a' assumbed to be a label image with integer values that encode object ids. id 0 denotes background."""
+
+    n_rays >= 3 or _raise(ValueError("need 'n_rays' >= 3"))
+
+    if mode == 'python':
+        return _py_star_dist(a, n_rays, grid=grid)
+    elif mode == 'cpp':
+        return _cpp_star_dist(a, n_rays, grid=grid)
+    elif mode == 'opencl':
+        return _ocl_star_dist(a, n_rays, grid=grid)
+    else:
+        _raise(ValueError("Unknown mode %s" % mode))
+
+
+def _dist_to_coord_old(rhos, grid=(1,1)):
+    """convert from polar to cartesian coordinates for a single image (3-D array) or multiple images (4-D array)"""
+
+    grid = _normalize_grid(grid,2)
+    is_single_image = rhos.ndim == 3
+    if is_single_image:
+        rhos = np.expand_dims(rhos,0)
+    assert rhos.ndim == 4
+
+    n_images,h,w,n_rays = rhos.shape
+    coord = np.empty((n_images,h,w,2,n_rays),dtype=rhos.dtype)
+
+    start = np.indices((h,w))
+    for i in range(2):
+        coord[...,i,:] = grid[i] * np.broadcast_to(start[i].reshape(1,h,w,1), (n_images,h,w,n_rays))
+
+    phis = ray_angles(n_rays).reshape(1,1,1,n_rays)
+
+    coord[...,0,:] += rhos * np.sin(phis) # row coordinate
+    coord[...,1,:] += rhos * np.cos(phis) # col coordinate
+
+    return coord[0] if is_single_image else coord
+
+
+def _polygons_to_label_old(coord, prob, points, shape=None, thr=-np.inf):
+    sh = coord.shape[:2] if shape is None else shape
+    lbl = np.zeros(sh,np.int32)
+    # sort points with increasing probability
+    ind = np.argsort([ prob[p[0],p[1]] for p in points ])
+    points = points[ind]
+
+    i = 1
+    for p in points:
+        if prob[p[0],p[1]] < thr:
+            continue
+        rr,cc = polygon(coord[p[0],p[1],0], coord[p[0],p[1],1], sh)
+        lbl[rr,cc] = i
+        i += 1
+
+    return lbl
+
+
+def dist_to_coord(dist, points, scale_dist=(1,1)):
+    """convert from polar to cartesian coordinates for a list of distances and center points
+    dist.shape   = (n_polys, n_rays)
+    points.shape = (n_polys, 2)
+    len(scale_dist) = 2
+    return coord.shape = (n_polys,2,n_rays)
+    """
+    dist = np.asarray(dist)
+    points = np.asarray(points)
+    assert dist.ndim==2 and points.ndim==2 and len(dist)==len(points) \
+        and points.shape[1]==2 and len(scale_dist)==2
+    n_rays = dist.shape[1]
+    phis = ray_angles(n_rays)
+    coord = (dist[:,np.newaxis]*np.array([np.sin(phis),np.cos(phis)])).astype(np.float32)
+    coord *= np.asarray(scale_dist).reshape(1,2,1)    
+    coord += points[...,np.newaxis] 
+    return coord
+
+
+def polygons_to_label_coord(coord, shape, labels=None):
+    """renders polygons to image of given shape
+
+    coord.shape   = (n_polys, n_rays)
+    """
+    coord = np.asarray(coord)
+    if labels is None: labels = np.arange(len(coord))
+
+    _check_label_array(labels, "labels")
+    assert coord.ndim==3 and coord.shape[1]==2 and len(coord)==len(labels)
+
+    lbl = np.zeros(shape,np.int32)
+
+    for i,c in zip(labels,coord):
+        rr,cc = polygon(*c, shape)
+        lbl[rr,cc] = i+1
+
+    return lbl
+
+
+def polygons_to_label(dist, points, shape, prob=None, thr=-np.inf, scale_dist=(1,1)):
+    """converts distances and center points to label image
+
+    dist.shape   = (n_polys, n_rays)
+    points.shape = (n_polys, 2)
+
+    label ids will be consecutive and adhere to the order given
+    """
+    dist = np.asarray(dist)
+    points = np.asarray(points)
+    prob = np.inf*np.ones(len(points)) if prob is None else np.asarray(prob)
+
+    assert dist.ndim==2 and points.ndim==2 and len(dist)==len(points)
+    assert len(points)==len(prob) and points.shape[1]==2 and prob.ndim==1
+
+    n_rays = dist.shape[1]
+
+    ind = prob>thr
+    points = points[ind]
+    dist = dist[ind]
+    prob = prob[ind]
+
+    ind = np.argsort(prob, kind='stable')
+    points = points[ind]
+    dist = dist[ind]
+
+    coord = dist_to_coord(dist, points, scale_dist=scale_dist)
+
+    return polygons_to_label_coord(coord, shape=shape, labels=ind)
+
+
+def relabel_image_stardist(lbl, n_rays, **kwargs):
+    """relabel each label region in `lbl` with its star representation"""
+    _check_label_array(lbl, "lbl")
+    if not lbl.ndim==2:
+        raise ValueError("lbl image should be 2 dimensional")
+    dist = star_dist(lbl, n_rays, **kwargs)
+    points = np.array(tuple(np.array(r.centroid).astype(int) for r in regionprops(lbl)))
+    dist = dist[tuple(points.T)]
+    return polygons_to_label(dist, points, shape=lbl.shape)
+
+
+def ray_angles(n_rays=32):
+    return np.linspace(0,2*np.pi,n_rays,endpoint=False)

stardist_pkg/kernels/stardist2d.cl

@@ -0,0 +1,51 @@
+#ifndef M_PI
+#define M_PI 3.141592653589793
+#endif
+
+__constant sampler_t sampler = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;
+
+inline float2 pol2cart(const float rho, const float phi) {
+    const float x = rho * cos(phi);
+    const float y = rho * sin(phi);
+    return (float2)(x,y);
+}
+
+__kernel void star_dist(__global float* dst, read_only image2d_t src, const int grid_y, const int grid_x) {
+
+    const int i = get_global_id(0), j = get_global_id(1);
+    const int Nx = get_global_size(0), Ny = get_global_size(1);
+    const float2 grid = (float2)(grid_x, grid_y);
+
+    const float2 origin = (float2)(i,j) * grid;
+    const int value = read_imageui(src,sampler,origin).x;
+
+    if (value == 0) {
+        // background pixel -> nothing to do, write all zeros
+        for (int k = 0; k < N_RAYS; k++) {
+            dst[k + i*N_RAYS + j*N_RAYS*Nx] = 0;
+        }
+    } else {
+        float st_rays = (2*M_PI) / N_RAYS; // step size for ray angles
+        // for all rays
+        for (int k = 0; k < N_RAYS; k++) {
+            const float phi = k*st_rays; // current ray angle phi
+            const float2 dir = pol2cart(1,phi); // small vector in direction of ray
+            float2 offset = 0; // offset vector to be added to origin
+            // find radius that leaves current object
+            while (1) {
+                offset += dir;
+                const int offset_value = read_imageui(src,sampler,round(origin+offset)).x;
+                if (offset_value != value) {
+                  // small correction as we overshoot the boundary
+                  const float t_corr = .5f/fmax(fabs(dir.x),fabs(dir.y));
+                  offset += (t_corr-1.f)*dir;
+
+                  const float dist = sqrt(offset.x*offset.x + offset.y*offset.y);
+                  dst[k + i*N_RAYS + j*N_RAYS*Nx] = dist;
+                  break;
+                }
+            }
+        }
+    }
+
+}

stardist_pkg/kernels/stardist3d.cl

@@ -0,0 +1,63 @@
+#ifndef M_PI
+#define M_PI 3.141592653589793
+#endif
+
+__constant sampler_t sampler = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;
+
+inline int round_to_int(float r) {
+    return (int)rint(r);
+}
+
+
+__kernel void stardist3d(read_only image3d_t lbl, __constant float * rays, __global float* dist, const int grid_z, const int grid_y, const int grid_x) {
+
+  const int i = get_global_id(0);
+  const int j = get_global_id(1);
+  const int k = get_global_id(2);
+
+  const int Nx = get_global_size(0);
+  const int Ny = get_global_size(1);
+  const int Nz = get_global_size(2);
+
+  const float4 grid = (float4)(grid_x, grid_y, grid_z, 1);
+  const float4 origin = (float4)(i,j,k,0) * grid;
+  const int value = read_imageui(lbl,sampler,origin).x;
+
+  if (value == 0) {
+	// background pixel -> nothing to do, write all zeros
+	for (int m = 0; m < N_RAYS; m++) {
+	  dist[m + i*N_RAYS + j*N_RAYS*Nx+k*N_RAYS*Nx*Ny] = 0;
+	}
+	
+  }
+  else {
+	for (int m = 0; m < N_RAYS; m++) {
+	
+	  const float4 dx = (float4)(rays[3*m+2],rays[3*m+1],rays[3*m],0);
+	  // if ((i==Nx/2)&&(j==Ny/2)&(k==Nz/2)){
+	  // 	printf("kernel: %.2f %.2f  %.2f  \n",dx.x,dx.y,dx.z);
+	  // }
+	  float4 x = (float4)(0,0,0,0);
+
+	  // move along ray
+	  while (1) {
+		x += dx;
+		// if ((i==10)&&(j==10)&(k==10)){
+		//   printf("kernel run: %.2f %.2f  %.2f value %d \n",x.x,x.y,x.z, read_imageui(lbl,sampler,origin+x).x);
+		// }
+		
+		// to make it equivalent to the cpp version... 
+		const float4 x_int = (float4)(round_to_int(x.x),
+										 round_to_int(x.y),
+										 round_to_int(x.z),
+										 0);
+
+		if (value != read_imageui(lbl,sampler,origin+x_int).x){
+
+		  dist[m + i*N_RAYS + j*N_RAYS*Nx+k*N_RAYS*Nx*Ny] = length(x_int);
+		  break;		  
+		}
+	  }
+	}
+  }
+}

stardist_pkg/matching.py

@@ -0,0 +1,483 @@
+import numpy as np
+
+from numba import jit
+from tqdm import tqdm
+from scipy.optimize import linear_sum_assignment
+from skimage.measure import regionprops
+from collections import namedtuple
+from csbdeep.utils import _raise
+
+matching_criteria = dict()
+
+
+def label_are_sequential(y):
+    """ returns true if y has only sequential labels from 1... """
+    labels = np.unique(y)
+    return (set(labels)-{0}) == set(range(1,1+labels.max()))
+
+
+def is_array_of_integers(y):
+    return isinstance(y,np.ndarray) and np.issubdtype(y.dtype, np.integer)
+
+
+def _check_label_array(y, name=None, check_sequential=False):
+    err = ValueError("{label} must be an array of {integers}.".format(
+        label = 'labels' if name is None else name,
+        integers = ('sequential ' if check_sequential else '') + 'non-negative integers',
+    ))
+    is_array_of_integers(y) or _raise(err)
+    if len(y) == 0:
+        return True
+    if check_sequential:
+        label_are_sequential(y) or _raise(err)
+    else:
+        y.min() >= 0 or _raise(err)
+    return True
+
+
+def label_overlap(x, y, check=True):
+    if check:
+        _check_label_array(x,'x',True)
+        _check_label_array(y,'y',True)
+        x.shape == y.shape or _raise(ValueError("x and y must have the same shape"))
+    return _label_overlap(x, y)
+
+@jit(nopython=True)
+def _label_overlap(x, y):
+    x = x.ravel()
+    y = y.ravel()
+    overlap = np.zeros((1+x.max(),1+y.max()), dtype=np.uint)
+    for i in range(len(x)):
+        overlap[x[i],y[i]] += 1
+    return overlap
+
+
+def _safe_divide(x,y, eps=1e-10):
+    """computes a safe divide which returns 0 if y is zero"""
+    if np.isscalar(x) and np.isscalar(y):
+        return x/y if np.abs(y)>eps else 0.0
+    else:
+        out = np.zeros(np.broadcast(x,y).shape, np.float32)
+        np.divide(x,y, out=out, where=np.abs(y)>eps)
+        return out
+
+
+def intersection_over_union(overlap):
+    _check_label_array(overlap,'overlap')
+    if np.sum(overlap) == 0:
+        return overlap
+    n_pixels_pred = np.sum(overlap, axis=0, keepdims=True)
+    n_pixels_true = np.sum(overlap, axis=1, keepdims=True)
+    return _safe_divide(overlap, (n_pixels_pred + n_pixels_true - overlap))
+
+matching_criteria['iou'] = intersection_over_union
+
+
+def intersection_over_true(overlap):
+    _check_label_array(overlap,'overlap')
+    if np.sum(overlap) == 0:
+        return overlap
+    n_pixels_true = np.sum(overlap, axis=1, keepdims=True)
+    return _safe_divide(overlap, n_pixels_true)
+
+matching_criteria['iot'] = intersection_over_true
+
+
+def intersection_over_pred(overlap):
+    _check_label_array(overlap,'overlap')
+    if np.sum(overlap) == 0:
+        return overlap
+    n_pixels_pred = np.sum(overlap, axis=0, keepdims=True)
+    return _safe_divide(overlap, n_pixels_pred)
+
+matching_criteria['iop'] = intersection_over_pred
+
+
+def precision(tp,fp,fn):
+    return tp/(tp+fp) if tp > 0 else 0
+def recall(tp,fp,fn):
+    return tp/(tp+fn) if tp > 0 else 0
+def accuracy(tp,fp,fn):
+    # also known as "average precision" (?)
+    # -> https://www.kaggle.com/c/data-science-bowl-2018#evaluation
+    return tp/(tp+fp+fn) if tp > 0 else 0
+def f1(tp,fp,fn):
+    # also known as "dice coefficient"
+    return (2*tp)/(2*tp+fp+fn) if tp > 0 else 0
+
+
+def matching(y_true, y_pred, thresh=0.5, criterion='iou', report_matches=False):
+    """Calculate detection/instance segmentation metrics between ground truth and predicted label images.
+
+    Currently, the following metrics are implemented:
+
+    'fp', 'tp', 'fn', 'precision', 'recall', 'accuracy', 'f1', 'criterion', 'thresh', 'n_true', 'n_pred', 'mean_true_score', 'mean_matched_score', 'panoptic_quality'
+
+    Corresponding objects of y_true and y_pred are counted as true positives (tp), false positives (fp), and false negatives (fn)
+    whether their intersection over union (IoU) >= thresh (for criterion='iou', which can be changed)
+
+    * mean_matched_score is the mean IoUs of matched true positives
+
+    * mean_true_score is the mean IoUs of matched true positives but normalized by the total number of GT objects
+
+    * panoptic_quality defined as in Eq. 1 of Kirillov et al. "Panoptic Segmentation", CVPR 2019
+
+    Parameters
+    ----------
+    y_true: ndarray
+        ground truth label image (integer valued)
+    y_pred: ndarray
+        predicted label image (integer valued)
+    thresh: float
+        threshold for matching criterion (default 0.5)
+    criterion: string
+        matching criterion (default IoU)
+    report_matches: bool
+        if True, additionally calculate matched_pairs and matched_scores (note, that this returns even gt-pred pairs whose scores are below  'thresh')
+
+    Returns
+    -------
+    Matching object with different metrics as attributes
+
+    Examples
+    --------
+    >>> y_true = np.zeros((100,100), np.uint16)
+    >>> y_true[10:20,10:20] = 1
+    >>> y_pred = np.roll(y_true,5,axis = 0)
+
+    >>> stats = matching(y_true, y_pred)
+    >>> print(stats)
+    Matching(criterion='iou', thresh=0.5, fp=1, tp=0, fn=1, precision=0, recall=0, accuracy=0, f1=0, n_true=1, n_pred=1, mean_true_score=0.0, mean_matched_score=0.0, panoptic_quality=0.0)
+
+    """
+    _check_label_array(y_true,'y_true')
+    _check_label_array(y_pred,'y_pred')
+    y_true.shape == y_pred.shape or _raise(ValueError("y_true ({y_true.shape}) and y_pred ({y_pred.shape}) have different shapes".format(y_true=y_true, y_pred=y_pred)))
+    criterion in matching_criteria or _raise(ValueError("Matching criterion '%s' not supported." % criterion))
+    if thresh is None: thresh = 0
+    thresh = float(thresh) if np.isscalar(thresh) else map(float,thresh)
+
+    y_true, _, map_rev_true = relabel_sequential(y_true)
+    y_pred, _, map_rev_pred = relabel_sequential(y_pred)
+
+    overlap = label_overlap(y_true, y_pred, check=False)
+    scores = matching_criteria[criterion](overlap)
+    assert 0 <= np.min(scores) <= np.max(scores) <= 1
+
+    # ignoring background
+    scores = scores[1:,1:]
+    n_true, n_pred = scores.shape
+    n_matched = min(n_true, n_pred)
+
+    def _single(thr):
+        # not_trivial = n_matched > 0 and np.any(scores >= thr)
+        not_trivial = n_matched > 0
+        if not_trivial:
+            # compute optimal matching with scores as tie-breaker
+            costs = -(scores >= thr).astype(float) - scores / (2*n_matched)
+            true_ind, pred_ind = linear_sum_assignment(costs)
+            assert n_matched == len(true_ind) == len(pred_ind)
+            match_ok = scores[true_ind,pred_ind] >= thr
+            tp = np.count_nonzero(match_ok)
+        else:
+            tp = 0
+        fp = n_pred - tp
+        fn = n_true - tp
+        # assert tp+fp == n_pred
+        # assert tp+fn == n_true
+
+        # the score sum over all matched objects (tp)
+        sum_matched_score = np.sum(scores[true_ind,pred_ind][match_ok]) if not_trivial else 0.0
+
+        # the score average over all matched objects (tp)
+        mean_matched_score = _safe_divide(sum_matched_score, tp)
+        # the score average over all gt/true objects
+        mean_true_score    = _safe_divide(sum_matched_score, n_true)
+        panoptic_quality   = _safe_divide(sum_matched_score, tp+fp/2+fn/2)
+
+        stats_dict = dict (
+            criterion          = criterion,
+            thresh             = thr,
+            fp                 = fp,
+            tp                 = tp,
+            fn                 = fn,
+            precision          = precision(tp,fp,fn),
+            recall             = recall(tp,fp,fn),
+            accuracy           = accuracy(tp,fp,fn),
+            f1                 = f1(tp,fp,fn),
+            n_true             = n_true,
+            n_pred             = n_pred,
+            mean_true_score    = mean_true_score,
+            mean_matched_score = mean_matched_score,
+            panoptic_quality   = panoptic_quality,
+        )
+        if bool(report_matches):
+            if not_trivial:
+                stats_dict.update (
+                    # int() to be json serializable
+                    matched_pairs  = tuple((int(map_rev_true[i]),int(map_rev_pred[j])) for i,j in zip(1+true_ind,1+pred_ind)),
+                    matched_scores = tuple(scores[true_ind,pred_ind]),
+                    matched_tps    = tuple(map(int,np.flatnonzero(match_ok))),
+                )
+            else:
+                stats_dict.update (
+                    matched_pairs  = (),
+                    matched_scores = (),
+                    matched_tps    = (),
+                )
+        return namedtuple('Matching',stats_dict.keys())(*stats_dict.values())
+
+    return _single(thresh) if np.isscalar(thresh) else tuple(map(_single,thresh))
+
+
+
+def matching_dataset(y_true, y_pred, thresh=0.5, criterion='iou', by_image=False, show_progress=True, parallel=False):
+    """matching metrics for list of images, see `stardist.matching.matching`
+    """
+    len(y_true) == len(y_pred) or _raise(ValueError("y_true and y_pred must have the same length."))
+    return matching_dataset_lazy (
+        tuple(zip(y_true,y_pred)), thresh=thresh, criterion=criterion, by_image=by_image, show_progress=show_progress, parallel=parallel,
+    )
+
+
+
+def matching_dataset_lazy(y_gen, thresh=0.5, criterion='iou', by_image=False, show_progress=True, parallel=False):
+
+    expected_keys = set(('fp', 'tp', 'fn', 'precision', 'recall', 'accuracy', 'f1', 'criterion', 'thresh', 'n_true', 'n_pred', 'mean_true_score', 'mean_matched_score', 'panoptic_quality'))
+
+    single_thresh = False
+    if np.isscalar(thresh):
+        single_thresh = True
+        thresh = (thresh,)
+
+    tqdm_kwargs = {}
+    tqdm_kwargs['disable'] = not bool(show_progress)
+    if int(show_progress) > 1:
+        tqdm_kwargs['total'] = int(show_progress)
+
+    # compute matching stats for every pair of label images
+    if parallel:
+        from concurrent.futures import ThreadPoolExecutor
+        fn = lambda pair: matching(*pair, thresh=thresh, criterion=criterion, report_matches=False)
+        with ThreadPoolExecutor() as pool:
+            stats_all = tuple(pool.map(fn, tqdm(y_gen,**tqdm_kwargs)))
+    else:
+        stats_all = tuple (
+            matching(y_t, y_p, thresh=thresh, criterion=criterion, report_matches=False)
+            for y_t,y_p in tqdm(y_gen,**tqdm_kwargs)
+        )
+
+    # accumulate results over all images for each threshold separately
+    n_images, n_threshs = len(stats_all), len(thresh)
+    accumulate = [{} for _ in range(n_threshs)]
+    for stats in stats_all:
+        for i,s in enumerate(stats):
+            acc = accumulate[i]
+            for k,v in s._asdict().items():
+                if k == 'mean_true_score' and not bool(by_image):
+                    # convert mean_true_score to "sum_matched_score"
+                    acc[k] = acc.setdefault(k,0) + v * s.n_true
+                else:
+                    try:
+                        acc[k] = acc.setdefault(k,0) + v
+                    except TypeError:
+                        pass
+
+    # normalize/compute 'precision', 'recall', 'accuracy', 'f1'
+    for thr,acc in zip(thresh,accumulate):
+        set(acc.keys()) == expected_keys or _raise(ValueError("unexpected keys"))
+        acc['criterion'] = criterion
+        acc['thresh'] = thr
+        acc['by_image'] = bool(by_image)
+        if bool(by_image):
+            for k in ('precision', 'recall', 'accuracy', 'f1', 'mean_true_score', 'mean_matched_score', 'panoptic_quality'):
+                acc[k] /= n_images
+        else:
+            tp, fp, fn, n_true = acc['tp'], acc['fp'], acc['fn'], acc['n_true']
+            sum_matched_score = acc['mean_true_score']
+
+            mean_matched_score = _safe_divide(sum_matched_score, tp)
+            mean_true_score    = _safe_divide(sum_matched_score, n_true)
+            panoptic_quality   = _safe_divide(sum_matched_score, tp+fp/2+fn/2)
+
+            acc.update(
+                precision          = precision(tp,fp,fn),
+                recall             = recall(tp,fp,fn),
+                accuracy           = accuracy(tp,fp,fn),
+                f1                 = f1(tp,fp,fn),
+                mean_true_score    = mean_true_score,
+                mean_matched_score = mean_matched_score,
+                panoptic_quality   = panoptic_quality,
+            )
+
+    accumulate = tuple(namedtuple('DatasetMatching',acc.keys())(*acc.values()) for acc in accumulate)
+    return accumulate[0] if single_thresh else accumulate
+
+
+
+# copied from scikit-image master for now (remove when part of a release)
+def relabel_sequential(label_field, offset=1):
+    """Relabel arbitrary labels to {`offset`, ... `offset` + number_of_labels}.
+
+    This function also returns the forward map (mapping the original labels to
+    the reduced labels) and the inverse map (mapping the reduced labels back
+    to the original ones).
+
+    Parameters
+    ----------
+    label_field : numpy array of int, arbitrary shape
+        An array of labels, which must be non-negative integers.
+    offset : int, optional
+        The return labels will start at `offset`, which should be
+        strictly positive.
+
+    Returns
+    -------
+    relabeled : numpy array of int, same shape as `label_field`
+        The input label field with labels mapped to
+        {offset, ..., number_of_labels + offset - 1}.
+        The data type will be the same as `label_field`, except when
+        offset + number_of_labels causes overflow of the current data type.
+    forward_map : numpy array of int, shape ``(label_field.max() + 1,)``
+        The map from the original label space to the returned label
+        space. Can be used to re-apply the same mapping. See examples
+        for usage. The data type will be the same as `relabeled`.
+    inverse_map : 1D numpy array of int, of length offset + number of labels
+        The map from the new label space to the original space. This
+        can be used to reconstruct the original label field from the
+        relabeled one. The data type will be the same as `relabeled`.
+
+    Notes
+    -----
+    The label 0 is assumed to denote the background and is never remapped.
+
+    The forward map can be extremely big for some inputs, since its
+    length is given by the maximum of the label field. However, in most
+    situations, ``label_field.max()`` is much smaller than
+    ``label_field.size``, and in these cases the forward map is
+    guaranteed to be smaller than either the input or output images.
+
+    Examples
+    --------
+    >>> from skimage.segmentation import relabel_sequential
+    >>> label_field = np.array([1, 1, 5, 5, 8, 99, 42])
+    >>> relab, fw, inv = relabel_sequential(label_field)
+    >>> relab
+    array([1, 1, 2, 2, 3, 5, 4])
+    >>> fw
+    array([0, 1, 0, 0, 0, 2, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 0,
+           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 5])
+    >>> inv
+    array([ 0,  1,  5,  8, 42, 99])
+    >>> (fw[label_field] == relab).all()
+    True
+    >>> (inv[relab] == label_field).all()
+    True
+    >>> relab, fw, inv = relabel_sequential(label_field, offset=5)
+    >>> relab
+    array([5, 5, 6, 6, 7, 9, 8])
+    """
+    offset = int(offset)
+    if offset <= 0:
+        raise ValueError("Offset must be strictly positive.")
+    if np.min(label_field) < 0:
+        raise ValueError("Cannot relabel array that contains negative values.")
+    max_label = int(label_field.max()) # Ensure max_label is an integer
+    if not np.issubdtype(label_field.dtype, np.integer):
+        new_type = np.min_scalar_type(max_label)
+        label_field = label_field.astype(new_type)
+    labels = np.unique(label_field)
+    labels0 = labels[labels != 0]
+    new_max_label = offset - 1 + len(labels0)
+    new_labels0 = np.arange(offset, new_max_label + 1)
+    output_type = label_field.dtype
+    required_type = np.min_scalar_type(new_max_label)
+    if np.dtype(required_type).itemsize > np.dtype(label_field.dtype).itemsize:
+        output_type = required_type
+    forward_map = np.zeros(max_label + 1, dtype=output_type)
+    forward_map[labels0] = new_labels0
+    inverse_map = np.zeros(new_max_label + 1, dtype=output_type)
+    inverse_map[offset:] = labels0
+    relabeled = forward_map[label_field]
+    return relabeled, forward_map, inverse_map
+
+
+
+def group_matching_labels(ys, thresh=1e-10, criterion='iou'):
+    """
+    Group matching objects (i.e. assign the same label id) in a
+    list of label images (e.g. consecutive frames of a time-lapse).
+
+    Uses function `matching` (with provided `criterion` and `thresh`) to
+    iteratively/greedily match and group objects/labels in consecutive images of `ys`.
+    To that end, matching objects are grouped together by assigning the same label id,
+    whereas unmatched objects are assigned a new label id.
+    At the end of this process, each label group will have been assigned a unique id.
+
+    Note that the label images `ys` will not be modified. Instead, they will initially
+    be duplicated and converted to data type `np.int32` before objects are grouped and the result
+    is returned. (Note that `np.int32` limits the number of label groups to at most 2147483647.)
+
+    Example
+    -------
+    import numpy as np
+    from stardist.data import test_image_nuclei_2d
+    from stardist.matching import group_matching_labels
+
+    _y = test_image_nuclei_2d(return_mask=True)[1]
+    labels = np.stack([_y, 2*np.roll(_y,10)], axis=0)
+
+    labels_new = group_matching_labels(labels)
+
+    Parameters
+    ----------
+    ys : np.ndarray or list/tuple of np.ndarray
+        list/array of integer labels (2D or 3D)
+    
+    """
+    # check 'ys' without making a copy
+    len(ys) > 1 or _raise(ValueError("'ys' must have 2 or more entries"))
+    if isinstance(ys, np.ndarray):
+        _check_label_array(ys, 'ys')
+        ys.ndim > 1 or _raise(ValueError("'ys' must be at least 2-dimensional"))
+        ys_grouped = np.empty_like(ys, dtype=np.int32)
+    else:
+        all(_check_label_array(y, 'ys') for y in ys) or _raise(ValueError("'ys' must be a list of label images"))
+        all(y.shape==ys[0].shape for y in ys) or _raise(ValueError("all label images must have the same shape"))
+        ys_grouped = np.empty((len(ys),)+ys[0].shape, dtype=np.int32)
+
+    def _match_single(y_prev, y, next_id):
+        y = y.astype(np.int32, copy=False)
+        res = matching(y_prev, y, report_matches=True, thresh=thresh, criterion=criterion)
+        # relabel dict (for matching labels) that maps label ids from y -> y_prev 
+        relabel = dict(reversed(res.matched_pairs[i]) for i in res.matched_tps)
+        y_grouped = np.zeros_like(y)
+        for r in regionprops(y):
+            m = (y[r.slice] == r.label)
+            if r.label in relabel:
+                y_grouped[r.slice][m] = relabel[r.label]
+            else:
+                y_grouped[r.slice][m] = next_id
+                next_id += 1
+        return y_grouped, next_id
+
+    ys_grouped[0] = ys[0]
+    next_id = ys_grouped[0].max() + 1
+    for i in range(len(ys)-1):
+        ys_grouped[i+1], next_id = _match_single(ys_grouped[i], ys[i+1], next_id)
+    return ys_grouped
+
+
+
+def _shuffle_labels(y):
+    _check_label_array(y, 'y')
+    y2 = np.zeros_like(y)
+    ids = tuple(set(np.unique(y)) - {0})
+    relabel = dict(zip(ids,np.random.permutation(ids)))
+    for r in regionprops(y):
+        m = (y[r.slice] == r.label)
+        y2[r.slice][m] = relabel[r.label]
+    return y2

stardist_pkg/models/__init__.py

@@ -0,0 +1,27 @@
+from __future__ import absolute_import, print_function
+
+from .model2d import Config2D, StarDist2D, StarDistData2D
+
+from csbdeep.utils import backend_channels_last
+from csbdeep.utils.tf import keras_import
+K = keras_import('backend')
+if not backend_channels_last():
+    raise NotImplementedError(
+        "Keras is configured to use the '%s' image data format, which is currently not supported. "
+        "Please change it to use 'channels_last' instead: "
+        "https://keras.io/getting-started/faq/#where-is-the-keras-configuration-file-stored" % K.image_data_format()
+    )
+del backend_channels_last, K
+
+from csbdeep.models import register_model, register_aliases, clear_models_and_aliases
+# register pre-trained models and aliases (TODO: replace with updatable solution)
+clear_models_and_aliases(StarDist2D, StarDist3D)
+register_model(StarDist2D,   '2D_versatile_fluo', 'https://github.com/stardist/stardist-models/releases/download/v0.1/python_2D_versatile_fluo.zip', '8db40dacb5a1311b8d2c447ad934fb8a')
+register_model(StarDist2D,   '2D_versatile_he',   'https://github.com/stardist/stardist-models/releases/download/v0.1/python_2D_versatile_he.zip', 'bf34cb3c0e5b3435971e18d66778a4ec')
+register_model(StarDist2D,   '2D_paper_dsb2018',  'https://github.com/stardist/stardist-models/releases/download/v0.1/python_2D_paper_dsb2018.zip', '6287bf283f85c058ec3e7094b41039b5')
+register_model(StarDist2D,   '2D_demo',           'https://github.com/stardist/stardist-models/releases/download/v0.1/python_2D_demo.zip', '31f70402f58c50dd231ec31b4375ea2c')
+
+register_aliases(StarDist2D, '2D_paper_dsb2018',  'DSB 2018 (from StarDist 2D paper)')
+register_aliases(StarDist2D, '2D_versatile_fluo', 'Versatile (fluorescent nuclei)')
+register_aliases(StarDist2D, '2D_versatile_he',   'Versatile (H&E nuclei)')
+del register_model, register_aliases, clear_models_and_aliases

stardist_pkg/models/base.py

@@ -0,0 +1,1196 @@
+from __future__ import print_function, unicode_literals, absolute_import, division
+
+import numpy as np
+import sys
+import warnings
+import math
+from tqdm import tqdm
+from collections import namedtuple
+from pathlib import Path
+import threading
+import functools
+import scipy.ndimage as ndi
+import numbers
+
+from csbdeep.models.base_model import BaseModel
+from csbdeep.utils.tf import export_SavedModel, keras_import, IS_TF_1, CARETensorBoard
+
+import tensorflow as tf
+K = keras_import('backend')
+Sequence = keras_import('utils', 'Sequence')
+Adam = keras_import('optimizers', 'Adam')
+ReduceLROnPlateau, TensorBoard = keras_import('callbacks', 'ReduceLROnPlateau', 'TensorBoard')
+
+from csbdeep.utils import _raise, backend_channels_last, axes_check_and_normalize, axes_dict, load_json, save_json
+from csbdeep.internals.predict import tile_iterator, total_n_tiles
+from csbdeep.internals.train import RollingSequence
+from csbdeep.data import Resizer
+
+from ..sample_patches import get_valid_inds
+from ..nms import _ind_prob_thresh
+from ..utils import _is_power_of_2,  _is_floatarray, optimize_threshold
+
+# TODO: helper function to check if receptive field of cnn is sufficient for object sizes in GT
+
+def generic_masked_loss(mask, loss, weights=1, norm_by_mask=True, reg_weight=0, reg_penalty=K.abs):
+    def _loss(y_true, y_pred):
+        actual_loss = K.mean(mask * weights * loss(y_true, y_pred), axis=-1)
+        norm_mask = (K.mean(mask) + K.epsilon()) if norm_by_mask else 1
+        if reg_weight > 0:
+            reg_loss = K.mean((1-mask) * reg_penalty(y_pred), axis=-1)
+            return actual_loss / norm_mask + reg_weight * reg_loss
+        else:
+            return actual_loss / norm_mask
+    return _loss
+
+def masked_loss(mask, penalty, reg_weight, norm_by_mask):
+    loss = lambda y_true, y_pred: penalty(y_true - y_pred)
+    return generic_masked_loss(mask, loss, reg_weight=reg_weight, norm_by_mask=norm_by_mask)
+
+# TODO: should we use norm_by_mask=True in the loss or only in a metric?
+#       previous 2D behavior was norm_by_mask=False
+#       same question for reg_weight? use 1e-4 (as in 3D) or 0 (as in 2D)?
+
+def masked_loss_mae(mask, reg_weight=0, norm_by_mask=True):
+    return masked_loss(mask, K.abs, reg_weight=reg_weight, norm_by_mask=norm_by_mask)
+
+def masked_loss_mse(mask, reg_weight=0, norm_by_mask=True):
+    return masked_loss(mask, K.square, reg_weight=reg_weight, norm_by_mask=norm_by_mask)
+
+def masked_metric_mae(mask):
+    def relevant_mae(y_true, y_pred):
+        return masked_loss(mask, K.abs, reg_weight=0, norm_by_mask=True)(y_true, y_pred)
+    return relevant_mae
+
+def masked_metric_mse(mask):
+    def relevant_mse(y_true, y_pred):
+        return masked_loss(mask, K.square, reg_weight=0, norm_by_mask=True)(y_true, y_pred)
+    return relevant_mse
+
+def kld(y_true, y_pred):
+    y_true = K.clip(y_true, K.epsilon(), 1)
+    y_pred = K.clip(y_pred, K.epsilon(), 1)
+    return K.mean(K.binary_crossentropy(y_true, y_pred) - K.binary_crossentropy(y_true, y_true), axis=-1)
+
+
+def masked_loss_iou(mask, reg_weight=0, norm_by_mask=True):
+    def iou_loss(y_true, y_pred):
+        axis = -1 if backend_channels_last() else 1
+        # y_pred can be negative (since not constrained) -> 'inter' can be very large for y_pred << 0
+        # - clipping y_pred values at 0 can lead to vanishing gradients
+        # - 'K.sign(y_pred)' term fixes issue by enforcing that y_pred values >= 0 always lead to larger 'inter' (lower loss)
+        inter = K.mean(K.sign(y_pred)*K.square(K.minimum(y_true,y_pred)), axis=axis)
+        union = K.mean(K.square(K.maximum(y_true,y_pred)), axis=axis)
+        iou = inter/(union+K.epsilon())
+        iou = K.expand_dims(iou,axis)
+        loss = 1. - iou # + 0.005*K.abs(y_true-y_pred)
+        return loss
+    return generic_masked_loss(mask, iou_loss, reg_weight=reg_weight, norm_by_mask=norm_by_mask)
+
+def masked_metric_iou(mask, reg_weight=0, norm_by_mask=True):
+    def iou_metric(y_true, y_pred):
+        axis = -1 if backend_channels_last() else 1
+        y_pred = K.maximum(0., y_pred)
+        inter = K.mean(K.square(K.minimum(y_true,y_pred)), axis=axis)
+        union = K.mean(K.square(K.maximum(y_true,y_pred)), axis=axis)
+        iou = inter/(union+K.epsilon())
+        loss = K.expand_dims(iou,axis)
+        return loss
+    return generic_masked_loss(mask, iou_metric, reg_weight=reg_weight, norm_by_mask=norm_by_mask)
+
+
+def weighted_categorical_crossentropy(weights, ndim):
+    """ ndim = (2,3) """
+
+    axis = -1 if backend_channels_last() else 1
+    shape = [1]*(ndim+2)
+    shape[axis] = len(weights)
+    weights = np.broadcast_to(weights, shape)
+    weights = K.constant(weights)
+
+    def weighted_cce(y_true, y_pred):
+        # ignore pixels that have y_true (prob_class) < 0
+        mask = K.cast(y_true>=0, K.floatx())
+        y_pred /= K.sum(y_pred+K.epsilon(), axis=axis, keepdims=True)
+        y_pred = K.clip(y_pred, K.epsilon(), 1. - K.epsilon())
+        loss = - K.sum(weights*mask*y_true*K.log(y_pred), axis = axis)
+        return loss
+
+    return weighted_cce
+
+
+class StarDistDataBase(RollingSequence):
+
+    def __init__(self, X, Y, n_rays, grid, batch_size, patch_size, length,
+                 n_classes=None, classes=None,
+                 use_gpu=False, sample_ind_cache=True, maxfilter_patch_size=None, augmenter=None, foreground_prob=0):
+
+        super().__init__(data_size=len(X), batch_size=batch_size, length=length, shuffle=True)
+
+        if isinstance(X, (np.ndarray, tuple, list)):
+            X = [x.astype(np.float32, copy=False) for x in X]
+
+        # sanity checks
+        len(X)==len(Y) and len(X)>0 or _raise(ValueError("X and Y can't be empty and must have same length"))
+
+        if classes is None:
+            # set classes to None for all images (i.e. defaults to every object instance assigned the same class)
+            classes = (None,)*len(X)
+        else:
+            n_classes is not None or warnings.warn("Ignoring classes since n_classes is None")
+
+        len(classes)==len(X) or _raise(ValueError("X and classes must have same length"))
+
+        self.n_classes, self.classes = n_classes, classes
+
+        nD = len(patch_size)
+        assert nD in (2,3)
+        x_ndim = X[0].ndim
+        assert x_ndim in (nD,nD+1)
+
+        if isinstance(X, (np.ndarray, tuple, list)) and \
+           isinstance(Y, (np.ndarray, tuple, list)):
+            all(y.ndim==nD and x.ndim==x_ndim and x.shape[:nD]==y.shape for x,y in zip(X,Y)) or _raise(ValueError("images and masks should have corresponding shapes/dimensions"))
+            all(x.shape[:nD]>=tuple(patch_size) for x in X) or _raise(ValueError("Some images are too small for given patch_size {patch_size}".format(patch_size=patch_size)))
+
+        if x_ndim == nD:
+            self.n_channel = None
+        else:
+            self.n_channel = X[0].shape[-1]
+            if isinstance(X, (np.ndarray, tuple, list)):
+                assert all(x.shape[-1]==self.n_channel for x in X)
+
+        assert 0 <= foreground_prob <= 1
+
+        self.X, self.Y = X, Y
+        # self.batch_size = batch_size
+        self.n_rays = n_rays
+        self.patch_size = patch_size
+        self.ss_grid = (slice(None),) + tuple(slice(0, None, g) for g in grid)
+        self.grid = tuple(grid)
+        self.use_gpu = bool(use_gpu)
+        if augmenter is None:
+            augmenter = lambda *args: args
+        callable(augmenter) or _raise(ValueError("augmenter must be None or callable"))
+        self.augmenter = augmenter
+        self.foreground_prob = foreground_prob
+
+        if self.use_gpu:
+            from gputools import max_filter
+            self.max_filter = lambda y, patch_size: max_filter(y.astype(np.float32), patch_size)
+        else:
+            from scipy.ndimage.filters import maximum_filter
+            self.max_filter = lambda y, patch_size: maximum_filter(y, patch_size, mode='constant')
+
+        self.maxfilter_patch_size = maxfilter_patch_size if maxfilter_patch_size is not None else self.patch_size
+
+        self.sample_ind_cache = sample_ind_cache
+        self._ind_cache_fg  = {}
+        self._ind_cache_all = {}
+        self.lock = threading.Lock()
+
+
+    def get_valid_inds(self, k, foreground_prob=None):
+        if foreground_prob is None:
+            foreground_prob = self.foreground_prob
+        foreground_only = np.random.uniform() < foreground_prob
+        _ind_cache = self._ind_cache_fg if foreground_only else self._ind_cache_all
+        if k in _ind_cache:
+            inds = _ind_cache[k]
+        else:
+            patch_filter = (lambda y,p: self.max_filter(y, self.maxfilter_patch_size) > 0) if foreground_only else None
+            inds = get_valid_inds(self.Y[k], self.patch_size, patch_filter=patch_filter)
+            if self.sample_ind_cache:
+                with self.lock:
+                    _ind_cache[k] = inds
+        if foreground_only and len(inds[0])==0:
+            # no foreground pixels available
+            return self.get_valid_inds(k, foreground_prob=0)
+        return inds
+
+
+    def channels_as_tuple(self, x):
+        if self.n_channel is None:
+            return (x,)
+        else:
+            return tuple(x[...,i] for i in range(self.n_channel))
+
+
+
+class StarDistBase(BaseModel):
+
+    def __init__(self, config, name=None, basedir='.'):
+        super().__init__(config=config, name=name, basedir=basedir)
+        threshs = dict(prob=None, nms=None)
+        if basedir is not None:
+            try:
+                threshs = load_json(str(self.logdir / 'thresholds.json'))
+                print("Loading thresholds from 'thresholds.json'.")
+                if threshs.get('prob') is None or not (0 < threshs.get('prob') < 1):
+                    print("- Invalid 'prob' threshold (%s), using default value." % str(threshs.get('prob')))
+                    threshs['prob'] = None
+                if threshs.get('nms') is None or not (0 < threshs.get('nms') < 1):
+                    print("- Invalid 'nms' threshold (%s), using default value." % str(threshs.get('nms')))
+                    threshs['nms'] = None
+            except FileNotFoundError:
+                if config is None and len(tuple(self.logdir.glob('*.h5'))) > 0:
+                    print("Couldn't load thresholds from 'thresholds.json', using default values. "
+                          "(Call 'optimize_thresholds' to change that.)")
+
+        self.thresholds = dict (
+            prob = 0.5 if threshs['prob'] is None else threshs['prob'],
+            nms  = 0.4 if threshs['nms']  is None else threshs['nms'],
+        )
+        print("Using default values: prob_thresh={prob:g}, nms_thresh={nms:g}.".format(prob=self.thresholds.prob, nms=self.thresholds.nms))
+
+
+    @property
+    def thresholds(self):
+        return self._thresholds
+
+    def _is_multiclass(self):
+        return (self.config.n_classes is not None)
+
+    def _parse_classes_arg(self, classes, length):
+        """ creates a proper classes tuple from different possible "classes" arguments in model.train()
+
+        classes can be
+          "auto" -> all objects will be assigned to the first foreground class (unless n_classes is None)
+          single integer -> all objects will be assigned that class
+          tuple, list, ndarray -> do nothing (needs to be of given length)
+
+        returns a tuple of given length
+        """
+        if isinstance(classes, str):
+            classes == "auto" or _raise(ValueError(f"classes = '{classes}': only 'auto' supported as string argument for classes"))
+            if self.config.n_classes is None:
+                classes = None
+            elif self.config.n_classes == 1:
+                classes = (1,)*length
+            else:
+                raise ValueError("using classes = 'auto' for n_classes > 1 not supported")
+        elif isinstance(classes, (tuple, list, np.ndarray)):
+            len(classes) == length or _raise(ValueError(f"len(classes) should be {length}!"))
+        else:
+            raise ValueError("classes should either be 'auto' or a list of scalars/label dicts")
+        return classes
+
+    @thresholds.setter
+    def thresholds(self, d):
+        self._thresholds = namedtuple('Thresholds',d.keys())(*d.values())
+
+
+    def prepare_for_training(self, optimizer=None):
+        """Prepare for neural network training.
+
+        Compiles the model and creates
+        `Keras Callbacks <https://keras.io/callbacks/>`_ to be used for training.
+
+        Note that this method will be implicitly called once by :func:`train`
+        (with default arguments) if not done so explicitly beforehand.
+
+        Parameters
+        ----------
+        optimizer : obj or None
+            Instance of a `Keras Optimizer <https://keras.io/optimizers/>`_ to be used for training.
+            If ``None`` (default), uses ``Adam`` with the learning rate specified in ``config``.
+
+        """
+        if optimizer is None:
+            optimizer = Adam(self.config.train_learning_rate)
+
+        masked_dist_loss = {'mse': masked_loss_mse,
+                            'mae': masked_loss_mae,
+                            'iou': masked_loss_iou,
+                            }[self.config.train_dist_loss]
+        prob_loss = 'binary_crossentropy'
+
+
+        def split_dist_true_mask(dist_true_mask):
+            return tf.split(dist_true_mask, num_or_size_splits=[self.config.n_rays,-1], axis=-1)
+
+        def dist_loss(dist_true_mask, dist_pred):
+            dist_true, dist_mask = split_dist_true_mask(dist_true_mask)
+            return masked_dist_loss(dist_mask, reg_weight=self.config.train_background_reg)(dist_true, dist_pred)
+
+        def dist_iou_metric(dist_true_mask, dist_pred):
+            dist_true, dist_mask = split_dist_true_mask(dist_true_mask)
+            return masked_metric_iou(dist_mask, reg_weight=0)(dist_true, dist_pred)
+
+        def relevant_mae(dist_true_mask, dist_pred):
+            dist_true, dist_mask = split_dist_true_mask(dist_true_mask)
+            return masked_metric_mae(dist_mask)(dist_true, dist_pred)
+
+        def relevant_mse(dist_true_mask, dist_pred):
+            dist_true, dist_mask = split_dist_true_mask(dist_true_mask)
+            return masked_metric_mse(dist_mask)(dist_true, dist_pred)
+
+
+        if self._is_multiclass():
+            prob_class_loss = weighted_categorical_crossentropy(self.config.train_class_weights, ndim=self.config.n_dim)
+            loss = [prob_loss, dist_loss, prob_class_loss]
+        else:
+            loss = [prob_loss, dist_loss]
+
+        self.keras_model.compile(optimizer, loss         = loss,
+                                            loss_weights = list(self.config.train_loss_weights),
+                                            metrics      = {'prob': kld,
+                                                            'dist': [relevant_mae, relevant_mse, dist_iou_metric]})
+
+        self.callbacks = []
+        if self.basedir is not None:
+            self.callbacks += self._checkpoint_callbacks()
+
+            if self.config.train_tensorboard:
+                if IS_TF_1:
+                    self.callbacks.append(CARETensorBoard(log_dir=str(self.logdir), prefix_with_timestamp=False, n_images=3, write_images=True, prob_out=False))
+                else:
+                    self.callbacks.append(TensorBoard(log_dir=str(self.logdir/'logs'), write_graph=False, profile_batch=0))
+
+        if self.config.train_reduce_lr is not None:
+            rlrop_params = self.config.train_reduce_lr
+            if 'verbose' not in rlrop_params:
+                rlrop_params['verbose'] = True
+            # TF2: add as first callback to put 'lr' in the logs for TensorBoard
+            self.callbacks.insert(0,ReduceLROnPlateau(**rlrop_params))
+
+        self._model_prepared = True
+
+
+    def _predict_setup(self, img, axes, normalizer, n_tiles, show_tile_progress, predict_kwargs):
+        """ Shared setup code between `predict` and `predict_sparse` """
+        if n_tiles is None:
+            n_tiles = [1]*img.ndim
+        try:
+            n_tiles = tuple(n_tiles)
+            img.ndim == len(n_tiles) or _raise(TypeError())
+        except TypeError:
+            raise ValueError("n_tiles must be an iterable of length %d" % img.ndim)
+        all(np.isscalar(t) and 1<=t and int(t)==t for t in n_tiles) or _raise(
+            ValueError("all values of n_tiles must be integer values >= 1"))
+
+        n_tiles = tuple(map(int,n_tiles))
+
+        axes     = self._normalize_axes(img, axes)
+        axes_net = self.config.axes
+
+        _permute_axes = self._make_permute_axes(axes, axes_net)
+        x = _permute_axes(img) # x has axes_net semantics
+
+        channel = axes_dict(axes_net)['C']
+        self.config.n_channel_in == x.shape[channel] or _raise(ValueError())
+        axes_net_div_by = self._axes_div_by(axes_net)
+
+        grid = tuple(self.config.grid)
+        len(grid) == len(axes_net)-1 or _raise(ValueError())
+        grid_dict = dict(zip(axes_net.replace('C',''),grid))
+
+        normalizer = self._check_normalizer_resizer(normalizer, None)[0]
+        resizer = StarDistPadAndCropResizer(grid=grid_dict)
+
+        x = normalizer.before(x, axes_net)
+        x = resizer.before(x, axes_net, axes_net_div_by)
+
+        if not _is_floatarray(x):
+            warnings.warn("Predicting on non-float input... ( forgot to normalize? )")
+
+        def predict_direct(x):
+            ys = self.keras_model.predict(x[np.newaxis], **predict_kwargs)
+            return tuple(y[0] for y in ys)
+
+        def tiling_setup():
+            assert np.prod(n_tiles) > 1
+            tiling_axes   = axes_net.replace('C','') # axes eligible for tiling
+            x_tiling_axis = tuple(axes_dict(axes_net)[a] for a in tiling_axes) # numerical axis ids for x
+            axes_net_tile_overlaps = self._axes_tile_overlap(axes_net)
+            # hack: permute tiling axis in the same way as img -> x was permuted
+            _n_tiles = _permute_axes(np.empty(n_tiles,bool)).shape
+            (all(_n_tiles[i] == 1 for i in range(x.ndim) if i not in x_tiling_axis) or
+                _raise(ValueError("entry of n_tiles > 1 only allowed for axes '%s'" % tiling_axes)))
+
+            sh = [s//grid_dict.get(a,1) for a,s in zip(axes_net,x.shape)]
+            sh[channel] = None
+            def create_empty_output(n_channel, dtype=np.float32):
+                sh[channel] = n_channel
+                return np.empty(sh,dtype)
+
+            if callable(show_tile_progress):
+                progress, _show_tile_progress = show_tile_progress, True
+            else:
+                progress, _show_tile_progress = tqdm, show_tile_progress
+
+            n_block_overlaps = [int(np.ceil(overlap/blocksize)) for overlap, blocksize
+                                in zip(axes_net_tile_overlaps, axes_net_div_by)]
+
+            num_tiles_used = total_n_tiles(x, _n_tiles, block_sizes=axes_net_div_by, n_block_overlaps=n_block_overlaps)
+
+            tile_generator = progress(tile_iterator(x, _n_tiles, block_sizes=axes_net_div_by, n_block_overlaps=n_block_overlaps),
+                                                    disable=(not _show_tile_progress), total=num_tiles_used)
+
+            return tile_generator, tuple(sh), create_empty_output
+
+        return x, axes, axes_net, axes_net_div_by, _permute_axes, resizer, n_tiles, grid, grid_dict, channel, predict_direct, tiling_setup
+
+
+    def _predict_generator(self, img, axes=None, normalizer=None, n_tiles=None, show_tile_progress=True, **predict_kwargs):
+        """Predict.
+
+        Parameters
+        ----------
+        img : :class:`numpy.ndarray`
+            Input image
+        axes : str or None
+            Axes of the input ``img``.
+            ``None`` denotes that axes of img are the same as denoted in the config.
+        normalizer : :class:`csbdeep.data.Normalizer` or None
+            (Optional) normalization of input image before prediction.
+            Note that the default (``None``) assumes ``img`` to be already normalized.
+        n_tiles : iterable or None
+            Out of memory (OOM) errors can occur if the input image is too large.
+            To avoid this problem, the input image is broken up into (overlapping) tiles
+            that are processed independently and re-assembled.
+            This parameter denotes a tuple of the number of tiles for every image axis (see ``axes``).
+            ``None`` denotes that no tiling should be used.
+        show_tile_progress: bool or callable
+            If boolean, indicates whether to show progress (via tqdm) during tiled prediction.
+            If callable, must be a drop-in replacement for tqdm.
+        show_tile_progress: bool
+            Whether to show progress during tiled prediction.
+        predict_kwargs: dict
+            Keyword arguments for ``predict`` function of Keras model.
+
+        Returns
+        -------
+        (:class:`numpy.ndarray`, :class:`numpy.ndarray`, [:class:`numpy.ndarray`])
+            Returns the tuple (`prob`, `dist`, [`prob_class`]) of per-pixel object probabilities and star-convex polygon/polyhedra distances.
+            In multiclass prediction mode, `prob_class` is the probability map for each of the 1+'n_classes' classes (first class is background)
+
+        """
+
+        x, axes, axes_net, axes_net_div_by, _permute_axes, resizer, n_tiles, grid, grid_dict, channel, predict_direct, tiling_setup = \
+            self._predict_setup(img, axes, normalizer, n_tiles, show_tile_progress, predict_kwargs)
+
+        if np.prod(n_tiles) > 1:
+            tile_generator, output_shape, create_empty_output = tiling_setup()
+
+            prob = create_empty_output(1)
+            dist = create_empty_output(self.config.n_rays)
+            if self._is_multiclass():
+                prob_class = create_empty_output(self.config.n_classes+1)
+                result = (prob, dist, prob_class)
+            else:
+                result = (prob, dist)
+
+            for tile, s_src, s_dst in tile_generator:
+                # predict_direct -> prob, dist, [prob_class if multi_class]
+                result_tile = predict_direct(tile)
+                # account for grid
+                s_src = [slice(s.start//grid_dict.get(a,1),s.stop//grid_dict.get(a,1)) for s,a in zip(s_src,axes_net)]
+                s_dst = [slice(s.start//grid_dict.get(a,1),s.stop//grid_dict.get(a,1)) for s,a in zip(s_dst,axes_net)]
+                # prob and dist have different channel dimensionality than image x
+                s_src[channel] = slice(None)
+                s_dst[channel] = slice(None)
+                s_src, s_dst = tuple(s_src), tuple(s_dst)
+                # print(s_src,s_dst)
+                for part, part_tile in zip(result, result_tile):
+                    part[s_dst] = part_tile[s_src]
+                yield  # yield None after each processed tile
+        else:
+            # predict_direct -> prob, dist, [prob_class if multi_class]
+            result = predict_direct(x)
+
+        result = [resizer.after(part, axes_net) for part in result]
+
+        # result = (prob, dist) for legacy or (prob, dist, prob_class) for multiclass
+
+        # prob
+        result[0] = np.take(result[0],0,axis=channel)
+        # dist
+        result[1] = np.maximum(1e-3, result[1]) # avoid small dist values to prevent problems with Qhull
+        result[1] = np.moveaxis(result[1],channel,-1)
+
+        if self._is_multiclass():
+            # prob_class
+            result[2] = np.moveaxis(result[2],channel,-1)
+
+        # last "yield" is the actual output that would have been "return"ed if this was a regular function
+        yield tuple(result)
+
+
+    @functools.wraps(_predict_generator)
+    def predict(self, *args, **kwargs):
+        # return last "yield"ed value of generator
+        r = None
+        for r in self._predict_generator(*args, **kwargs):
+            pass
+        return r
+
+
+    def _predict_sparse_generator(self, img, prob_thresh=None, axes=None, normalizer=None, n_tiles=None, show_tile_progress=True, b=2, **predict_kwargs):
+        """ Sparse version of model.predict()
+        Returns
+        -------
+        (prob, dist, [prob_class], points)   flat list of probs, dists, (optional prob_class) and points
+        """
+        if prob_thresh is None: prob_thresh = self.thresholds.prob
+
+        x, axes, axes_net, axes_net_div_by, _permute_axes, resizer, n_tiles, grid, grid_dict, channel, predict_direct, tiling_setup = \
+            self._predict_setup(img, axes, normalizer, n_tiles, show_tile_progress, predict_kwargs)
+
+        def _prep(prob, dist):
+            prob = np.take(prob,0,axis=channel)
+            dist = np.moveaxis(dist,channel,-1)
+            dist = np.maximum(1e-3, dist)
+            return prob, dist
+
+        proba, dista, pointsa, prob_class = [],[],[], []
+
+        if np.prod(n_tiles) > 1:
+            tile_generator, output_shape, create_empty_output = tiling_setup()
+
+            sh = list(output_shape)
+            sh[channel] = 1;
+
+            proba, dista, pointsa, prob_classa = [], [], [], []
+
+            for tile, s_src, s_dst in tile_generator:
+
+                results_tile = predict_direct(tile)
+
+                # account for grid
+                s_src = [slice(s.start//grid_dict.get(a,1),s.stop//grid_dict.get(a,1)) for s,a in zip(s_src,axes_net)]
+                s_dst = [slice(s.start//grid_dict.get(a,1),s.stop//grid_dict.get(a,1)) for s,a in zip(s_dst,axes_net)]
+                s_src[channel] = slice(None)
+                s_dst[channel] = slice(None)
+                s_src, s_dst = tuple(s_src), tuple(s_dst)
+
+                prob_tile, dist_tile = results_tile[:2]
+                prob_tile, dist_tile = _prep(prob_tile[s_src], dist_tile[s_src])
+
+                bs = list((b if s.start==0 else -1, b if s.stop==_sh else -1) for s,_sh in zip(s_dst, sh))
+                bs.pop(channel)
+                inds   = _ind_prob_thresh(prob_tile, prob_thresh, b=bs)
+                proba.extend(prob_tile[inds].copy())
+                dista.extend(dist_tile[inds].copy())
+                _points = np.stack(np.where(inds), axis=1)
+                offset = list(s.start for i,s in enumerate(s_dst))
+                offset.pop(channel)
+                _points = _points + np.array(offset).reshape((1,len(offset)))
+                _points = _points * np.array(self.config.grid).reshape((1,len(self.config.grid)))
+                pointsa.extend(_points)
+
+                if self._is_multiclass():
+                    p = results_tile[2][s_src].copy()
+                    p = np.moveaxis(p,channel,-1)
+                    prob_classa.extend(p[inds])
+                yield  # yield None after each processed tile
+
+        else:
+            # predict_direct -> prob, dist, [prob_class if multi_class]
+            results = predict_direct(x)
+            prob, dist = results[:2]
+            prob, dist = _prep(prob, dist)
+            inds   = _ind_prob_thresh(prob, prob_thresh, b=b)
+            proba = prob[inds].copy()
+            dista = dist[inds].copy()
+            _points = np.stack(np.where(inds), axis=1)
+            pointsa = (_points * np.array(self.config.grid).reshape((1,len(self.config.grid))))
+
+            if self._is_multiclass():
+                p = np.moveaxis(results[2],channel,-1)
+                prob_classa = p[inds].copy()
+
+
+        proba = np.asarray(proba)
+        dista = np.asarray(dista).reshape((-1,self.config.n_rays))
+        pointsa = np.asarray(pointsa).reshape((-1,self.config.n_dim))
+
+        idx = resizer.filter_points(x.ndim, pointsa, axes_net)
+        proba = proba[idx]
+        dista = dista[idx]
+        pointsa = pointsa[idx]
+
+        # last "yield" is the actual output that would have been "return"ed if this was a regular function
+        if self._is_multiclass():
+            prob_classa = np.asarray(prob_classa).reshape((-1,self.config.n_classes+1))
+            prob_classa = prob_classa[idx]
+            yield proba, dista, prob_classa, pointsa
+        else:
+            prob_classa = None
+            yield proba, dista, pointsa
+
+
+    @functools.wraps(_predict_sparse_generator)
+    def predict_sparse(self, *args, **kwargs):
+        # return last "yield"ed value of generator
+        r = None
+        for r in self._predict_sparse_generator(*args, **kwargs):
+            pass
+        return r
+
+
+    def _predict_instances_generator(self, img, axes=None, normalizer=None,
+                                     sparse=True,
+                                     prob_thresh=None, nms_thresh=None,
+                                     scale=None,
+                                     n_tiles=None, show_tile_progress=True,
+                                     verbose=False,
+                                     return_labels=True,
+                                     predict_kwargs=None, nms_kwargs=None,
+                                     overlap_label=None, return_predict=False):
+        """Predict instance segmentation from input image.
+
+        Parameters
+        ----------
+        img : :class:`numpy.ndarray`
+            Input image
+        axes : str or None
+            Axes of the input ``img``.
+            ``None`` denotes that axes of img are the same as denoted in the config.
+        normalizer : :class:`csbdeep.data.Normalizer` or None
+            (Optional) normalization of input image before prediction.
+            Note that the default (``None``) assumes ``img`` to be already normalized.
+        sparse: bool
+            If true, aggregate probabilities/distances sparsely during tiled
+            prediction to save memory (recommended).
+        prob_thresh : float or None
+            Consider only object candidates from pixels with predicted object probability
+            above this threshold (also see `optimize_thresholds`).
+        nms_thresh : float or None
+            Perform non-maximum suppression that considers two objects to be the same
+            when their area/surface overlap exceeds this threshold (also see `optimize_thresholds`).
+        scale: None or float or iterable
+            Scale the input image internally by this factor and rescale the output accordingly.
+            All spatial axes (X,Y,Z) will be scaled if a scalar value is provided.
+            Alternatively, multiple scale values (compatible with input `axes`) can be used
+            for more fine-grained control (scale values for non-spatial axes must be 1).
+        n_tiles : iterable or None
+            Out of memory (OOM) errors can occur if the input image is too large.
+            To avoid this problem, the input image is broken up into (overlapping) tiles
+            that are processed independently and re-assembled.
+            This parameter denotes a tuple of the number of tiles for every image axis (see ``axes``).
+            ``None`` denotes that no tiling should be used.
+        show_tile_progress: bool
+            Whether to show progress during tiled prediction.
+        verbose: bool
+            Whether to print some info messages.
+        return_labels: bool
+            Whether to create a label image, otherwise return None in its place.
+        predict_kwargs: dict
+            Keyword arguments for ``predict`` function of Keras model.
+        nms_kwargs: dict
+            Keyword arguments for non-maximum suppression.
+        overlap_label: scalar or None
+            if not None, label the regions where polygons overlap with that value
+        return_predict: bool
+            Also return the outputs of :func:`predict` (in a separate tuple)
+            If True, implies sparse = False
+
+        Returns
+        -------
+        (:class:`numpy.ndarray`, dict), (optional: return tuple of :func:`predict`)
+            Returns a tuple of the label instances image and also
+            a dictionary with the details (coordinates, etc.) of all remaining polygons/polyhedra.
+
+        """
+        if predict_kwargs is None:
+            predict_kwargs = {}
+        if nms_kwargs is None:
+            nms_kwargs = {}
+
+        if return_predict and sparse:
+            sparse = False
+            warnings.warn("Setting sparse to False because return_predict is True")
+
+        nms_kwargs.setdefault("verbose", verbose)
+
+        _axes         = self._normalize_axes(img, axes)
+        _axes_net     = self.config.axes
+        _permute_axes = self._make_permute_axes(_axes, _axes_net)
+        _shape_inst   = tuple(s for s,a in zip(_permute_axes(img).shape, _axes_net) if a != 'C')
+
+        if scale is not None:
+            if isinstance(scale, numbers.Number):
+                scale = tuple(scale if a in 'XYZ' else 1 for a in _axes)
+            scale = tuple(scale)
+            len(scale) == len(_axes) or _raise(ValueError(f"scale {scale} must be of length {len(_axes)}, i.e. one value for each of the axes {_axes}"))
+            for s,a in zip(scale,_axes):
+                s > 0 or _raise(ValueError("scale values must be greater than 0"))
+                (s in (1,None) or a in 'XYZ') or warnings.warn(f"replacing scale value {s} for non-spatial axis {a} with 1")
+            scale = tuple(s if a in 'XYZ' else 1 for s,a in zip(scale,_axes))
+            verbose and print(f"scaling image by factors {scale} for axes {_axes}")
+            img = ndi.zoom(img, scale, order=1)
+
+        yield 'predict'  # indicate that prediction is starting
+        res = None
+        if sparse:
+            for res in self._predict_sparse_generator(img, axes=axes, normalizer=normalizer, n_tiles=n_tiles,
+                                                      prob_thresh=prob_thresh, show_tile_progress=show_tile_progress, **predict_kwargs):
+                if res is None:
+                    yield 'tile'  # yield 'tile' each time a tile has been processed
+        else:
+            for res in self._predict_generator(img, axes=axes, normalizer=normalizer, n_tiles=n_tiles,
+                                               show_tile_progress=show_tile_progress, **predict_kwargs):
+                if res is None:
+                    yield 'tile'  # yield 'tile' each time a tile has been processed
+            res = tuple(res) + (None,)
+
+        if self._is_multiclass():
+            prob, dist, prob_class, points = res
+        else:
+            prob, dist, points = res
+            prob_class = None
+
+        yield 'nms'  # indicate that non-maximum suppression is starting
+        res_instances = self._instances_from_prediction(_shape_inst, prob, dist,
+                                                        points=points,
+                                                        prob_class=prob_class,
+                                                        prob_thresh=prob_thresh,
+                                                        nms_thresh=nms_thresh,
+                                                        scale=(None if scale is None else dict(zip(_axes,scale))),
+                                                        return_labels=return_labels,
+                                                        overlap_label=overlap_label,
+                                                        **nms_kwargs)
+
+        # last "yield" is the actual output that would have been "return"ed if this was a regular function
+        if return_predict:
+            yield res_instances, tuple(res[:-1])
+        else:
+            yield res_instances
+
+
+    @functools.wraps(_predict_instances_generator)
+    def predict_instances(self, *args, **kwargs):
+        # the reason why the actual computation happens as a generator function
+        # (in '_predict_instances_generator') is that the generator is called
+        # from the stardist napari plugin, which has its benefits regarding
+        # control flow and progress display. however, typical use cases should
+        # almost always use this function ('predict_instances'), and shouldn't
+        # even notice (thanks to @functools.wraps) that it wraps the generator
+        # function. note that similar reasoning applies to 'predict' and
+        # 'predict_sparse'.
+
+        # return last "yield"ed value of generator
+        r = None
+        for r in self._predict_instances_generator(*args, **kwargs):
+            pass
+        return r
+
+
+    # def _predict_instances_old(self, img, axes=None, normalizer=None,
+    #                       sparse = False,
+    #                       prob_thresh=None, nms_thresh=None,
+    #                       n_tiles=None, show_tile_progress=True,
+    #                       verbose = False,
+    #                       predict_kwargs=None, nms_kwargs=None, overlap_label=None):
+    #     """
+    #     old version, should be removed....
+    #     """
+    #     if predict_kwargs is None:
+    #         predict_kwargs = {}
+    #     if nms_kwargs is None:
+    #         nms_kwargs = {}
+
+    #     nms_kwargs.setdefault("verbose", verbose)
+
+    #     _axes         = self._normalize_axes(img, axes)
+    #     _axes_net     = self.config.axes
+    #     _permute_axes = self._make_permute_axes(_axes, _axes_net)
+    #     _shape_inst   = tuple(s for s,a in zip(_permute_axes(img).shape, _axes_net) if a != 'C')
+
+
+    #     res = self.predict(img, axes=axes, normalizer=normalizer,
+    #                                   n_tiles=n_tiles,
+    #                                   show_tile_progress=show_tile_progress,
+    #                                   **predict_kwargs)
+
+    #     res = tuple(res) + (None,)
+
+    #     if self._is_multiclass():
+    #         prob, dist, prob_class, points = res
+    #     else:
+    #         prob, dist, points = res
+    #         prob_class = None
+
+
+    #     return self._instances_from_prediction_old(_shape_inst, prob, dist,
+    #                                            points = points,
+    #                                            prob_class = prob_class,
+    #                                            prob_thresh=prob_thresh,
+    #                                            nms_thresh=nms_thresh,
+    #                                            overlap_label=overlap_label,
+    #                                            **nms_kwargs)
+
+
+    def predict_instances_big(self, img, axes, block_size, min_overlap, context=None,
+                              labels_out=None, labels_out_dtype=np.int32, show_progress=True, **kwargs):
+        """Predict instance segmentation from very large input images.
+
+        Intended to be used when `predict_instances` cannot be used due to memory limitations.
+        This function will break the input image into blocks and process them individually
+        via `predict_instances` and assemble all the partial results. If used as intended, the result
+        should be the same as if `predict_instances` was used directly on the whole image.
+
+        **Important**: The crucial assumption is that all predicted object instances are smaller than
+                       the provided `min_overlap`. Also, it must hold that: min_overlap + 2*context < block_size.
+
+        Example
+        -------
+        >>> img.shape
+        (20000, 20000)
+        >>> labels, polys = model.predict_instances_big(img, axes='YX', block_size=4096,
+                                                        min_overlap=128, context=128, n_tiles=(4,4))
+
+        Parameters
+        ----------
+        img: :class:`numpy.ndarray` or similar
+            Input image
+        axes: str
+            Axes of the input ``img`` (such as 'YX', 'ZYX', 'YXC', etc.)
+        block_size: int or iterable of int
+            Process input image in blocks of the provided shape.
+            (If a scalar value is given, it is used for all spatial image dimensions.)
+        min_overlap: int or iterable of int
+            Amount of guaranteed overlap between blocks.
+            (If a scalar value is given, it is used for all spatial image dimensions.)
+        context: int or iterable of int, or None
+            Amount of image context on all sides of a block, which is discarded.
+            If None, uses an automatic estimate that should work in many cases.
+            (If a scalar value is given, it is used for all spatial image dimensions.)
+        labels_out: :class:`numpy.ndarray` or similar, or None, or False
+            numpy array or similar (must be of correct shape) to which the label image is written.
+            If None, will allocate a numpy array of the correct shape and data type ``labels_out_dtype``.
+            If False, will not write the label image (useful if only the dictionary is needed).
+        labels_out_dtype: str or dtype
+            Data type of returned label image if ``labels_out=None`` (has no effect otherwise).
+        show_progress: bool
+            Show progress bar for block processing.
+        kwargs: dict
+            Keyword arguments for ``predict_instances``.
+
+        Returns
+        -------
+        (:class:`numpy.ndarray` or False, dict)
+            Returns the label image and a dictionary with the details (coordinates, etc.) of the polygons/polyhedra.
+
+        """
+        from ..big import _grid_divisible, BlockND, OBJECT_KEYS#, repaint_labels
+        from ..matching import relabel_sequential
+
+        n = img.ndim
+        axes = axes_check_and_normalize(axes, length=n)
+        grid = self._axes_div_by(axes)
+        axes_out = self._axes_out.replace('C','')
+        shape_dict = dict(zip(axes,img.shape))
+        shape_out = tuple(shape_dict[a] for a in axes_out)
+
+        if context is None:
+            context = self._axes_tile_overlap(axes)
+
+        if np.isscalar(block_size):  block_size  = n*[block_size]
+        if np.isscalar(min_overlap): min_overlap = n*[min_overlap]
+        if np.isscalar(context):     context     = n*[context]
+        block_size, min_overlap, context = list(block_size), list(min_overlap), list(context)
+        assert n == len(block_size) == len(min_overlap) == len(context)
+
+        if 'C' in axes:
+            # single block for channel axis
+            i = axes_dict(axes)['C']
+            # if (block_size[i], min_overlap[i], context[i]) != (None, None, None):
+            #     print("Ignoring values of 'block_size', 'min_overlap', and 'context' for channel axis " +
+            #           "(set to 'None' to avoid this warning).", file=sys.stderr, flush=True)
+            block_size[i] = img.shape[i]
+            min_overlap[i] = context[i] = 0
+
+        block_size  = tuple(_grid_divisible(g, v, name='block_size',  verbose=False) for v,g,a in zip(block_size, grid,axes))
+        min_overlap = tuple(_grid_divisible(g, v, name='min_overlap', verbose=False) for v,g,a in zip(min_overlap,grid,axes))
+        context     = tuple(_grid_divisible(g, v, name='context',     verbose=False) for v,g,a in zip(context,    grid,axes))
+
+        # print(f"input: shape {img.shape} with axes {axes}")
+        print(f'effective: block_size={block_size}, min_overlap={min_overlap}, context={context}', flush=True)
+
+        for a,c,o in zip(axes,context,self._axes_tile_overlap(axes)):
+            if c < o:
+                print(f"{a}: context of {c} is small, recommended to use at least {o}", flush=True)
+
+        # create block cover
+        blocks = BlockND.cover(img.shape, axes, block_size, min_overlap, context, grid)
+
+        if np.isscalar(labels_out) and bool(labels_out) is False:
+            labels_out = None
+        else:
+            if labels_out is None:
+                labels_out = np.zeros(shape_out, dtype=labels_out_dtype)
+            else:
+                labels_out.shape == shape_out or _raise(ValueError(f"'labels_out' must have shape {shape_out} (axes {axes_out})."))
+
+        polys_all = {}
+        # problem_ids = []
+        label_offset = 1
+
+        kwargs_override = dict(axes=axes, overlap_label=None, return_labels=True, return_predict=False)
+        if show_progress:
+            kwargs_override['show_tile_progress'] = False # disable progress for predict_instances
+        for k,v in kwargs_override.items():
+            if k in kwargs: print(f"changing '{k}' from {kwargs[k]} to {v}", flush=True)
+            kwargs[k] = v
+
+        blocks = tqdm(blocks, disable=(not show_progress))
+        # actual computation
+        for block in blocks:
+            labels, polys = self.predict_instances(block.read(img, axes=axes), **kwargs)
+            labels = block.crop_context(labels, axes=axes_out)
+            labels, polys = block.filter_objects(labels, polys, axes=axes_out)
+            # TODO: relabel_sequential is not very memory-efficient (will allocate memory proportional to label_offset)
+            # this should not change the order of labels
+            labels = relabel_sequential(labels, label_offset)[0]
+
+            # labels, fwd_map, _ = relabel_sequential(labels, label_offset)
+            # if len(incomplete) > 0:
+            #     problem_ids.extend([fwd_map[i] for i in incomplete])
+            #     if show_progress:
+            #         blocks.set_postfix_str(f"found {len(problem_ids)} problematic {'object' if len(problem_ids)==1 else 'objects'}")
+            if labels_out is not None:
+                block.write(labels_out, labels, axes=axes_out)
+
+            for k,v in polys.items():
+                polys_all.setdefault(k,[]).append(v)
+
+            label_offset += len(polys['prob'])
+            del labels
+
+        polys_all = {k: (np.concatenate(v) if k in OBJECT_KEYS else v[0]) for k,v in polys_all.items()}
+
+        # if labels_out is not None and len(problem_ids) > 0:
+        #     # if show_progress:
+        #     #     blocks.write('')
+        #     # print(f"Found {len(problem_ids)} objects that violate the 'min_overlap' assumption.", file=sys.stderr, flush=True)
+        #     repaint_labels(labels_out, problem_ids, polys_all, show_progress=False)
+
+        return labels_out, polys_all#, tuple(problem_ids)
+
+
+    def optimize_thresholds(self, X_val, Y_val, nms_threshs=[0.3,0.4,0.5], iou_threshs=[0.3,0.5,0.7], predict_kwargs=None, optimize_kwargs=None, save_to_json=True):
+        """Optimize two thresholds (probability, NMS overlap) necessary for predicting object instances.
+
+        Note that the default thresholds yield good results in many cases, but optimizing
+        the thresholds for a particular dataset can further improve performance.
+
+        The optimized thresholds are automatically used for all further predictions
+        and also written to the model directory.
+
+        See ``utils.optimize_threshold`` for details and possible choices for ``optimize_kwargs``.
+
+        Parameters
+        ----------
+        X_val : list of ndarray
+            (Validation) input images (must be normalized) to use for threshold tuning.
+        Y_val : list of ndarray
+            (Validation) label images to use for threshold tuning.
+        nms_threshs : list of float
+            List of overlap thresholds to be considered for NMS.
+            For each value in this list, optimization is run to find a corresponding prob_thresh value.
+        iou_threshs : list of float
+            List of intersection over union (IOU) thresholds for which
+            the (average) matching performance is considered to tune the thresholds.
+        predict_kwargs: dict
+            Keyword arguments for ``predict`` function of this class.
+            (If not provided, will guess value for `n_tiles` to prevent out of memory errors.)
+        optimize_kwargs: dict
+            Keyword arguments for ``utils.optimize_threshold`` function.
+
+        """
+        if predict_kwargs is None:
+            predict_kwargs = {}
+        if optimize_kwargs is None:
+            optimize_kwargs = {}
+
+        def _predict_kwargs(x):
+            if 'n_tiles' in predict_kwargs:
+                return predict_kwargs
+            else:
+                return {**predict_kwargs, 'n_tiles': self._guess_n_tiles(x), 'show_tile_progress': False}
+
+        # only take first two elements of predict in case multi class is activated
+        Yhat_val = [self.predict(x, **_predict_kwargs(x))[:2] for x in X_val]
+
+        opt_prob_thresh, opt_measure, opt_nms_thresh = None, -np.inf, None
+        for _opt_nms_thresh in nms_threshs:
+            _opt_prob_thresh, _opt_measure = optimize_threshold(Y_val, Yhat_val, model=self, nms_thresh=_opt_nms_thresh, iou_threshs=iou_threshs, **optimize_kwargs)
+            if _opt_measure > opt_measure:
+                opt_prob_thresh, opt_measure, opt_nms_thresh = _opt_prob_thresh, _opt_measure, _opt_nms_thresh
+        opt_threshs = dict(prob=opt_prob_thresh, nms=opt_nms_thresh)
+
+        self.thresholds = opt_threshs
+        print(end='', file=sys.stderr, flush=True)
+        print("Using optimized values: prob_thresh={prob:g}, nms_thresh={nms:g}.".format(prob=self.thresholds.prob, nms=self.thresholds.nms))
+        if save_to_json and self.basedir is not None:
+            print("Saving to 'thresholds.json'.")
+            save_json(opt_threshs, str(self.logdir / 'thresholds.json'))
+        return opt_threshs
+
+
+    def _guess_n_tiles(self, img):
+        axes = self._normalize_axes(img, axes=None)
+        shape = list(img.shape)
+        if 'C' in axes:
+            del shape[axes_dict(axes)['C']]
+        b = self.config.train_batch_size**(1.0/self.config.n_dim)
+        n_tiles = [int(np.ceil(s/(p*b))) for s,p in zip(shape,self.config.train_patch_size)]
+        if 'C' in axes:
+            n_tiles.insert(axes_dict(axes)['C'],1)
+        return tuple(n_tiles)
+
+
+    def _normalize_axes(self, img, axes):
+        if axes is None:
+            axes = self.config.axes
+            assert 'C' in axes
+            if img.ndim == len(axes)-1 and self.config.n_channel_in == 1:
+                # img has no dedicated channel axis, but 'C' always part of config axes
+                axes = axes.replace('C','')
+        return axes_check_and_normalize(axes, img.ndim)
+
+
+    def _compute_receptive_field(self, img_size=None):
+        # TODO: good enough?
+        from scipy.ndimage import zoom
+        if img_size is None:
+            img_size = tuple(g*(128 if self.config.n_dim==2 else 64) for g in self.config.grid)
+        if np.isscalar(img_size):
+            img_size = (img_size,) * self.config.n_dim
+        img_size = tuple(img_size)
+        # print(img_size)
+        assert all(_is_power_of_2(s) for s in img_size)
+        mid = tuple(s//2 for s in img_size)
+        x = np.zeros((1,)+img_size+(self.config.n_channel_in,), dtype=np.float32)
+        z = np.zeros_like(x)
+        x[(0,)+mid+(slice(None),)] = 1
+        y  = self.keras_model.predict(x)[0][0,...,0]
+        y0 = self.keras_model.predict(z)[0][0,...,0]
+        grid = tuple((np.array(x.shape[1:-1])/np.array(y.shape)).astype(int))
+        assert grid == self.config.grid
+        y  = zoom(y, grid,order=0)
+        y0 = zoom(y0,grid,order=0)
+        ind = np.where(np.abs(y-y0)>0)
+        return [(m-np.min(i), np.max(i)-m) for (m,i) in zip(mid,ind)]
+
+
+    def _axes_tile_overlap(self, query_axes):
+        query_axes = axes_check_and_normalize(query_axes)
+        try:
+            self._tile_overlap
+        except AttributeError:
+            self._tile_overlap = self._compute_receptive_field()
+        overlap = dict(zip(
+            self.config.axes.replace('C',''),
+            tuple(max(rf) for rf in self._tile_overlap)
+        ))
+        return tuple(overlap.get(a,0) for a in query_axes)
+
+
+    def export_TF(self, fname=None, single_output=True, upsample_grid=True):
+        """Export model to TensorFlow's SavedModel format that can be used e.g. in the Fiji plugin
+
+        Parameters
+        ----------
+        fname : str
+            Path of the zip file to store the model
+            If None, the default path "<modeldir>/TF_SavedModel.zip" is used
+        single_output: bool
+            If set, concatenates the two model outputs into a single output (note: this is currently mandatory for further use in Fiji)
+        upsample_grid: bool
+            If set, upsamples the output to the input shape (note: this is currently mandatory for further use in Fiji)
+        """
+        Concatenate, UpSampling2D, UpSampling3D, Conv2DTranspose, Conv3DTranspose = keras_import('layers', 'Concatenate', 'UpSampling2D', 'UpSampling3D', 'Conv2DTranspose', 'Conv3DTranspose')
+        Model = keras_import('models', 'Model')
+
+        if self.basedir is None and fname is None:
+            raise ValueError("Need explicit 'fname', since model directory not available (basedir=None).")
+
+        if self._is_multiclass():
+            warnings.warn("multi-class mode not supported yet, removing classification output from exported model")
+
+        grid = self.config.grid
+        prob = self.keras_model.outputs[0]
+        dist = self.keras_model.outputs[1]
+        assert self.config.n_dim in (2,3)
+
+        if upsample_grid and any(g>1 for g in grid):
+            # CSBDeep Fiji plugin needs same size input/output
+            # -> we need to upsample the outputs if grid > (1,1)
+            # note: upsampling prob with a transposed convolution creates sparse
+            #       prob output with less candidates than with standard upsampling
+            conv_transpose = Conv2DTranspose if self.config.n_dim==2 else Conv3DTranspose
+            upsampling     = UpSampling2D    if self.config.n_dim==2 else UpSampling3D
+            prob = conv_transpose(1, (1,)*self.config.n_dim,
+                                  strides=grid, padding='same',
+                                  kernel_initializer='ones', use_bias=False)(prob)
+            dist = upsampling(grid)(dist)
+
+        inputs  = self.keras_model.inputs[0]
+        outputs = Concatenate()([prob,dist]) if single_output else [prob,dist]
+        csbdeep_model = Model(inputs, outputs)
+
+        fname = (self.logdir / 'TF_SavedModel.zip') if fname is None else Path(fname)
+        export_SavedModel(csbdeep_model, str(fname))
+        return csbdeep_model
+
+
+
+class StarDistPadAndCropResizer(Resizer):
+
+    # TODO: check correctness
+    def __init__(self, grid, mode='reflect', **kwargs):
+        assert isinstance(grid, dict)
+        self.mode = mode
+        self.grid = grid
+        self.kwargs = kwargs
+
+
+    def before(self, x, axes, axes_div_by):
+        assert all(a%g==0 for g,a in zip((self.grid.get(a,1) for a in axes), axes_div_by))
+        axes = axes_check_and_normalize(axes,x.ndim)
+        def _split(v):
+            return 0, v # only pad at the end
+        self.pad = {
+            a : _split((div_n-s%div_n)%div_n)
+            for a, div_n, s in zip(axes, axes_div_by, x.shape)
+        }
+        x_pad = np.pad(x, tuple(self.pad[a] for a in axes), mode=self.mode, **self.kwargs)
+        self.padded_shape = dict(zip(axes,x_pad.shape))
+        if 'C' in self.padded_shape: del self.padded_shape['C']
+        return x_pad
+
+
+    def after(self, x, axes):
+        # axes can include 'C', which may not have been present in before()
+        axes = axes_check_and_normalize(axes,x.ndim)
+        assert all(s_pad == s * g for s,s_pad,g in zip(x.shape,
+                                                       (self.padded_shape.get(a,_s) for a,_s in zip(axes,x.shape)),
+                                                       (self.grid.get(a,1) for a in axes)))
+        # print(self.padded_shape)
+        # print(self.pad)
+        # print(self.grid)
+        crop = tuple (
+            slice(0, -(math.floor(p[1]/g)) if p[1]>=g else None)
+            for p,g in zip((self.pad.get(a,(0,0)) for a in axes),(self.grid.get(a,1) for a in axes))
+        )
+        # print(crop)
+        return x[crop]
+
+
+    def filter_points(self, ndim, points, axes):
+        """ returns indices of points inside crop region """
+        assert points.ndim==2
+        axes = axes_check_and_normalize(axes,ndim)
+
+        bounds = np.array(tuple(self.padded_shape[a]-self.pad[a][1] for a in axes if a.lower() in ('z','y','x')))
+        idx = np.where(np.all(points< bounds, 1))
+        return idx
+
+
+
+def _tf_version_at_least(version_string="1.0.0"):
+    from packaging import version
+    return version.parse(tf.__version__) >= version.parse(version_string)

stardist_pkg/models/model2d.py

@@ -0,0 +1,570 @@
+from __future__ import print_function, unicode_literals, absolute_import, division
+
+import numpy as np
+import warnings
+import math
+from tqdm import tqdm
+
+from csbdeep.models import BaseConfig
+from csbdeep.internals.blocks import unet_block
+from csbdeep.utils import _raise, backend_channels_last, axes_check_and_normalize, axes_dict
+from csbdeep.utils.tf import keras_import, IS_TF_1, CARETensorBoard, CARETensorBoardImage
+from skimage.segmentation import clear_border
+from skimage.measure import regionprops
+from scipy.ndimage import zoom
+from distutils.version import LooseVersion
+
+keras = keras_import()
+K = keras_import('backend')
+Input, Conv2D, MaxPooling2D = keras_import('layers', 'Input', 'Conv2D', 'MaxPooling2D')
+Model = keras_import('models', 'Model')
+
+from .base import StarDistBase, StarDistDataBase, _tf_version_at_least
+from ..sample_patches import sample_patches
+from ..utils import edt_prob, _normalize_grid, mask_to_categorical
+from ..geometry import star_dist, dist_to_coord, polygons_to_label
+from ..nms import non_maximum_suppression, non_maximum_suppression_sparse
+
+
+class StarDistData2D(StarDistDataBase):
+
+    def __init__(self, X, Y, batch_size, n_rays, length,
+                 n_classes=None, classes=None,
+                 patch_size=(256,256), b=32, grid=(1,1), shape_completion=False, augmenter=None, foreground_prob=0, **kwargs):
+
+        super().__init__(X=X, Y=Y, n_rays=n_rays, grid=grid,
+                         n_classes=n_classes, classes=classes,
+                         batch_size=batch_size, patch_size=patch_size, length=length,
+                         augmenter=augmenter, foreground_prob=foreground_prob, **kwargs)
+
+        self.shape_completion = bool(shape_completion)
+        if self.shape_completion and b > 0:
+            self.b = slice(b,-b),slice(b,-b)
+        else:
+            self.b = slice(None),slice(None)
+
+        self.sd_mode = 'opencl' if self.use_gpu else 'cpp'
+
+
+    def __getitem__(self, i):
+        idx = self.batch(i)
+        arrays = [sample_patches((self.Y[k],) + self.channels_as_tuple(self.X[k]),
+                                 patch_size=self.patch_size, n_samples=1,
+                                 valid_inds=self.get_valid_inds(k)) for k in idx]
+
+        if self.n_channel is None:
+            X, Y = list(zip(*[(x[0][self.b],y[0]) for y,x in arrays]))
+        else:
+            X, Y = list(zip(*[(np.stack([_x[0] for _x in x],axis=-1)[self.b], y[0]) for y,*x in arrays]))
+
+        X, Y = tuple(zip(*tuple(self.augmenter(_x, _y) for _x, _y in zip(X,Y))))
+
+
+        prob = np.stack([edt_prob(lbl[self.b][self.ss_grid[1:3]]) for lbl in Y])
+        # prob = np.stack([edt_prob(lbl[self.b]) for lbl in Y])
+        # prob = prob[self.ss_grid]
+
+        if self.shape_completion:
+            Y_cleared = [clear_border(lbl) for lbl in Y]
+            _dist     = np.stack([star_dist(lbl,self.n_rays,mode=self.sd_mode)[self.b+(slice(None),)] for lbl in Y_cleared])
+            dist      = _dist[self.ss_grid]
+            dist_mask = np.stack([edt_prob(lbl[self.b][self.ss_grid[1:3]]) for lbl in Y_cleared])
+        else:
+            # directly subsample with grid
+            dist      = np.stack([star_dist(lbl,self.n_rays,mode=self.sd_mode, grid=self.grid) for lbl in Y])
+            dist_mask = prob
+
+        X = np.stack(X)
+        if X.ndim == 3: # input image has no channel axis
+            X = np.expand_dims(X,-1)
+        prob = np.expand_dims(prob,-1)
+        dist_mask = np.expand_dims(dist_mask,-1)
+
+        # subsample wth given grid
+        # dist_mask = dist_mask[self.ss_grid]
+        # prob      = prob[self.ss_grid]
+
+        # append dist_mask to dist as additional channel
+        # dist_and_mask = np.concatenate([dist,dist_mask],axis=-1)
+        # faster than concatenate
+        dist_and_mask = np.empty(dist.shape[:-1]+(self.n_rays+1,), np.float32)
+        dist_and_mask[...,:-1] = dist
+        dist_and_mask[...,-1:] = dist_mask
+
+
+        if self.n_classes is None:
+            return [X], [prob,dist_and_mask]
+        else:
+            prob_class = np.stack(tuple((mask_to_categorical(y, self.n_classes, self.classes[k]) for y,k in zip(Y, idx))))
+
+            # TODO: investigate downsampling via simple indexing vs. using 'zoom'
+            # prob_class = prob_class[self.ss_grid]
+            # 'zoom' might lead to better registered maps (especially if upscaled later)
+            prob_class = zoom(prob_class, (1,)+tuple(1/g for g in self.grid)+(1,), order=0)
+
+            return [X], [prob,dist_and_mask, prob_class]
+
+
+
+class Config2D(BaseConfig):
+    """Configuration for a :class:`StarDist2D` model.
+
+    Parameters
+    ----------
+    axes : str or None
+        Axes of the input images.
+    n_rays : int
+        Number of radial directions for the star-convex polygon.
+        Recommended to use a power of 2 (default: 32).
+    n_channel_in : int
+        Number of channels of given input image (default: 1).
+    grid : (int,int)
+        Subsampling factors (must be powers of 2) for each of the axes.
+        Model will predict on a subsampled grid for increased efficiency and larger field of view.
+    n_classes : None or int
+        Number of object classes to use for multi-class predection (use None to disable)
+    backbone : str
+        Name of the neural network architecture to be used as backbone.
+    kwargs : dict
+        Overwrite (or add) configuration attributes (see below).
+
+
+    Attributes
+    ----------
+    unet_n_depth : int
+        Number of U-Net resolution levels (down/up-sampling layers).
+    unet_kernel_size : (int,int)
+        Convolution kernel size for all (U-Net) convolution layers.
+    unet_n_filter_base : int
+        Number of convolution kernels (feature channels) for first U-Net layer.
+        Doubled after each down-sampling layer.
+    unet_pool : (int,int)
+        Maxpooling size for all (U-Net) convolution layers.
+    net_conv_after_unet : int
+        Number of filters of the extra convolution layer after U-Net (0 to disable).
+    unet_* : *
+        Additional parameters for U-net backbone.
+    train_shape_completion : bool
+        Train model to predict complete shapes for partially visible objects at image boundary.
+    train_completion_crop : int
+        If 'train_shape_completion' is set to True, specify number of pixels to crop at boundary of training patches.
+        Should be chosen based on (largest) object sizes.
+    train_patch_size : (int,int)
+        Size of patches to be cropped from provided training images.
+    train_background_reg : float
+        Regularizer to encourage distance predictions on background regions to be 0.
+    train_foreground_only : float
+        Fraction (0..1) of patches that will only be sampled from regions that contain foreground pixels.
+    train_sample_cache : bool
+        Activate caching of valid patch regions for all training images (disable to save memory for large datasets)
+    train_dist_loss : str
+        Training loss for star-convex polygon distances ('mse' or 'mae').
+    train_loss_weights : tuple of float
+        Weights for losses relating to (probability, distance)
+    train_epochs : int
+        Number of training epochs.
+    train_steps_per_epoch : int
+        Number of parameter update steps per epoch.
+    train_learning_rate : float
+        Learning rate for training.
+    train_batch_size : int
+        Batch size for training.
+    train_n_val_patches : int
+        Number of patches to be extracted from validation images (``None`` = one patch per image).
+    train_tensorboard : bool
+        Enable TensorBoard for monitoring training progress.
+    train_reduce_lr : dict
+        Parameter :class:`dict` of ReduceLROnPlateau_ callback; set to ``None`` to disable.
+    use_gpu : bool
+        Indicate that the data generator should use OpenCL to do computations on the GPU.
+
+        .. _ReduceLROnPlateau: https://keras.io/api/callbacks/reduce_lr_on_plateau/
+    """
+
+    def __init__(self, axes='YX', n_rays=32, n_channel_in=1, grid=(1,1), n_classes=None, backbone='unet', **kwargs):
+        """See class docstring."""
+
+        super().__init__(axes=axes, n_channel_in=n_channel_in, n_channel_out=1+n_rays)
+
+        # directly set by parameters
+        self.n_rays                    = int(n_rays)
+        self.grid                      = _normalize_grid(grid,2)
+        self.backbone                  = str(backbone).lower()
+        self.n_classes                 = None if n_classes is None else int(n_classes)
+
+        # default config (can be overwritten by kwargs below)
+        if self.backbone == 'unet':
+            self.unet_n_depth          = 3
+            self.unet_kernel_size      = 3,3
+            self.unet_n_filter_base    = 32
+            self.unet_n_conv_per_depth = 2
+            self.unet_pool             = 2,2
+            self.unet_activation       = 'relu'
+            self.unet_last_activation  = 'relu'
+            self.unet_batch_norm       = False
+            self.unet_dropout          = 0.0
+            self.unet_prefix           = ''
+            self.net_conv_after_unet   = 128
+        else:
+            # TODO: resnet backbone for 2D model?
+            raise ValueError("backbone '%s' not supported." % self.backbone)
+
+        # net_mask_shape not needed but kept for legacy reasons
+        if backend_channels_last():
+            self.net_input_shape       = None,None,self.n_channel_in
+            self.net_mask_shape        = None,None,1
+        else:
+            self.net_input_shape       = self.n_channel_in,None,None
+            self.net_mask_shape        = 1,None,None
+
+        self.train_shape_completion    = False
+        self.train_completion_crop     = 32
+        self.train_patch_size          = 256,256
+        self.train_background_reg      = 1e-4
+        self.train_foreground_only     = 0.9
+        self.train_sample_cache        = True
+
+        self.train_dist_loss           = 'mae'
+        self.train_loss_weights        = (1,0.2) if self.n_classes is None else (1,0.2,1)
+        self.train_class_weights       = (1,1) if self.n_classes is None else (1,)*(self.n_classes+1)
+        self.train_epochs              = 400
+        self.train_steps_per_epoch     = 100
+        self.train_learning_rate       = 0.0003
+        self.train_batch_size          = 4
+        self.train_n_val_patches       = None
+        self.train_tensorboard         = True
+        # the parameter 'min_delta' was called 'epsilon' for keras<=2.1.5
+        min_delta_key = 'epsilon' if LooseVersion(keras.__version__)<=LooseVersion('2.1.5') else 'min_delta'
+        self.train_reduce_lr           = {'factor': 0.5, 'patience': 40, min_delta_key: 0}
+
+        self.use_gpu                   = False
+
+        # remove derived attributes that shouldn't be overwritten
+        for k in ('n_dim', 'n_channel_out'):
+            try: del kwargs[k]
+            except KeyError: pass
+
+        self.update_parameters(False, **kwargs)
+
+        # FIXME: put into is_valid()
+        if not len(self.train_loss_weights) == (2 if self.n_classes is None else 3):
+            raise ValueError(f"train_loss_weights {self.train_loss_weights} not compatible with n_classes ({self.n_classes}): must be 3 weights if n_classes is not None, otherwise 2")
+
+        if not len(self.train_class_weights) == (2 if self.n_classes is None else self.n_classes+1):
+            raise ValueError(f"train_class_weights {self.train_class_weights} not compatible with n_classes ({self.n_classes}): must be 'n_classes + 1' weights if n_classes is not None, otherwise 2")
+
+
+
+class StarDist2D(StarDistBase):
+    """StarDist2D model.
+
+    Parameters
+    ----------
+    config : :class:`Config` or None
+        Will be saved to disk as JSON (``config.json``).
+        If set to ``None``, will be loaded from disk (must exist).
+    name : str or None
+        Model name. Uses a timestamp if set to ``None`` (default).
+    basedir : str
+        Directory that contains (or will contain) a folder with the given model name.
+
+    Raises
+    ------
+    FileNotFoundError
+        If ``config=None`` and config cannot be loaded from disk.
+    ValueError
+        Illegal arguments, including invalid configuration.
+
+    Attributes
+    ----------
+    config : :class:`Config`
+        Configuration, as provided during instantiation.
+    keras_model : `Keras model <https://keras.io/getting-started/functional-api-guide/>`_
+        Keras neural network model.
+    name : str
+        Model name.
+    logdir : :class:`pathlib.Path`
+        Path to model folder (which stores configuration, weights, etc.)
+    """
+
+    def __init__(self, config=Config2D(), name=None, basedir='.'):
+        """See class docstring."""
+        super().__init__(config, name=name, basedir=basedir)
+
+
+    def _build(self):
+        self.config.backbone == 'unet' or _raise(NotImplementedError())
+        unet_kwargs = {k[len('unet_'):]:v for (k,v) in vars(self.config).items() if k.startswith('unet_')}
+
+        input_img = Input(self.config.net_input_shape, name='input')
+
+        # maxpool input image to grid size
+        pooled = np.array([1,1])
+        pooled_img = input_img
+        while tuple(pooled) != tuple(self.config.grid):
+            pool = 1 + (np.asarray(self.config.grid) > pooled)
+            pooled *= pool
+            for _ in range(self.config.unet_n_conv_per_depth):
+                pooled_img = Conv2D(self.config.unet_n_filter_base, self.config.unet_kernel_size,
+                                    padding='same', activation=self.config.unet_activation)(pooled_img)
+            pooled_img = MaxPooling2D(pool)(pooled_img)
+
+        unet_base = unet_block(**unet_kwargs)(pooled_img)
+
+        if self.config.net_conv_after_unet > 0:
+            unet = Conv2D(self.config.net_conv_after_unet, self.config.unet_kernel_size,
+                          name='features', padding='same', activation=self.config.unet_activation)(unet_base)
+        else:
+            unet = unet_base
+
+        output_prob = Conv2D(                 1, (1,1), name='prob', padding='same', activation='sigmoid')(unet)
+        output_dist = Conv2D(self.config.n_rays, (1,1), name='dist', padding='same', activation='linear')(unet)
+
+        # attach extra classification head when self.n_classes is given
+        if self._is_multiclass():
+            if self.config.net_conv_after_unet > 0:
+                unet_class  = Conv2D(self.config.net_conv_after_unet, self.config.unet_kernel_size,
+                                     name='features_class', padding='same', activation=self.config.unet_activation)(unet_base)
+            else:
+                unet_class  = unet_base
+
+            output_prob_class  = Conv2D(self.config.n_classes+1, (1,1), name='prob_class', padding='same', activation='softmax')(unet_class)
+            return Model([input_img], [output_prob,output_dist,output_prob_class])
+        else:
+            return Model([input_img], [output_prob,output_dist])
+
+
+    def train(self, X, Y, validation_data, classes='auto', augmenter=None, seed=None, epochs=None, steps_per_epoch=None, workers=1):
+        """Train the neural network with the given data.
+
+        Parameters
+        ----------
+        X : tuple, list, `numpy.ndarray`, `keras.utils.Sequence`
+            Input images
+        Y : tuple, list, `numpy.ndarray`, `keras.utils.Sequence`
+            Label masks
+        classes (optional): 'auto' or iterable of same length as X
+             label id -> class id mapping for each label mask of Y if multiclass prediction is activated (n_classes > 0)
+             list of dicts with label id -> class id (1,...,n_classes)
+             'auto' -> all objects will be assigned to the first non-background class,
+                       or will be ignored if config.n_classes is None
+        validation_data : tuple(:class:`numpy.ndarray`, :class:`numpy.ndarray`) or triple (if multiclass)
+            Tuple (triple if multiclass) of X,Y,[classes] validation data.
+        augmenter : None or callable
+            Function with expected signature ``xt, yt = augmenter(x, y)``
+            that takes in a single pair of input/label image (x,y) and returns
+            the transformed images (xt, yt) for the purpose of data augmentation
+            during training. Not applied to validation images.
+            Example:
+            def simple_augmenter(x,y):
+                x = x + 0.05*np.random.normal(0,1,x.shape)
+                return x,y
+        seed : int
+            Convenience to set ``np.random.seed(seed)``. (To obtain reproducible validation patches, etc.)
+        epochs : int
+            Optional argument to use instead of the value from ``config``.
+        steps_per_epoch : int
+            Optional argument to use instead of the value from ``config``.
+
+        Returns
+        -------
+        ``History`` object
+            See `Keras training history <https://keras.io/models/model/#fit>`_.
+
+        """
+        if seed is not None:
+            # https://keras.io/getting-started/faq/#how-can-i-obtain-reproducible-results-using-keras-during-development
+            np.random.seed(seed)
+        if epochs is None:
+            epochs = self.config.train_epochs
+        if steps_per_epoch is None:
+            steps_per_epoch = self.config.train_steps_per_epoch
+
+        classes = self._parse_classes_arg(classes, len(X))
+
+        if not self._is_multiclass() and classes is not None:
+            warnings.warn("Ignoring given classes as n_classes is set to None")
+
+        isinstance(validation_data,(list,tuple)) or _raise(ValueError())
+        if self._is_multiclass() and len(validation_data) == 2:
+            validation_data = tuple(validation_data) + ('auto',)
+        ((len(validation_data) == (3 if self._is_multiclass() else 2))
+            or _raise(ValueError(f'len(validation_data) = {len(validation_data)}, but should be {3 if self._is_multiclass() else 2}')))
+
+        patch_size = self.config.train_patch_size
+        axes = self.config.axes.replace('C','')
+        b = self.config.train_completion_crop if self.config.train_shape_completion else 0
+        div_by = self._axes_div_by(axes)
+        [(p-2*b) % d == 0 or _raise(ValueError(
+            "'train_patch_size' - 2*'train_completion_crop' must be divisible by {d} along axis '{a}'".format(a=a,d=d) if self.config.train_shape_completion else
+            "'train_patch_size' must be divisible by {d} along axis '{a}'".format(a=a,d=d)
+         )) for p,d,a in zip(patch_size,div_by,axes)]
+
+        if not self._model_prepared:
+            self.prepare_for_training()
+
+        data_kwargs = dict (
+            n_rays           = self.config.n_rays,
+            patch_size       = self.config.train_patch_size,
+            grid             = self.config.grid,
+            shape_completion = self.config.train_shape_completion,
+            b                = self.config.train_completion_crop,
+            use_gpu          = self.config.use_gpu,
+            foreground_prob  = self.config.train_foreground_only,
+            n_classes        = self.config.n_classes,
+            sample_ind_cache = self.config.train_sample_cache,
+        )
+
+        # generate validation data and store in numpy arrays
+        n_data_val = len(validation_data[0])
+        classes_val = self._parse_classes_arg(validation_data[2], n_data_val) if self._is_multiclass() else None
+        n_take = self.config.train_n_val_patches if self.config.train_n_val_patches is not None else n_data_val
+        _data_val = StarDistData2D(validation_data[0],validation_data[1], classes=classes_val, batch_size=n_take, length=1, **data_kwargs)
+        data_val = _data_val[0]
+
+        # expose data generator as member for general diagnostics
+        self.data_train = StarDistData2D(X, Y, classes=classes, batch_size=self.config.train_batch_size,
+                                         augmenter=augmenter, length=epochs*steps_per_epoch, **data_kwargs)
+
+        if self.config.train_tensorboard:
+            # show dist for three rays
+            _n = min(3, self.config.n_rays)
+            channel = axes_dict(self.config.axes)['C']
+            output_slices = [[slice(None)]*4,[slice(None)]*4]
+            output_slices[1][1+channel] = slice(0,(self.config.n_rays//_n)*_n, self.config.n_rays//_n)
+            if self._is_multiclass():
+                _n = min(3, self.config.n_classes)
+                output_slices += [[slice(None)]*4]
+                output_slices[2][1+channel] = slice(1,1+(self.config.n_classes//_n)*_n, self.config.n_classes//_n)
+
+            if IS_TF_1:
+                for cb in self.callbacks:
+                    if isinstance(cb,CARETensorBoard):
+                        cb.output_slices = output_slices
+                        # target image for dist includes dist_mask and thus has more channels than dist output
+                        cb.output_target_shapes = [None,[None]*4,None]
+                        cb.output_target_shapes[1][1+channel] = data_val[1][1].shape[1+channel]
+            elif self.basedir is not None and not any(isinstance(cb,CARETensorBoardImage) for cb in self.callbacks):
+                self.callbacks.append(CARETensorBoardImage(model=self.keras_model, data=data_val, log_dir=str(self.logdir/'logs'/'images'),
+                                                           n_images=3, prob_out=False, output_slices=output_slices))
+
+        fit = self.keras_model.fit_generator if IS_TF_1 else self.keras_model.fit
+        history = fit(iter(self.data_train), validation_data=data_val,
+                      epochs=epochs, steps_per_epoch=steps_per_epoch,
+                      workers=workers, use_multiprocessing=workers>1,
+                      callbacks=self.callbacks, verbose=1,
+                      # set validation batchsize to training batchsize (only works for tf >= 2.2)
+                      **(dict(validation_batch_size = self.config.train_batch_size) if _tf_version_at_least("2.2.0") else {}))
+        self._training_finished()
+
+        return history
+
+
+    # def _instances_from_prediction_old(self, img_shape, prob, dist,points = None, prob_class = None,  prob_thresh=None, nms_thresh=None, overlap_label = None, **nms_kwargs):
+    #     from stardist.geometry.geom2d import _polygons_to_label_old, _dist_to_coord_old
+    #     from stardist.nms import _non_maximum_suppression_old
+
+    #     if prob_thresh is None: prob_thresh = self.thresholds.prob
+    #     if nms_thresh  is None: nms_thresh  = self.thresholds.nms
+    #     if overlap_label is not None: raise NotImplementedError("overlap_label not supported for 2D yet!")
+
+    #     coord = _dist_to_coord_old(dist, grid=self.config.grid)
+    #     inds = _non_maximum_suppression_old(coord, prob, grid=self.config.grid,
+    #                                    prob_thresh=prob_thresh, nms_thresh=nms_thresh, **nms_kwargs)
+    #     labels = _polygons_to_label_old(coord, prob, inds, shape=img_shape)
+    #     # sort 'inds' such that ids in 'labels' map to entries in polygon dictionary entries
+    #     inds = inds[np.argsort(prob[inds[:,0],inds[:,1]])]
+    #     # adjust for grid
+    #     points = inds*np.array(self.config.grid)
+
+    #     res_dict = dict(coord=coord[inds[:,0],inds[:,1]], points=points, prob=prob[inds[:,0],inds[:,1]])
+
+    #     if prob_class is not None:
+    #         prob_class = np.asarray(prob_class)
+    #         res_dict.update(dict(class_prob = prob_class))
+
+    #     return labels, res_dict
+
+
+    def _instances_from_prediction(self, img_shape, prob, dist, points=None, prob_class=None, prob_thresh=None, nms_thresh=None, overlap_label=None, return_labels=True, scale=None, **nms_kwargs):
+        """
+        if points is None     -> dense prediction
+        if points is not None -> sparse prediction
+
+        if prob_class is None     -> single class prediction
+        if prob_class is not None -> multi  class prediction
+        """
+        if prob_thresh is None: prob_thresh = self.thresholds.prob
+        if nms_thresh  is None: nms_thresh  = self.thresholds.nms
+        if overlap_label is not None: raise NotImplementedError("overlap_label not supported for 2D yet!")
+
+        # sparse prediction
+        if points is not None:
+            points, probi, disti, indsi = non_maximum_suppression_sparse(dist, prob, points, nms_thresh=nms_thresh, **nms_kwargs)
+            if prob_class is not None:
+                prob_class = prob_class[indsi]
+
+        # dense prediction
+        else:
+            points, probi, disti = non_maximum_suppression(dist, prob, grid=self.config.grid,
+                                                           prob_thresh=prob_thresh, nms_thresh=nms_thresh, **nms_kwargs)
+            if prob_class is not None:
+                inds = tuple(p//g for p,g in zip(points.T, self.config.grid))
+                prob_class = prob_class[inds]
+
+        if scale is not None:
+            # need to undo the scaling given by the scale dict, e.g. scale = dict(X=0.5,Y=0.5):
+            #   1. re-scale points (origins of polygons)
+            #   2. re-scale coordinates (computed from distances) of (zero-origin) polygons 
+            if not (isinstance(scale,dict) and 'X' in scale and 'Y' in scale):
+                raise ValueError("scale must be a dictionary with entries for 'X' and 'Y'")
+            rescale = (1/scale['Y'],1/scale['X'])
+            points = points * np.array(rescale).reshape(1,2)
+        else:
+            rescale = (1,1)
+
+        if return_labels:
+            labels = polygons_to_label(disti, points, prob=probi, shape=img_shape, scale_dist=rescale)
+        else:
+            labels = None
+
+        coord = dist_to_coord(disti, points, scale_dist=rescale)
+        res_dict = dict(coord=coord, points=points, prob=probi)
+
+        # multi class prediction
+        if prob_class is not None:
+            prob_class = np.asarray(prob_class)
+            class_id = np.argmax(prob_class, axis=-1)
+            res_dict.update(dict(class_prob=prob_class, class_id=class_id))
+
+        return labels, res_dict
+
+
+    def _axes_div_by(self, query_axes):
+        self.config.backbone == 'unet' or _raise(NotImplementedError())
+        query_axes = axes_check_and_normalize(query_axes)
+        assert len(self.config.unet_pool) == len(self.config.grid)
+        div_by = dict(zip(
+            self.config.axes.replace('C',''),
+            tuple(p**self.config.unet_n_depth * g for p,g in zip(self.config.unet_pool,self.config.grid))
+        ))
+        return tuple(div_by.get(a,1) for a in query_axes)
+
+
+    # def _axes_tile_overlap(self, query_axes):
+    #     self.config.backbone == 'unet' or _raise(NotImplementedError())
+    #     query_axes = axes_check_and_normalize(query_axes)
+    #     assert len(self.config.unet_pool) == len(self.config.grid) == len(self.config.unet_kernel_size)
+    #     # TODO: compute this properly when any value of grid > 1
+    #     # all(g==1 for g in self.config.grid) or warnings.warn('FIXME')
+    #     overlap = dict(zip(
+    #         self.config.axes.replace('C',''),
+    #         tuple(tile_overlap(self.config.unet_n_depth + int(np.log2(g)), k, p)
+    #               for p,k,g in zip(self.config.unet_pool,self.config.unet_kernel_size,self.config.grid))
+    #     ))
+    #     return tuple(overlap.get(a,0) for a in query_axes)
+
+
+    @property
+    def _config_class(self):
+        return Config2D

stardist_pkg/nms.py

@@ -0,0 +1,387 @@
+from __future__ import print_function, unicode_literals, absolute_import, division
+import numpy as np
+from time import time
+from .utils import _normalize_grid
+
+def _ind_prob_thresh(prob, prob_thresh, b=2):
+    if b is not None and np.isscalar(b):
+        b = ((b,b),)*prob.ndim
+
+    ind_thresh = prob > prob_thresh
+    if b is not None:
+        _ind_thresh = np.zeros_like(ind_thresh)
+        ss = tuple(slice(_bs[0] if _bs[0]>0 else None,
+                         -_bs[1] if _bs[1]>0 else None)  for _bs in b)
+        _ind_thresh[ss] = True
+        ind_thresh &= _ind_thresh
+    return ind_thresh
+
+
+def _non_maximum_suppression_old(coord, prob, grid=(1,1), b=2, nms_thresh=0.5, prob_thresh=0.5, verbose=False, max_bbox_search=True):
+    """2D coordinates of the polys that survive from a given prediction (prob, coord)
+
+    prob.shape = (Ny,Nx)
+    coord.shape = (Ny,Nx,2,n_rays)
+
+    b: don't use pixel closer than b pixels to the image boundary
+
+    returns retained points
+    """
+    from .lib.stardist2d import c_non_max_suppression_inds_old
+
+    # TODO: using b>0 with grid>1 can suppress small/cropped objects at the image boundary
+
+    assert prob.ndim == 2
+    assert coord.ndim == 4
+    grid = _normalize_grid(grid,2)
+
+    # mask = prob > prob_thresh
+    # if b is not None and b > 0:
+    #     _mask = np.zeros_like(mask)
+    #     _mask[b:-b,b:-b] = True
+    #     mask &= _mask
+
+    mask = _ind_prob_thresh(prob, prob_thresh, b)
+
+    polygons = coord[mask]
+    scores   = prob[mask]
+
+    # sort scores descendingly
+    ind = np.argsort(scores)[::-1]
+    survivors = np.zeros(len(ind), bool)
+    polygons = polygons[ind]
+    scores = scores[ind]
+
+    if max_bbox_search:
+        # map pixel indices to ids of sorted polygons (-1 => polygon at that pixel not a candidate)
+        mapping = -np.ones(mask.shape,np.int32)
+        mapping.flat[ np.flatnonzero(mask)[ind] ] = range(len(ind))
+    else:
+        mapping = np.empty((0,0),np.int32)
+
+    if verbose:
+        t = time()
+
+    survivors[ind] = c_non_max_suppression_inds_old(np.ascontiguousarray(polygons.astype(np.int32)),
+                    mapping, np.float32(nms_thresh), np.int32(max_bbox_search),
+                    np.int32(grid[0]), np.int32(grid[1]),np.int32(verbose))
+
+    if verbose:
+        print("keeping %s/%s polygons" % (np.count_nonzero(survivors), len(polygons)))
+        print("NMS took %.4f s" % (time() - t))
+
+    points = np.stack([ii[survivors] for ii in np.nonzero(mask)],axis=-1)
+    return points
+
+
+def non_maximum_suppression(dist, prob, grid=(1,1), b=2, nms_thresh=0.5, prob_thresh=0.5,
+                            use_bbox=True, use_kdtree=True, verbose=False,cut=False):
+    """Non-Maximum-Supression of 2D polygons
+
+    Retains only polygons whose overlap is smaller than nms_thresh
+
+    dist.shape = (Ny,Nx, n_rays)
+    prob.shape = (Ny,Nx)
+
+    returns the retained points, probabilities, and distances:
+
+    points, prob, dist = non_maximum_suppression(dist, prob, ....
+
+    """
+
+    # TODO: using b>0 with grid>1 can suppress small/cropped objects at the image boundary
+
+    assert prob.ndim == 2 and dist.ndim == 3  and prob.shape == dist.shape[:2]
+    dist = np.asarray(dist)
+    prob = np.asarray(prob)
+    n_rays = dist.shape[-1]
+
+    grid = _normalize_grid(grid,2)
+
+    # mask = prob > prob_thresh
+    # if b is not None and b > 0:
+    #     _mask = np.zeros_like(mask)
+    #     _mask[b:-b,b:-b] = True
+    #     mask &= _mask
+
+    mask = _ind_prob_thresh(prob, prob_thresh, b)
+    points = np.stack(np.where(mask), axis=1)
+
+    dist   = dist[mask]
+    scores = prob[mask]
+
+    # sort scores descendingly
+    ind = np.argsort(scores)[::-1]
+    if cut is True and ind.shape[0] > 20000:
+    #if cut is True and :
+        ind = ind[:round(ind.shape[0]*0.5)]
+    dist   = dist[ind]
+    scores = scores[ind]
+    points = points[ind]
+
+    points = (points * np.array(grid).reshape((1,2)))
+
+    if verbose:
+        t = time()
+
+    inds = non_maximum_suppression_inds(dist, points.astype(np.int32, copy=False), scores=scores,
+                                        use_bbox=use_bbox, use_kdtree=use_kdtree,
+                                        thresh=nms_thresh, verbose=verbose)
+
+    if verbose:
+        print("keeping %s/%s polygons" % (np.count_nonzero(inds), len(inds)))
+        print("NMS took %.4f s" % (time() - t))
+
+    return points[inds], scores[inds], dist[inds]
+
+
+def non_maximum_suppression_sparse(dist, prob, points, b=2, nms_thresh=0.5,
+                                   use_bbox=True, use_kdtree = True, verbose=False):
+    """Non-Maximum-Supression of 2D polygons from a list of dists, probs (scores), and points
+
+    Retains only polyhedra whose overlap is smaller than nms_thresh
+
+    dist.shape = (n_polys, n_rays)
+    prob.shape = (n_polys,)
+    points.shape = (n_polys,2)
+
+    returns the retained instances
+
+    (pointsi, probi, disti, indsi)
+
+    with
+    pointsi = points[indsi] ...
+
+    """
+
+    # TODO: using b>0 with grid>1 can suppress small/cropped objects at the image boundary
+
+    dist = np.asarray(dist)
+    prob = np.asarray(prob)
+    points = np.asarray(points)
+    n_rays = dist.shape[-1]
+
+    assert dist.ndim == 2 and prob.ndim == 1 and points.ndim == 2 and \
+        points.shape[-1]==2 and len(prob) == len(dist) == len(points)
+
+    verbose and print("predicting instances with nms_thresh = {nms_thresh}".format(nms_thresh=nms_thresh), flush=True)
+
+    inds_original = np.arange(len(prob))
+    _sorted = np.argsort(prob)[::-1]
+    probi = prob[_sorted]
+    disti = dist[_sorted]
+    pointsi = points[_sorted]
+    inds_original = inds_original[_sorted]
+
+    if verbose:
+        print("non-maximum suppression...")
+        t = time()
+
+    inds = non_maximum_suppression_inds(disti, pointsi, scores=probi, thresh=nms_thresh, use_kdtree = use_kdtree, verbose=verbose)
+
+    if verbose:
+        print("keeping %s/%s polyhedra" % (np.count_nonzero(inds), len(inds)))
+        print("NMS took %.4f s" % (time() - t))
+
+    return pointsi[inds], probi[inds], disti[inds], inds_original[inds]
+
+
+def non_maximum_suppression_inds(dist, points, scores, thresh=0.5, use_bbox=True, use_kdtree = True, verbose=1):
+    """
+    Applies non maximum supression to ray-convex polygons given by dists and points
+    sorted by scores and IoU threshold
+
+    P1 will suppress P2, if IoU(P1,P2) > thresh
+
+    with IoU(P1,P2) = Ainter(P1,P2) / min(A(P1),A(P2))
+
+    i.e. the smaller thresh, the more polygons will be supressed
+
+    dist.shape = (n_poly, n_rays)
+    point.shape = (n_poly, 2)
+    score.shape = (n_poly,)
+
+    returns indices of selected polygons
+    """
+
+    from stardist.lib.stardist2d import c_non_max_suppression_inds
+
+    assert dist.ndim == 2
+    assert points.ndim == 2
+
+    n_poly = dist.shape[0]
+
+    if scores is None:
+        scores = np.ones(n_poly)
+
+    assert len(scores) == n_poly
+    assert points.shape[0] == n_poly
+
+    def _prep(x, dtype):
+        return np.ascontiguousarray(x.astype(dtype, copy=False))
+
+    inds = c_non_max_suppression_inds(_prep(dist,  np.float32),
+                                      _prep(points, np.float32),
+                                      int(use_kdtree),
+                                      int(use_bbox),
+                                      int(verbose),
+                                      np.float32(thresh))
+
+    return inds
+
+
+#########
+
+
+def non_maximum_suppression_3d(dist, prob, rays, grid=(1,1,1), b=2, nms_thresh=0.5, prob_thresh=0.5, use_bbox=True, use_kdtree=True, verbose=False):
+    """Non-Maximum-Supression of 3D polyhedra
+
+    Retains only polyhedra whose overlap is smaller than nms_thresh
+
+    dist.shape = (Nz,Ny,Nx, n_rays)
+    prob.shape = (Nz,Ny,Nx)
+
+    returns the retained points, probabilities, and distances:
+
+    points, prob, dist = non_maximum_suppression_3d(dist, prob, ....
+    """
+
+    # TODO: using b>0 with grid>1 can suppress small/cropped objects at the image boundary
+
+    dist = np.asarray(dist)
+    prob = np.asarray(prob)
+
+    assert prob.ndim == 3 and dist.ndim == 4 and dist.shape[-1] == len(rays) and prob.shape == dist.shape[:3]
+
+    grid = _normalize_grid(grid,3)
+
+    verbose and print("predicting instances with prob_thresh = {prob_thresh} and nms_thresh = {nms_thresh}".format(prob_thresh=prob_thresh, nms_thresh=nms_thresh), flush=True)
+
+    # ind_thresh = prob > prob_thresh
+    # if b is not None and b > 0:
+    #     _ind_thresh = np.zeros_like(ind_thresh)
+    #     _ind_thresh[b:-b,b:-b,b:-b] = True
+    #     ind_thresh &= _ind_thresh
+
+    ind_thresh = _ind_prob_thresh(prob, prob_thresh, b)
+    points = np.stack(np.where(ind_thresh), axis=1)
+    verbose and print("found %s candidates"%len(points))
+    probi = prob[ind_thresh]
+    disti = dist[ind_thresh]
+
+    _sorted = np.argsort(probi)[::-1]
+    probi = probi[_sorted]
+    disti = disti[_sorted]
+    points = points[_sorted]
+
+    verbose and print("non-maximum suppression...")
+    points = (points * np.array(grid).reshape((1,3)))
+
+    inds = non_maximum_suppression_3d_inds(disti, points, rays=rays, scores=probi, thresh=nms_thresh,
+                                           use_bbox=use_bbox, use_kdtree = use_kdtree,
+                                           verbose=verbose)
+
+    verbose and print("keeping %s/%s polyhedra" % (np.count_nonzero(inds), len(inds)))
+    return points[inds], probi[inds], disti[inds]
+
+
+def non_maximum_suppression_3d_sparse(dist, prob, points, rays, b=2, nms_thresh=0.5, use_kdtree = True, verbose=False):
+    """Non-Maximum-Supression of 3D polyhedra from a list of dists, probs and points
+
+    Retains only polyhedra whose overlap is smaller than nms_thresh
+    dist.shape = (n_polys, n_rays)
+    prob.shape = (n_polys,)
+    points.shape = (n_polys,3)
+
+    returns the retained instances
+
+    (pointsi, probi, disti, indsi)
+
+    with
+    pointsi = points[indsi] ...
+    """
+
+    # TODO: using b>0 with grid>1 can suppress small/cropped objects at the image boundary
+
+    dist = np.asarray(dist)
+    prob = np.asarray(prob)
+    points = np.asarray(points)
+
+    assert dist.ndim == 2 and prob.ndim == 1 and points.ndim == 2 and \
+        dist.shape[-1] == len(rays) and points.shape[-1]==3 and len(prob) == len(dist) == len(points)
+
+    verbose and print("predicting instances with nms_thresh = {nms_thresh}".format(nms_thresh=nms_thresh), flush=True)
+
+    inds_original = np.arange(len(prob))
+    _sorted = np.argsort(prob)[::-1]
+    probi = prob[_sorted]
+    disti = dist[_sorted]
+    pointsi = points[_sorted]
+    inds_original = inds_original[_sorted]
+
+    verbose and print("non-maximum suppression...")
+
+    inds = non_maximum_suppression_3d_inds(disti, pointsi, rays=rays, scores=probi, thresh=nms_thresh, use_kdtree = use_kdtree, verbose=verbose)
+
+    verbose and print("keeping %s/%s polyhedra" % (np.count_nonzero(inds), len(inds)))
+    return pointsi[inds], probi[inds], disti[inds], inds_original[inds]
+
+
+def non_maximum_suppression_3d_inds(dist, points, rays, scores, thresh=0.5, use_bbox=True, use_kdtree = True, verbose=1):
+    """
+    Applies non maximum supression to ray-convex polyhedra given by dists and rays
+    sorted by scores and IoU threshold
+
+    P1 will suppress P2, if IoU(P1,P2) > thresh
+
+    with IoU(P1,P2) = Ainter(P1,P2) / min(A(P1),A(P2))
+
+    i.e. the smaller thresh, the more polygons will be supressed
+
+    dist.shape = (n_poly, n_rays)
+    point.shape = (n_poly, 3)
+    score.shape = (n_poly,)
+
+    returns indices of selected polygons
+    """
+    from .lib.stardist3d import c_non_max_suppression_inds
+
+    assert dist.ndim == 2
+    assert points.ndim == 2
+    assert dist.shape[1] == len(rays)
+
+    n_poly = dist.shape[0]
+
+    if scores is None:
+        scores = np.ones(n_poly)
+
+    assert len(scores) == n_poly
+    assert points.shape[0] == n_poly
+
+    # sort scores descendingly
+    ind = np.argsort(scores)[::-1]
+    survivors = np.ones(n_poly, bool)
+    dist = dist[ind]
+    points = points[ind]
+    scores = scores[ind]
+
+    def _prep(x, dtype):
+        return np.ascontiguousarray(x.astype(dtype, copy=False))
+
+    if verbose:
+        t = time()
+
+    survivors[ind] = c_non_max_suppression_inds(_prep(dist, np.float32),
+                                                _prep(points, np.float32),
+                                                _prep(rays.vertices, np.float32),
+                                                _prep(rays.faces, np.int32),
+                                                _prep(scores, np.float32),
+                                                int(use_bbox),
+                                                int(use_kdtree),
+                                                int(verbose),
+                                                np.float32(thresh))
+
+    if verbose:
+        print("NMS took %.4f s" % (time() - t))
+
+    return survivors

stardist_pkg/rays3d.py

@@ -0,0 +1,373 @@
+"""
+Ray factory
+
+classes that provide vertex and triangle information for rays on spheres
+
+Example:
+
+    rays = Rays_Tetra(n_level = 4)
+
+    print(rays.vertices)
+    print(rays.faces)
+
+"""
+from __future__ import print_function, unicode_literals, absolute_import, division
+import numpy as np
+from scipy.spatial import ConvexHull
+import copy
+import warnings
+
+class Rays_Base(object):
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+        self._vertices, self._faces = self.setup_vertices_faces()
+        self._vertices = np.asarray(self._vertices, np.float32)
+        self._faces = np.asarray(self._faces, int)
+        self._faces = np.asanyarray(self._faces)
+
+    def setup_vertices_faces(self):
+        """has to return
+
+         verts , faces
+
+         verts = ( (z_1,y_1,x_1), ... )
+         faces ( (0,1,2), (2,3,4), ... )
+
+         """
+        raise NotImplementedError()
+
+    @property
+    def vertices(self):
+        """read-only property"""
+        return self._vertices.copy()
+
+    @property
+    def faces(self):
+        """read-only property"""
+        return self._faces.copy()
+
+    def __getitem__(self, i):
+        return self.vertices[i]
+
+    def __len__(self):
+        return len(self._vertices)
+
+    def __repr__(self):
+        def _conv(x):
+            if isinstance(x,(tuple, list, np.ndarray)):
+                return "_".join(_conv(_x) for _x in x)
+            if isinstance(x,float):
+                return "%.2f"%x
+            return str(x)
+        return "%s_%s" % (self.__class__.__name__, "_".join("%s_%s" % (k, _conv(v)) for k, v in sorted(self.kwargs.items())))
+    
+    def to_json(self):
+        return {
+            "name": self.__class__.__name__,
+            "kwargs": self.kwargs
+        }
+
+    def dist_loss_weights(self, anisotropy = (1,1,1)):
+        """returns the anisotropy corrected weights for each ray"""
+        anisotropy = np.array(anisotropy)
+        assert anisotropy.shape == (3,)
+        return np.linalg.norm(self.vertices*anisotropy, axis = -1)
+
+    def volume(self, dist=None):
+        """volume of the starconvex polyhedron spanned by dist (if None, uses dist=1)
+        dist can be a nD array, but the last dimension has to be of length n_rays
+        """
+        if dist is None: dist = np.ones_like(self.vertices)
+
+        dist = np.asarray(dist)
+        
+        if not dist.shape[-1]==len(self.vertices):
+            raise ValueError("last dimension of dist should have length len(rays.vertices)")
+        # all the shuffling below is to allow dist to be an arbitrary sized array (with last dim n_rays)
+        # self.vertices -> (n_rays,3)
+        # dist -> (m,n,..., n_rays)
+        
+        # dist  -> (m,n,..., n_rays, 3)
+        dist = np.repeat(np.expand_dims(dist,-1), 3, axis = -1)
+        # verts  -> (m,n,..., n_rays, 3)
+        verts = np.broadcast_to(self.vertices, dist.shape)
+
+        # dist, verts  -> (n_rays, m,n, ..., 3)        
+        dist = np.moveaxis(dist,-2,0)
+        verts = np.moveaxis(verts,-2,0)
+
+        # vs -> (n_faces, 3, m, n, ..., 3)
+        vs = (dist*verts)[self.faces]
+        # vs -> (n_faces, m, n, ..., 3, 3)
+        vs = np.moveaxis(vs, 1,-2)
+        # vs -> (n_faces * m * n, 3, 3)        
+        vs = vs.reshape((len(self.faces)*int(np.prod(dist.shape[1:-1])),3,3))
+        d = np.linalg.det(list(vs)).reshape((len(self.faces),)+dist.shape[1:-1])
+        
+        return -1./6*np.sum(d, axis = 0)
+    
+    def surface(self, dist=None):
+        """surface area of the starconvex polyhedron spanned by dist (if None, uses dist=1)"""
+        dist = np.asarray(dist)
+        
+        if not dist.shape[-1]==len(self.vertices):
+            raise ValueError("last dimension of dist should have length len(rays.vertices)")
+
+        # self.vertices -> (n_rays,3)
+        # dist -> (m,n,..., n_rays)
+        
+        # all the shuffling below is to allow dist to be an arbitrary sized array (with last dim n_rays)
+        
+        # dist  -> (m,n,..., n_rays, 3)
+        dist = np.repeat(np.expand_dims(dist,-1), 3, axis = -1)
+        # verts  -> (m,n,..., n_rays, 3)
+        verts = np.broadcast_to(self.vertices, dist.shape)
+
+        # dist, verts  -> (n_rays, m,n, ..., 3)        
+        dist = np.moveaxis(dist,-2,0)
+        verts = np.moveaxis(verts,-2,0)
+
+        # vs -> (n_faces, 3, m, n, ..., 3)
+        vs = (dist*verts)[self.faces]
+        # vs -> (n_faces, m, n, ..., 3, 3)
+        vs = np.moveaxis(vs, 1,-2)
+        # vs -> (n_faces * m * n, 3, 3)        
+        vs = vs.reshape((len(self.faces)*int(np.prod(dist.shape[1:-1])),3,3))
+       
+        pa = vs[...,1,:]-vs[...,0,:]
+        pb = vs[...,2,:]-vs[...,0,:]
+
+        d = .5*np.linalg.norm(np.cross(list(pa), list(pb)), axis = -1)
+        d = d.reshape((len(self.faces),)+dist.shape[1:-1])
+        return np.sum(d, axis = 0)
+
+    
+    def copy(self, scale=(1,1,1)):
+        """ returns a copy whose vertices are scaled by given factor"""
+        scale = np.asarray(scale)
+        assert scale.shape == (3,)
+        res = copy.deepcopy(self)
+        res._vertices *= scale[np.newaxis]
+        return res 
+
+
+
+    
+def rays_from_json(d):
+    return eval(d["name"])(**d["kwargs"])
+
+
+################################################################
+
+class Rays_Explicit(Rays_Base):
+    def __init__(self, vertices0, faces0):
+        self.vertices0, self.faces0 = vertices0, faces0
+        super().__init__(vertices0=list(vertices0), faces0=list(faces0))
+        
+    def setup_vertices_faces(self):
+        return self.vertices0, self.faces0
+    
+
+class Rays_Cartesian(Rays_Base):
+    def __init__(self, n_rays_x=11, n_rays_z=5):
+        super().__init__(n_rays_x=n_rays_x, n_rays_z=n_rays_z)
+
+    def setup_vertices_faces(self):
+        """has to return list of ( (z_1,y_1,x_1), ... )  _"""
+        n_rays_x, n_rays_z = self.kwargs["n_rays_x"], self.kwargs["n_rays_z"]
+        dphi = np.float32(2. * np.pi / n_rays_x)
+        dtheta = np.float32(np.pi / n_rays_z)
+
+        verts = []
+        for mz in range(n_rays_z):
+            for mx in range(n_rays_x):
+                phi = mx * dphi
+                theta = mz * dtheta
+                if mz == 0:
+                    theta = 1e-12
+                if mz == n_rays_z - 1:
+                    theta = np.pi - 1e-12
+                dx = np.cos(phi) * np.sin(theta)
+                dy = np.sin(phi) * np.sin(theta)
+                dz = np.cos(theta)
+                if mz == 0 or mz == n_rays_z - 1:
+                    dx += 1e-12
+                    dy += 1e-12
+                verts.append([dz, dy, dx])
+
+        verts = np.array(verts)
+
+        def _ind(mz, mx):
+            return mz * n_rays_x + mx
+
+        faces = []
+
+        for mz in range(n_rays_z - 1):
+            for mx in range(n_rays_x):
+                faces.append([_ind(mz, mx), _ind(mz + 1, (mx + 1) % n_rays_x), _ind(mz, (mx + 1) % n_rays_x)])
+                faces.append([_ind(mz, mx), _ind(mz + 1, mx), _ind(mz + 1, (mx + 1) % n_rays_x)])
+
+        faces = np.array(faces)
+
+        return verts, faces
+
+
+class Rays_SubDivide(Rays_Base):
+    """
+    Subdivision polyehdra
+
+    n_level = 1 -> base polyhedra
+    n_level = 2 -> 1x subdivision
+    n_level = 3 -> 2x subdivision
+                ...
+    """
+
+    def __init__(self, n_level=4):
+        super().__init__(n_level=n_level)
+
+    def base_polyhedron(self):
+        raise NotImplementedError()
+
+    def setup_vertices_faces(self):
+        n_level = self.kwargs["n_level"]
+        verts0, faces0 = self.base_polyhedron()
+        return self._recursive_split(verts0, faces0, n_level)
+
+    def _recursive_split(self, verts, faces, n_level):
+        if n_level <= 1:
+            return verts, faces
+        else:
+            verts, faces = Rays_SubDivide.split(verts, faces)
+            return self._recursive_split(verts, faces, n_level - 1)
+
+    @classmethod
+    def split(self, verts0, faces0):
+        """split a level"""
+
+        split_edges = dict()
+        verts = list(verts0[:])
+        faces = []
+
+        def _add(a, b):
+            """ returns index of middle point and adds vertex if not already added"""
+            edge = tuple(sorted((a, b)))
+            if not edge in split_edges:
+                v = .5 * (verts[a] + verts[b])
+                v *= 1. / np.linalg.norm(v)
+                verts.append(v)
+                split_edges[edge] = len(verts) - 1
+            return split_edges[edge]
+
+        for v1, v2, v3 in faces0:
+            ind1 = _add(v1, v2)
+            ind2 = _add(v2, v3)
+            ind3 = _add(v3, v1)
+            faces.append([v1, ind1, ind3])
+            faces.append([v2, ind2, ind1])
+            faces.append([v3, ind3, ind2])
+            faces.append([ind1, ind2, ind3])
+
+        return verts, faces
+
+
+class Rays_Tetra(Rays_SubDivide):
+    """
+    Subdivision of a tetrahedron
+
+    n_level = 1 -> normal tetrahedron (4 vertices)
+    n_level = 2 -> 1x subdivision (10 vertices)
+    n_level = 3 -> 2x subdivision (34 vertices)
+                ...
+    """
+
+    def base_polyhedron(self):
+        verts = np.array([
+            [np.sqrt(8. / 9), 0., -1. / 3],
+            [-np.sqrt(2. / 9), np.sqrt(2. / 3), -1. / 3],
+            [-np.sqrt(2. / 9), -np.sqrt(2. / 3), -1. / 3],
+            [0., 0., 1.]
+        ])
+        faces = [[0, 1, 2],
+                 [0, 3, 1],
+                 [0, 2, 3],
+                 [1, 3, 2]]
+
+        return verts, faces
+
+
+class Rays_Octo(Rays_SubDivide):
+    """
+    Subdivision of a tetrahedron
+
+    n_level = 1 -> normal Octahedron (6 vertices)
+    n_level = 2 -> 1x subdivision (18 vertices)
+    n_level = 3 -> 2x subdivision (66 vertices)
+
+    """
+
+    def base_polyhedron(self):
+        verts = np.array([
+            [0, 0, 1],
+            [0, 1, 0],
+            [0, 0, -1],
+            [0, -1, 0],
+            [1, 0, 0],
+            [-1, 0, 0]])
+
+        faces = [[0, 1, 4],
+                 [0, 5, 1],
+                 [1, 2, 4],
+                 [1, 5, 2],
+                 [2, 3, 4],
+                 [2, 5, 3],
+                 [3, 0, 4],
+                 [3, 5, 0],
+                 ]
+
+        return verts, faces
+
+
+def reorder_faces(verts, faces):
+    """reorder faces such that their orientation points outward"""
+    def _single(face):
+        return face[::-1] if np.linalg.det(verts[face])>0 else face
+    return tuple(map(_single, faces))
+
+
+class Rays_GoldenSpiral(Rays_Base):
+    def __init__(self, n=70, anisotropy = None):
+        if n<4:
+            raise ValueError("At least 4 points have to be given!")
+        super().__init__(n=n, anisotropy = anisotropy if anisotropy is None else tuple(anisotropy))
+
+    def setup_vertices_faces(self):
+        n = self.kwargs["n"]
+        anisotropy = self.kwargs["anisotropy"]
+        if anisotropy is None:
+            anisotropy = np.ones(3)
+        else:
+            anisotropy = np.array(anisotropy)
+
+        # the smaller golden angle = 2pi * 0.3819...
+        g = (3. - np.sqrt(5.)) * np.pi
+        phi = g * np.arange(n)
+        # z = np.linspace(-1, 1, n + 2)[1:-1]
+        # rho = np.sqrt(1. - z ** 2)
+        # verts = np.stack([rho*np.cos(phi), rho*np.sin(phi),z]).T
+        #
+        z = np.linspace(-1, 1, n)
+        rho = np.sqrt(1. - z ** 2)
+        verts = np.stack([z, rho * np.sin(phi), rho * np.cos(phi)]).T
+
+        # warnings.warn("ray definition has changed! Old results are invalid!")
+
+        # correct for anisotropy
+        verts = verts/anisotropy
+        #verts /= np.linalg.norm(verts, axis=-1, keepdims=True)
+
+        hull = ConvexHull(verts)
+        faces = reorder_faces(verts,hull.simplices)
+
+        verts /= np.linalg.norm(verts, axis=-1, keepdims=True)
+
+        return verts, faces

stardist_pkg/sample_patches.py

@@ -0,0 +1,65 @@
+"""provides a faster sampling function"""
+
+import numpy as np
+from csbdeep.utils import _raise, choice
+
+
+def sample_patches(datas, patch_size, n_samples, valid_inds=None, verbose=False):
+    """optimized version of csbdeep.data.sample_patches_from_multiple_stacks
+    """
+
+    len(patch_size)==datas[0].ndim or _raise(ValueError())
+
+    if not all(( a.shape == datas[0].shape for a in datas )):
+        raise ValueError("all input shapes must be the same: %s" % (" / ".join(str(a.shape) for a in datas)))
+
+    if not all(( 0 < s <= d for s,d in zip(patch_size,datas[0].shape) )):
+        raise ValueError("patch_size %s negative or larger than data shape %s along some dimensions" % (str(patch_size), str(datas[0].shape)))
+
+    if valid_inds is None:
+        valid_inds = tuple(_s.ravel() for _s in np.meshgrid(*tuple(np.arange(p//2,s-p//2+1) for s,p in zip(datas[0].shape, patch_size))))
+
+    n_valid = len(valid_inds[0])
+
+    if n_valid == 0:
+        raise ValueError("no regions to sample from!")
+
+    idx = choice(range(n_valid), n_samples, replace=(n_valid < n_samples))
+    rand_inds = [v[idx] for v in valid_inds]
+    res = [np.stack([data[tuple(slice(_r-(_p//2),_r+_p-(_p//2)) for _r,_p in zip(r,patch_size))] for r in zip(*rand_inds)]) for data in datas]
+
+    return res
+
+
+def get_valid_inds(img, patch_size, patch_filter=None):
+    """
+    Returns all indices of an image that 
+    - can be used as center points for sampling patches of a given patch_size, and
+    - are part of the boolean mask given by the function patch_filter (if provided)
+
+    img: np.ndarray
+    patch_size: tuple of ints 
+        the width of patches per img dimension, 
+    patch_filter: None or callable
+        a function with signature patch_filter(img, patch_size) returning a boolean mask 
+    """
+
+    len(patch_size)==img.ndim or _raise(ValueError())
+
+    if not all(( 0 < s <= d for s,d in zip(patch_size,img.shape))):
+        raise ValueError("patch_size %s negative or larger than image shape %s along some dimensions" % (str(patch_size), str(img.shape)))
+
+    if patch_filter is None:
+        # only cut border indices (which is faster)
+        patch_mask = np.ones(img.shape,dtype=bool)
+        valid_inds = tuple(np.arange(p // 2, s - p + p // 2 + 1).astype(np.uint32) for p, s in zip(patch_size, img.shape))
+        valid_inds = tuple(s.ravel() for s in np.meshgrid(*valid_inds, indexing='ij'))
+    else:
+        patch_mask = patch_filter(img, patch_size)
+
+        # get the valid indices
+        border_slices = tuple([slice(p // 2, s - p + p // 2 + 1) for p, s in zip(patch_size, img.shape)])
+        valid_inds = np.where(patch_mask[border_slices])
+        valid_inds = tuple((v + s.start).astype(np.uint32) for s, v in zip(border_slices, valid_inds))
+
+    return valid_inds

stardist_pkg/utils.py

@@ -0,0 +1,394 @@
+from __future__ import print_function, unicode_literals, absolute_import, division
+
+import numpy as np
+import warnings
+import os
+import datetime
+from tqdm import tqdm
+from collections import defaultdict
+from zipfile import ZipFile, ZIP_DEFLATED
+from scipy.ndimage.morphology import distance_transform_edt, binary_fill_holes
+from scipy.ndimage.measurements import find_objects
+from scipy.optimize import minimize_scalar
+from skimage.measure import regionprops
+from csbdeep.utils import _raise
+from csbdeep.utils.six import Path
+from collections.abc import Iterable
+
+from .matching import matching_dataset, _check_label_array
+
+
+try:
+    from edt import edt
+    _edt_available = True
+    try:    _edt_parallel_max = len(os.sched_getaffinity(0))
+    except: _edt_parallel_max = 128
+    _edt_parallel_default = 4
+    _edt_parallel = os.environ.get('STARDIST_EDT_NUM_THREADS', _edt_parallel_default)
+    try:
+        _edt_parallel = min(_edt_parallel_max, int(_edt_parallel))
+    except ValueError as e:
+        warnings.warn(f"Invalid value ({_edt_parallel}) for STARDIST_EDT_NUM_THREADS. Using default value ({_edt_parallel_default}) instead.")
+        _edt_parallel = _edt_parallel_default
+    del _edt_parallel_default, _edt_parallel_max
+except ImportError:
+    _edt_available = False
+    # warnings.warn("Could not find package edt... \nConsider installing it with \n  pip install edt\nto improve training data generation performance.")
+    pass
+
+
+def gputools_available():
+    try:
+        import gputools
+    except:
+        return False
+    return True
+
+
+def path_absolute(path_relative):
+    """ Get absolute path to resource"""
+    base_path = os.path.abspath(os.path.dirname(__file__))
+    return os.path.join(base_path, path_relative)
+
+
+def _is_power_of_2(i):
+    assert i > 0
+    e = np.log2(i)
+    return e == int(e)
+
+
+def _normalize_grid(grid,n):
+    try:
+        grid = tuple(grid)
+        (len(grid) == n and
+         all(map(np.isscalar,grid)) and
+         all(map(_is_power_of_2,grid))) or _raise(TypeError())
+        return tuple(int(g) for g in grid)
+    except (TypeError, AssertionError):
+        raise ValueError("grid = {grid} must be a list/tuple of length {n} with values that are power of 2".format(grid=grid, n=n))
+
+
+def edt_prob(lbl_img, anisotropy=None):
+    if _edt_available:
+        return _edt_prob_edt(lbl_img, anisotropy=anisotropy)
+    else:
+        # warnings.warn("Could not find package edt... \nConsider installing it with \n  pip install edt\nto improve training data generation performance.")
+        return _edt_prob_scipy(lbl_img, anisotropy=anisotropy)
+
+def _edt_prob_edt(lbl_img, anisotropy=None):
+    """Perform EDT on each labeled object and normalize.
+    Internally uses https://github.com/seung-lab/euclidean-distance-transform-3d
+    that can handle multiple labels at once
+    """
+    lbl_img = np.ascontiguousarray(lbl_img)
+    constant_img = lbl_img.min() == lbl_img.max() and lbl_img.flat[0] > 0
+    if constant_img:
+        warnings.warn("EDT of constant label image is ill-defined. (Assuming background around it.)")
+    # we just need to compute the edt once but then normalize it for each object
+    prob = edt(lbl_img, anisotropy=anisotropy, black_border=constant_img, parallel=_edt_parallel)
+    objects = find_objects(lbl_img)
+    for i,sl in enumerate(objects,1):
+        # i: object label id, sl: slices of object in lbl_img
+        if sl is None: continue
+        _mask = lbl_img[sl]==i
+        # normalize it
+        prob[sl][_mask] /= np.max(prob[sl][_mask]+1e-10)
+    return prob
+
+def _edt_prob_scipy(lbl_img, anisotropy=None):
+    """Perform EDT on each labeled object and normalize."""
+    def grow(sl,interior):
+        return tuple(slice(s.start-int(w[0]),s.stop+int(w[1])) for s,w in zip(sl,interior))
+    def shrink(interior):
+        return tuple(slice(int(w[0]),(-1 if w[1] else None)) for w in interior)
+    constant_img = lbl_img.min() == lbl_img.max() and lbl_img.flat[0] > 0
+    if constant_img:
+        lbl_img = np.pad(lbl_img, ((1,1),)*lbl_img.ndim, mode='constant')
+        warnings.warn("EDT of constant label image is ill-defined. (Assuming background around it.)")
+    objects = find_objects(lbl_img)
+    prob = np.zeros(lbl_img.shape,np.float32)
+    for i,sl in enumerate(objects,1):
+        # i: object label id, sl: slices of object in lbl_img
+        if sl is None: continue
+        interior = [(s.start>0,s.stop<sz) for s,sz in zip(sl,lbl_img.shape)]
+        # 1. grow object slice by 1 for all interior object bounding boxes
+        # 2. perform (correct) EDT for object with label id i
+        # 3. extract EDT for object of original slice and normalize
+        # 4. store edt for object only for pixels of given label id i
+        shrink_slice = shrink(interior)
+        grown_mask = lbl_img[grow(sl,interior)]==i
+        mask = grown_mask[shrink_slice]
+        edt = distance_transform_edt(grown_mask, sampling=anisotropy)[shrink_slice][mask]
+        prob[sl][mask] = edt/(np.max(edt)+1e-10)
+    if constant_img:
+        prob = prob[(slice(1,-1),)*lbl_img.ndim].copy()
+    return prob
+
+
+def _fill_label_holes(lbl_img, **kwargs):
+    lbl_img_filled = np.zeros_like(lbl_img)
+    for l in (set(np.unique(lbl_img)) - set([0])):
+        mask = lbl_img==l
+        mask_filled = binary_fill_holes(mask,**kwargs)
+        lbl_img_filled[mask_filled] = l
+    return lbl_img_filled
+
+
+def fill_label_holes(lbl_img, **kwargs):
+    """Fill small holes in label image."""
+    # TODO: refactor 'fill_label_holes' and 'edt_prob' to share code
+    def grow(sl,interior):
+        return tuple(slice(s.start-int(w[0]),s.stop+int(w[1])) for s,w in zip(sl,interior))
+    def shrink(interior):
+        return tuple(slice(int(w[0]),(-1 if w[1] else None)) for w in interior)
+    objects = find_objects(lbl_img)
+    lbl_img_filled = np.zeros_like(lbl_img)
+    for i,sl in enumerate(objects,1):
+        if sl is None: continue
+        interior = [(s.start>0,s.stop<sz) for s,sz in zip(sl,lbl_img.shape)]
+        shrink_slice = shrink(interior)
+        grown_mask = lbl_img[grow(sl,interior)]==i
+        mask_filled = binary_fill_holes(grown_mask,**kwargs)[shrink_slice]
+        lbl_img_filled[sl][mask_filled] = i
+    return lbl_img_filled
+
+
+def sample_points(n_samples, mask, prob=None, b=2):
+    """sample points to draw some of the associated polygons"""
+    if b is not None and b > 0:
+        # ignore image boundary, since predictions may not be reliable
+        mask_b = np.zeros_like(mask)
+        mask_b[b:-b,b:-b] = True
+    else:
+        mask_b = True
+
+    points = np.nonzero(mask & mask_b)
+
+    if prob is not None:
+        # weighted sampling via prob
+        w = prob[points[0],points[1]].astype(np.float64)
+        w /= np.sum(w)
+        ind = np.random.choice(len(points[0]), n_samples, replace=True, p=w)
+    else:
+        ind = np.random.choice(len(points[0]), n_samples, replace=True)
+
+    points = points[0][ind], points[1][ind]
+    points = np.stack(points,axis=-1)
+    return points
+
+
+def calculate_extents(lbl, func=np.median):
+    """ Aggregate bounding box sizes of objects in label images. """
+    if (isinstance(lbl,np.ndarray) and lbl.ndim==4) or (not isinstance(lbl,np.ndarray) and  isinstance(lbl,Iterable)):
+        return func(np.stack([calculate_extents(_lbl,func) for _lbl in lbl], axis=0), axis=0)
+
+    n = lbl.ndim
+    n in (2,3) or _raise(ValueError("label image should be 2- or 3-dimensional (or pass a list of these)"))
+
+    regs = regionprops(lbl)
+    if len(regs) == 0:
+        return np.zeros(n)
+    else:
+        extents = np.array([np.array(r.bbox[n:])-np.array(r.bbox[:n]) for r in regs])
+        return func(extents, axis=0)
+
+
+def polyroi_bytearray(x,y,pos=None,subpixel=True):
+    """ Byte array of polygon roi with provided x and y coordinates
+        See https://github.com/imagej/imagej1/blob/master/ij/io/RoiDecoder.java
+    """
+    import struct
+    def _int16(x):
+        return int(x).to_bytes(2, byteorder='big', signed=True)
+    def _uint16(x):
+        return int(x).to_bytes(2, byteorder='big', signed=False)
+    def _int32(x):
+        return int(x).to_bytes(4, byteorder='big', signed=True)
+    def _float(x):
+        return struct.pack(">f", x)
+
+    subpixel = bool(subpixel)
+    # add offset since pixel center is at (0.5,0.5) in ImageJ
+    x_raw = np.asarray(x).ravel() + 0.5
+    y_raw = np.asarray(y).ravel() + 0.5
+    x = np.round(x_raw)
+    y = np.round(y_raw)
+    assert len(x) == len(y)
+    top, left, bottom, right = y.min(), x.min(), y.max(), x.max() # bbox
+
+    n_coords = len(x)
+    bytes_header = 64
+    bytes_total = bytes_header + n_coords*2*2 + subpixel*n_coords*2*4
+    B = [0] * bytes_total
+    B[ 0: 4] = map(ord,'Iout')   # magic start
+    B[ 4: 6] = _int16(227)       # version
+    B[ 6: 8] = _int16(0)         # roi type (0 = polygon)
+    B[ 8:10] = _int16(top)       # bbox top
+    B[10:12] = _int16(left)      # bbox left
+    B[12:14] = _int16(bottom)    # bbox bottom
+    B[14:16] = _int16(right)     # bbox right
+    B[16:18] = _uint16(n_coords) # number of coordinates
+    if subpixel:
+        B[50:52] = _int16(128)   # subpixel resolution (option flag)
+    if pos is not None:
+        B[56:60] = _int32(pos)   # position (C, Z, or T)
+
+    for i,(_x,_y) in enumerate(zip(x,y)):
+        xs = bytes_header + 2*i
+        ys = xs + 2*n_coords
+        B[xs:xs+2] = _int16(_x - left)
+        B[ys:ys+2] = _int16(_y - top)
+
+    if subpixel:
+        base1 = bytes_header + n_coords*2*2
+        base2 = base1 + n_coords*4
+        for i,(_x,_y) in enumerate(zip(x_raw,y_raw)):
+            xs = base1 + 4*i
+            ys = base2 + 4*i
+            B[xs:xs+4] = _float(_x)
+            B[ys:ys+4] = _float(_y)
+
+    return bytearray(B)
+
+
+def export_imagej_rois(fname, polygons, set_position=True, subpixel=True, compression=ZIP_DEFLATED):
+    """ polygons assumed to be a list of arrays with shape (id,2,c) """
+
+    if isinstance(polygons,np.ndarray):
+        polygons = (polygons,)
+
+    fname = Path(fname)
+    if fname.suffix == '.zip':
+        fname = fname.with_suffix('')
+
+    with ZipFile(str(fname)+'.zip', mode='w', compression=compression) as roizip:
+        for pos,polygroup in enumerate(polygons,start=1):
+            for i,poly in enumerate(polygroup,start=1):
+                roi = polyroi_bytearray(poly[1],poly[0], pos=(pos if set_position else None), subpixel=subpixel)
+                roizip.writestr('{pos:03d}_{i:03d}.roi'.format(pos=pos,i=i), roi)
+
+
+def optimize_threshold(Y, Yhat, model, nms_thresh, measure='accuracy', iou_threshs=[0.3,0.5,0.7], bracket=None, tol=1e-2, maxiter=20, verbose=1):
+    """ Tune prob_thresh for provided (fixed) nms_thresh to maximize matching score (for given measure and averaged over iou_threshs). """
+    np.isscalar(nms_thresh) or _raise(ValueError("nms_thresh must be a scalar"))
+    iou_threshs = [iou_threshs] if np.isscalar(iou_threshs) else iou_threshs
+    values = dict()
+
+    if bracket is None:
+        max_prob = max([np.max(prob) for prob, dist in Yhat])
+        bracket = max_prob/2, max_prob
+    # print("bracket =", bracket)
+
+    with tqdm(total=maxiter, disable=(verbose!=1), desc="NMS threshold = %g" % nms_thresh) as progress:
+
+        def fn(thr):
+            prob_thresh = np.clip(thr, *bracket)
+            value = values.get(prob_thresh)
+            if value is None:
+                Y_instances = [model._instances_from_prediction(y.shape, *prob_dist, prob_thresh=prob_thresh, nms_thresh=nms_thresh)[0] for y,prob_dist in zip(Y,Yhat)]
+                stats = matching_dataset(Y, Y_instances, thresh=iou_threshs, show_progress=False, parallel=True)
+                values[prob_thresh] = value = np.mean([s._asdict()[measure] for s in stats])
+            if verbose > 1:
+                print("{now}   thresh: {prob_thresh:f}   {measure}: {value:f}".format(
+                    now = datetime.datetime.now().strftime('%H:%M:%S'),
+                    prob_thresh = prob_thresh,
+                    measure = measure,
+                    value = value,
+                ), flush=True)
+            else:
+                progress.update()
+                progress.set_postfix_str("{prob_thresh:.3f} -> {value:.3f}".format(prob_thresh=prob_thresh, value=value))
+                progress.refresh()
+            return -value
+
+        opt = minimize_scalar(fn, method='golden', bracket=bracket, tol=tol, options={'maxiter': maxiter})
+
+    verbose > 1 and print('\n',opt, flush=True)
+    return opt.x, -opt.fun
+
+
+def _invert_dict(d):
+    """ return  v-> [k_1,k_2,k_3....] for k,v in d"""
+    res = defaultdict(list)
+    for k,v in d.items():
+        res[v].append(k)
+    return res
+
+
+def mask_to_categorical(y, n_classes, classes, return_cls_dict=False):
+    """generates a multi-channel categorical class map
+
+    Parameters
+    ----------
+    y : n-dimensional ndarray
+        integer label array
+    n_classes : int
+        Number of different classes (without background)
+    classes: dict, integer, or None
+        the label to class assignment
+        can be
+        - dict {label -> class_id}
+           the value of class_id can be
+                             0   -> background class
+                  1...n_classes  -> the respective object class (1 ... n_classes)
+                           None  -> ignore object (prob is set to -1 for the pixels of the object, except for background class)
+        - single integer value or None -> broadcast value to all labels
+
+    Returns
+    -------
+    probability map of shape y.shape+(n_classes+1,) (first channel is background)
+
+    """
+
+    _check_label_array(y, 'y')
+    if not (np.issubdtype(type(n_classes), np.integer) and n_classes>=1):
+        raise ValueError(f"n_classes is '{n_classes}' but should be a positive integer")
+
+    y_labels = np.unique(y[y>0]).tolist()
+
+    # build dict class_id -> labels (inverse of classes)
+    if np.issubdtype(type(classes), np.integer) or classes is None:
+        classes = dict((k,classes) for k in y_labels)
+    elif isinstance(classes, dict):
+        pass
+    else:
+        raise ValueError("classes should be dict, single scalar, or None!")
+
+    if not set(y_labels).issubset(set(classes.keys())):
+        raise ValueError(f"all gt labels should be present in class dict provided \ngt_labels found\n{set(y_labels)}\nclass dict labels provided\n{set(classes.keys())}")
+
+    cls_dict = _invert_dict(classes)
+
+    # prob map
+    y_mask = np.zeros(y.shape+(n_classes+1,), np.float32)
+
+    for cls, labels in cls_dict.items():
+        if cls is None:
+            # prob == -1 will be used in the loss to ignore object
+            y_mask[np.isin(y, labels)] = -1
+        elif np.issubdtype(type(cls), np.integer) and 0 <= cls <= n_classes:
+            y_mask[...,cls] = np.isin(y, labels)
+        else:
+            raise ValueError(f"Wrong class id '{cls}' (for n_classes={n_classes})")
+
+    # set 0/1 background prob (unaffected by None values for class ids)
+    y_mask[...,0] = (y==0)
+
+    if return_cls_dict:
+        return y_mask, cls_dict
+    else:
+        return y_mask
+
+
+def _is_floatarray(x):
+    return isinstance(x.dtype.type(0),np.floating)
+
+
+def abspath(root, relpath):
+    from pathlib import Path
+    root = Path(root)
+    if root.is_dir():
+        path = root/relpath
+    else:
+        path = root.parent/relpath
+    return str(path.absolute())

stardist_pkg/version.py

@@ -0,0 +1 @@
+__version__ = '0.8.3'

utils_modify.py

@@ -0,0 +1,743 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+from typing import Any, Callable, Dict, List, Mapping, Sequence, Tuple, Union
+import numpy as np
+import torch
+import torch.nn.functional as F
+from stardist_pkg.big import _grid_divisible, BlockND, OBJECT_KEYS#, repaint_labels
+from stardist_pkg.matching import relabel_sequential
+from stardist_pkg import dist_to_coord, non_maximum_suppression, polygons_to_label
+#from stardist_pkg import dist_to_coord, polygons_to_label
+from stardist_pkg import star_dist,edt_prob
+from monai.data.meta_tensor import MetaTensor
+from monai.data.utils import compute_importance_map, dense_patch_slices, get_valid_patch_size
+from monai.transforms import Resize
+from monai.utils import (
+    BlendMode,
+    PytorchPadMode,
+    convert_data_type,
+    convert_to_dst_type,
+    ensure_tuple,
+    fall_back_tuple,
+    look_up_option,
+    optional_import,
+)
+import cv2
+from scipy import ndimage
+from scipy.ndimage.filters import gaussian_filter
+from scipy.ndimage.interpolation import affine_transform, map_coordinates
+from skimage import morphology as morph
+from scipy.ndimage import filters, measurements
+from scipy.ndimage.morphology import (
+    binary_dilation,
+    binary_fill_holes,
+    distance_transform_cdt,
+    distance_transform_edt,
+)
+
+from skimage.segmentation import watershed
+tqdm, _ = optional_import("tqdm", name="tqdm")
+
+__all__ = ["sliding_window_inference"]
+
+
+####
+def normalize(mask, dtype=np.uint8):
+    return (255 * mask / np.amax(mask)).astype(dtype)
+
+def fix_mirror_padding(ann):
+    """Deal with duplicated instances due to mirroring in interpolation
+    during shape augmentation (scale, rotation etc.).
+
+    """
+    current_max_id = np.amax(ann)
+    inst_list = list(np.unique(ann))
+    if 0 in inst_list:
+        inst_list.remove(0)  # 0 is background
+    for inst_id in inst_list:
+        inst_map = np.array(ann == inst_id, np.uint8)
+        remapped_ids = measurements.label(inst_map)[0]
+        remapped_ids[remapped_ids > 1] += current_max_id
+        ann[remapped_ids > 1] = remapped_ids[remapped_ids > 1]
+        current_max_id = np.amax(ann)
+    return ann
+
+####
+def get_bounding_box(img):
+    """Get bounding box coordinate information."""
+    rows = np.any(img, axis=1)
+    cols = np.any(img, axis=0)
+    rmin, rmax = np.where(rows)[0][[0, -1]]
+    cmin, cmax = np.where(cols)[0][[0, -1]]
+    # due to python indexing, need to add 1 to max
+    # else accessing will be 1px in the box, not out
+    rmax += 1
+    cmax += 1
+    return [rmin, rmax, cmin, cmax]
+
+
+####
+def cropping_center(x, crop_shape, batch=False):
+    """Crop an input image at the centre.
+
+    Args:
+        x: input array
+        crop_shape: dimensions of cropped array
+    
+    Returns:
+        x: cropped array
+    
+    """
+    orig_shape = x.shape
+    if not batch:
+        h0 = int((orig_shape[0] - crop_shape[0]) * 0.5)
+        w0 = int((orig_shape[1] - crop_shape[1]) * 0.5)
+        x = x[h0 : h0 + crop_shape[0], w0 : w0 + crop_shape[1]]
+    else:
+        h0 = int((orig_shape[1] - crop_shape[0]) * 0.5)
+        w0 = int((orig_shape[2] - crop_shape[1]) * 0.5)
+        x = x[:, h0 : h0 + crop_shape[0], w0 : w0 + crop_shape[1]]
+    return x
+
+def gen_instance_hv_map(ann, crop_shape):
+    """Input annotation must be of original shape.
+    
+    The map is calculated only for instances within the crop portion
+    but based on the original shape in original image.
+
+    Perform following operation:
+    Obtain the horizontal and vertical distance maps for each
+    nuclear instance.
+
+    """
+    orig_ann = ann.copy()  # instance ID map
+    fixed_ann = fix_mirror_padding(orig_ann)
+    # re-cropping with fixed instance id map
+    crop_ann = cropping_center(fixed_ann, crop_shape)
+    # TODO: deal with 1 label warning
+    crop_ann = morph.remove_small_objects(crop_ann, min_size=30)
+
+    x_map = np.zeros(orig_ann.shape[:2], dtype=np.float32)
+    y_map = np.zeros(orig_ann.shape[:2], dtype=np.float32)
+
+    inst_list = list(np.unique(crop_ann))
+    if 0 in inst_list:
+        inst_list.remove(0)  # 0 is background
+    for inst_id in inst_list:
+        inst_map = np.array(fixed_ann == inst_id, np.uint8)
+        inst_box = get_bounding_box(inst_map) # rmin, rmax, cmin, cmax
+
+        # expand the box by 2px
+        # Because we first pad the ann at line 207, the bboxes
+        # will remain valid after expansion
+        inst_box[0] -= 2
+        inst_box[2] -= 2
+        inst_box[1] += 2
+        inst_box[3] += 2
+        
+        # fix inst_box
+        inst_box[0] = max(inst_box[0], 0)
+        inst_box[2] = max(inst_box[2], 0)
+        # inst_box[1] = min(inst_box[1], fixed_ann.shape[0]) 
+        # inst_box[3] = min(inst_box[3], fixed_ann.shape[1])
+
+        inst_map = inst_map[inst_box[0] : inst_box[1], inst_box[2] : inst_box[3]]
+
+        if inst_map.shape[0] < 2 or inst_map.shape[1] < 2:
+            print(f'inst_map.shape < 2: {inst_map.shape}, {inst_box}, {get_bounding_box(np.array(fixed_ann == inst_id, np.uint8))}')
+            continue
+
+        # instance center of mass, rounded to nearest pixel
+        inst_com = list(measurements.center_of_mass(inst_map))
+        if np.isnan(measurements.center_of_mass(inst_map)).any():
+            print(inst_id, fixed_ann.shape, np.array(fixed_ann == inst_id, np.uint8).shape)
+            print(get_bounding_box(np.array(fixed_ann == inst_id, np.uint8)))
+            print(inst_map)
+            print(inst_list)
+            print(inst_box)
+            print(np.count_nonzero(np.array(fixed_ann == inst_id, np.uint8)))
+
+        inst_com[0] = int(inst_com[0] + 0.5)
+        inst_com[1] = int(inst_com[1] + 0.5)
+
+        inst_x_range = np.arange(1, inst_map.shape[1] + 1)
+        inst_y_range = np.arange(1, inst_map.shape[0] + 1)
+        # shifting center of pixels grid to instance center of mass
+        inst_x_range -= inst_com[1]
+        inst_y_range -= inst_com[0]
+
+        inst_x, inst_y = np.meshgrid(inst_x_range, inst_y_range)
+
+        # remove coord outside of instance
+        inst_x[inst_map == 0] = 0
+        inst_y[inst_map == 0] = 0
+        inst_x = inst_x.astype("float32")
+        inst_y = inst_y.astype("float32")
+
+        # normalize min into -1 scale
+        if np.min(inst_x) < 0:
+            inst_x[inst_x < 0] /= -np.amin(inst_x[inst_x < 0])
+        if np.min(inst_y) < 0:
+            inst_y[inst_y < 0] /= -np.amin(inst_y[inst_y < 0])
+        # normalize max into +1 scale
+        if np.max(inst_x) > 0:
+            inst_x[inst_x > 0] /= np.amax(inst_x[inst_x > 0])
+        if np.max(inst_y) > 0:
+            inst_y[inst_y > 0] /= np.amax(inst_y[inst_y > 0])
+
+        ####
+        x_map_box = x_map[inst_box[0] : inst_box[1], inst_box[2] : inst_box[3]]
+        x_map_box[inst_map > 0] = inst_x[inst_map > 0]
+
+        y_map_box = y_map[inst_box[0] : inst_box[1], inst_box[2] : inst_box[3]]
+        y_map_box[inst_map > 0] = inst_y[inst_map > 0]
+
+    hv_map = np.dstack([x_map, y_map])
+    return hv_map
+
+def remove_small_objects(pred, min_size=64, connectivity=1):
+    """Remove connected components smaller than the specified size.
+
+    This function is taken from skimage.morphology.remove_small_objects, but the warning
+    is removed when a single label is provided.
+
+    Args:
+        pred: input labelled array
+        min_size: minimum size of instance in output array
+        connectivity: The connectivity defining the neighborhood of a pixel.
+
+    Returns:
+        out: output array with instances removed under min_size
+
+    """
+    out = pred
+
+    if min_size == 0:  # shortcut for efficiency
+        return out
+
+    if out.dtype == bool:
+        selem = ndimage.generate_binary_structure(pred.ndim, connectivity)
+        ccs = np.zeros_like(pred, dtype=np.int32)
+        ndimage.label(pred, selem, output=ccs)
+    else:
+        ccs = out
+
+    try:
+        component_sizes = np.bincount(ccs.ravel())
+    except ValueError:
+        raise ValueError(
+            "Negative value labels are not supported. Try "
+            "relabeling the input with `scipy.ndimage.label` or "
+            "`skimage.morphology.label`."
+        )
+
+    too_small = component_sizes < min_size
+    too_small_mask = too_small[ccs]
+    out[too_small_mask] = 0
+
+    return out
+
+####
+def gen_targets(ann, crop_shape, **kwargs):
+    """Generate the targets for the network."""
+    hv_map = gen_instance_hv_map(ann, crop_shape)
+    np_map = ann.copy()
+    np_map[np_map > 0] = 1
+
+    hv_map = cropping_center(hv_map, crop_shape)
+    np_map = cropping_center(np_map, crop_shape)
+
+    target_dict = {
+        "hv_map": hv_map,
+        "np_map": np_map,
+    }
+
+    return target_dict
+def __proc_np_hv(pred, np_thres=0.5, ksize=21, overall_thres=0.4, obj_size_thres=10):
+    """Process Nuclei Prediction with XY Coordinate Map.
+
+    Args:
+        pred: prediction output, assuming
+              channel 0 contain probability map of nuclei
+              channel 1 containing the regressed X-map
+              channel 2 containing the regressed Y-map
+
+    """
+    pred = np.array(pred, dtype=np.float32)
+
+    blb_raw = pred[..., 0]
+    h_dir_raw = pred[..., 1]
+    v_dir_raw = pred[..., 2]
+
+    # processing
+    blb = np.array(blb_raw >= np_thres, dtype=np.int32)
+
+    blb = measurements.label(blb)[0]
+    blb = remove_small_objects(blb, min_size=10)
+    blb[blb > 0] = 1  # background is 0 already
+
+    h_dir = cv2.normalize(
+        h_dir_raw, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F
+    )
+    v_dir = cv2.normalize(
+        v_dir_raw, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F
+    )
+
+    sobelh = cv2.Sobel(h_dir, cv2.CV_64F, 1, 0, ksize=ksize)
+    sobelv = cv2.Sobel(v_dir, cv2.CV_64F, 0, 1, ksize=ksize)
+
+    sobelh = 1 - (
+        cv2.normalize(
+            sobelh, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F
+        )
+    )
+    sobelv = 1 - (
+        cv2.normalize(
+            sobelv, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F
+        )
+    )
+
+    overall = np.maximum(sobelh, sobelv)
+    overall = overall - (1 - blb)
+    overall[overall < 0] = 0
+
+    dist = (1.0 - overall) * blb
+    ## nuclei values form mountains so inverse to get basins
+    dist = -cv2.GaussianBlur(dist, (3, 3), 0)
+
+    overall = np.array(overall >= overall_thres, dtype=np.int32)
+
+    marker = blb - overall
+    marker[marker < 0] = 0
+    marker = binary_fill_holes(marker).astype("uint8")
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    marker = cv2.morphologyEx(marker, cv2.MORPH_OPEN, kernel)
+    marker = measurements.label(marker)[0]
+    marker = remove_small_objects(marker, min_size=obj_size_thres)
+
+    proced_pred = watershed(dist, markers=marker, mask=blb)
+
+    return proced_pred
+
+####
+def colorize(ch, vmin, vmax):
+    """Will clamp value value outside the provided range to vmax and vmin."""
+    cmap = plt.get_cmap("jet")
+    ch = np.squeeze(ch.astype("float32"))
+    vmin = vmin if vmin is not None else ch.min()
+    vmax = vmax if vmax is not None else ch.max()
+    ch[ch > vmax] = vmax  # clamp value
+    ch[ch < vmin] = vmin
+    ch = (ch - vmin) / (vmax - vmin + 1.0e-16)
+    # take RGB from RGBA heat map
+    ch_cmap = (cmap(ch)[..., :3] * 255).astype("uint8")
+    return ch_cmap
+
+
+####
+def random_colors(N, bright=True):
+    """Generate random colors.
+
+    To get visually distinct colors, generate them in HSV space then
+    convert to RGB.
+    """
+    brightness = 1.0 if bright else 0.7
+    hsv = [(i / N, 1, brightness) for i in range(N)]
+    colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv))
+    random.shuffle(colors)
+    return colors
+
+
+####
+def visualize_instances_map(
+    input_image, inst_map, type_map=None, type_colour=None, line_thickness=2
+):
+    """Overlays segmentation results on image as contours.
+
+    Args:
+        input_image: input image
+        inst_map: instance mask with unique value for every object
+        type_map: type mask with unique value for every class
+        type_colour: a dict of {type : colour} , `type` is from 0-N
+                     and `colour` is a tuple of (R, G, B)
+        line_thickness: line thickness of contours
+
+    Returns:
+        overlay: output image with segmentation overlay as contours
+    """
+    overlay = np.copy((input_image).astype(np.uint8))
+
+    inst_list = list(np.unique(inst_map))  # get list of instances
+    inst_list.remove(0)  # remove background
+
+    inst_rng_colors = random_colors(len(inst_list))
+    inst_rng_colors = np.array(inst_rng_colors) * 255
+    inst_rng_colors = inst_rng_colors.astype(np.uint8)
+
+    for inst_idx, inst_id in enumerate(inst_list):
+        inst_map_mask = np.array(inst_map == inst_id, np.uint8)  # get single object
+        y1, y2, x1, x2 = get_bounding_box(inst_map_mask)
+        y1 = y1 - 2 if y1 - 2 >= 0 else y1
+        x1 = x1 - 2 if x1 - 2 >= 0 else x1
+        x2 = x2 + 2 if x2 + 2 <= inst_map.shape[1] - 1 else x2
+        y2 = y2 + 2 if y2 + 2 <= inst_map.shape[0] - 1 else y2
+        inst_map_crop = inst_map_mask[y1:y2, x1:x2]
+        contours_crop = cv2.findContours(
+            inst_map_crop, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+        )
+        # only has 1 instance per map, no need to check #contour detected by opencv
+        contours_crop = np.squeeze(
+            contours_crop[0][0].astype("int32")
+        )  # * opencv protocol format may break
+        contours_crop += np.asarray([[x1, y1]])  # index correction
+        if type_map is not None:
+            type_map_crop = type_map[y1:y2, x1:x2]
+            type_id = np.unique(type_map_crop).max()  # non-zero
+            inst_colour = type_colour[type_id]
+        else:
+            inst_colour = (inst_rng_colors[inst_idx]).tolist()
+        cv2.drawContours(overlay, [contours_crop], -1, inst_colour, line_thickness)
+    return overlay
+
+
+def sliding_window_inference_large(inputs,block_size,min_overlap,context,roi_size,sw_batch_size,predictor,device):
+    
+    h,w = inputs.shape[0],inputs.shape[1]
+    if h < 5000 or w < 5000:
+        test_tensor = torch.from_numpy(np.expand_dims(inputs, 0)).permute(0,3,1,2).type(torch.FloatTensor).to(device)
+        output_dist,output_prob = sliding_window_inference(test_tensor, roi_size, sw_batch_size, predictor)          
+        prob = output_prob[0][0].cpu().numpy()
+        dist = output_dist[0].cpu().numpy()
+        dist = np.transpose(dist,(1,2,0))
+        dist = np.maximum(1e-3, dist)
+        if h*w < 1500*1500:
+            points, probi, disti = non_maximum_suppression(dist,prob,prob_thresh=0.55, nms_thresh=0.4,cut=True)
+        else:
+            points, probi, disti = non_maximum_suppression(dist,prob,prob_thresh=0.5, nms_thresh=0.4)
+
+            
+        labels_out = polygons_to_label(disti, points, prob=probi,shape=prob.shape)
+    else:
+        n = inputs.ndim
+        axes = 'YXC'
+        grid = (1,1,1)
+        if np.isscalar(block_size):  block_size  = n*[block_size]
+        if np.isscalar(min_overlap): min_overlap = n*[min_overlap]
+        if np.isscalar(context):     context     = n*[context]
+        shape_out = (inputs.shape[0],inputs.shape[1])
+        labels_out = np.zeros(shape_out, dtype=np.uint64)
+    #print(inputs.dtype)
+        block_size[2] = inputs.shape[2]
+        min_overlap[2] = context[2] = 0
+        block_size  = tuple(_grid_divisible(g, v, name='block_size',  verbose=False) for v,g,a in zip(block_size, grid,axes))
+        min_overlap = tuple(_grid_divisible(g, v, name='min_overlap', verbose=False) for v,g,a in zip(min_overlap,grid,axes))
+        context     = tuple(_grid_divisible(g, v, name='context',     verbose=False) for v,g,a in zip(context,    grid,axes))
+        print(f'effective: block_size={block_size}, min_overlap={min_overlap}, context={context}', flush=True)
+        blocks = BlockND.cover(inputs.shape, axes, block_size, min_overlap, context)
+        label_offset = 1
+        blocks = tqdm(blocks)
+        for block in blocks:
+            image = block.read(inputs, axes=axes)
+            test_tensor = torch.from_numpy(np.expand_dims(image, 0)).permute(0,3,1,2).type(torch.FloatTensor).to(device)
+            output_dist,output_prob = sliding_window_inference(test_tensor, roi_size, sw_batch_size, predictor)
+            prob = output_prob[0][0].cpu().numpy()
+            dist = output_dist[0].cpu().numpy()
+            dist = np.transpose(dist,(1,2,0))
+            dist = np.maximum(1e-3, dist)
+            points, probi, disti = non_maximum_suppression(dist,prob,prob_thresh=0.5, nms_thresh=0.4)
+
+            coord = dist_to_coord(disti,points)
+            polys = dict(coord=coord, points=points, prob=probi)
+            labels = polygons_to_label(disti, points, prob=probi,shape=prob.shape)
+            labels = block.crop_context(labels, axes='YX')
+            labels, polys = block.filter_objects(labels, polys, axes='YX')
+            labels = relabel_sequential(labels, label_offset)[0]
+            if labels_out is not None:
+                block.write(labels_out, labels, axes='YX')
+        #for k,v in polys.items():
+            #polys_all.setdefault(k,[]).append(v)
+            label_offset += len(polys['prob'])
+            del labels
+    #polys_all = {k: (np.concatenate(v) if k in OBJECT_KEYS else v[0]) for k,v in polys_all.items()}
+    return labels_out
+def sliding_window_inference(
+    inputs: torch.Tensor,
+    roi_size: Union[Sequence[int], int],
+    sw_batch_size: int,
+    predictor: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor], Dict[Any, torch.Tensor]]],
+    overlap: float = 0.25,
+    mode: Union[BlendMode, str] = BlendMode.CONSTANT,
+    sigma_scale: Union[Sequence[float], float] = 0.125,
+    padding_mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT,
+    cval: float = 0.0,
+    sw_device: Union[torch.device, str, None] = None,
+    device: Union[torch.device, str, None] = None,
+    progress: bool = False,
+    roi_weight_map: Union[torch.Tensor, None] = None,
+    *args: Any,
+    **kwargs: Any,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, ...], Dict[Any, torch.Tensor]]:
+    """
+    Sliding window inference on `inputs` with `predictor`.
+
+    The outputs of `predictor` could be a tensor, a tuple, or a dictionary of tensors.
+    Each output in the tuple or dict value is allowed to have different resolutions with respect to the input.
+    e.g., the input patch spatial size is [128,128,128], the output (a tuple of two patches) patch sizes
+    could be ([128,64,256], [64,32,128]).
+    In this case, the parameter `overlap` and `roi_size` need to be carefully chosen to ensure the output ROI is still
+    an integer. If the predictor's input and output spatial sizes are not equal, we recommend choosing the parameters
+    so that `overlap*roi_size*output_size/input_size` is an integer (for each spatial dimension).
+
+    When roi_size is larger than the inputs' spatial size, the input image are padded during inference.
+    To maintain the same spatial sizes, the output image will be cropped to the original input size.
+
+    Args:
+        inputs: input image to be processed (assuming NCHW[D])
+        roi_size: the spatial window size for inferences.
+            When its components have None or non-positives, the corresponding inputs dimension will be used.
+            if the components of the `roi_size` are non-positive values, the transform will use the
+            corresponding components of img size. For example, `roi_size=(32, -1)` will be adapted
+            to `(32, 64)` if the second spatial dimension size of img is `64`.
+        sw_batch_size: the batch size to run window slices.
+        predictor: given input tensor ``patch_data`` in shape NCHW[D],
+            The outputs of the function call ``predictor(patch_data)`` should be a tensor, a tuple, or a dictionary
+            with Tensor values. Each output in the tuple or dict value should have the same batch_size, i.e. NM'H'W'[D'];
+            where H'W'[D'] represents the output patch's spatial size, M is the number of output channels,
+            N is `sw_batch_size`, e.g., the input shape is (7, 1, 128,128,128),
+            the output could be a tuple of two tensors, with shapes: ((7, 5, 128, 64, 256), (7, 4, 64, 32, 128)).
+            In this case, the parameter `overlap` and `roi_size` need to be carefully chosen
+            to ensure the scaled output ROI sizes are still integers.
+            If the `predictor`'s input and output spatial sizes are different,
+            we recommend choosing the parameters so that ``overlap*roi_size*zoom_scale`` is an integer for each dimension.
+        overlap: Amount of overlap between scans.
+        mode: {``"constant"``, ``"gaussian"``}
+            How to blend output of overlapping windows. Defaults to ``"constant"``.
+
+            - ``"constant``": gives equal weight to all predictions.
+            - ``"gaussian``": gives less weight to predictions on edges of windows.
+
+        sigma_scale: the standard deviation coefficient of the Gaussian window when `mode` is ``"gaussian"``.
+            Default: 0.125. Actual window sigma is ``sigma_scale`` * ``dim_size``.
+            When sigma_scale is a sequence of floats, the values denote sigma_scale at the corresponding
+            spatial dimensions.
+        padding_mode: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}
+            Padding mode for ``inputs``, when ``roi_size`` is larger than inputs. Defaults to ``"constant"``
+            See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html
+        cval: fill value for 'constant' padding mode. Default: 0
+        sw_device: device for the window data.
+            By default the device (and accordingly the memory) of the `inputs` is used.
+            Normally `sw_device` should be consistent with the device where `predictor` is defined.
+        device: device for the stitched output prediction.
+            By default the device (and accordingly the memory) of the `inputs` is used. If for example
+            set to device=torch.device('cpu') the gpu memory consumption is less and independent of the
+            `inputs` and `roi_size`. Output is on the `device`.
+        progress: whether to print a `tqdm` progress bar.
+        roi_weight_map: pre-computed (non-negative) weight map for each ROI.
+            If not given, and ``mode`` is not `constant`, this map will be computed on the fly.
+        args: optional args to be passed to ``predictor``.
+        kwargs: optional keyword args to be passed to ``predictor``.
+
+    Note:
+        - input must be channel-first and have a batch dim, supports N-D sliding window.
+
+    """
+    compute_dtype = inputs.dtype
+    num_spatial_dims = len(inputs.shape) - 2
+    if overlap < 0 or overlap >= 1:
+        raise ValueError("overlap must be >= 0 and < 1.")
+
+    # determine image spatial size and batch size
+    # Note: all input images must have the same image size and batch size
+    batch_size, _, *image_size_ = inputs.shape
+
+    if device is None:
+        device = inputs.device
+    if sw_device is None:
+        sw_device = inputs.device
+
+    roi_size = fall_back_tuple(roi_size, image_size_)
+    # in case that image size is smaller than roi size
+    image_size = tuple(max(image_size_[i], roi_size[i]) for i in range(num_spatial_dims))
+    pad_size = []
+    for k in range(len(inputs.shape) - 1, 1, -1):
+        diff = max(roi_size[k - 2] - inputs.shape[k], 0)
+        half = diff // 2
+        pad_size.extend([half, diff - half])
+    inputs = F.pad(inputs, pad=pad_size, mode=look_up_option(padding_mode, PytorchPadMode), value=cval)
+    #print('inputs',inputs.shape)
+    scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims, overlap)
+
+    # Store all slices in list
+    slices = dense_patch_slices(image_size, roi_size, scan_interval)
+    num_win = len(slices)  # number of windows per image
+    total_slices = num_win * batch_size  # total number of windows
+
+    # Create window-level importance map
+    valid_patch_size = get_valid_patch_size(image_size, roi_size)
+    if valid_patch_size == roi_size and (roi_weight_map is not None):
+        importance_map = roi_weight_map
+    else:
+        try:
+            importance_map = compute_importance_map(valid_patch_size, mode=mode, sigma_scale=sigma_scale, device=device)
+        except BaseException as e:
+            raise RuntimeError(
+                "Seems to be OOM. Please try smaller patch size or mode='constant' instead of mode='gaussian'."
+            ) from e
+    importance_map = convert_data_type(importance_map, torch.Tensor, device, compute_dtype)[0]  # type: ignore
+    # handle non-positive weights
+    min_non_zero = max(importance_map[importance_map != 0].min().item(), 1e-3)
+    importance_map = torch.clamp(importance_map.to(torch.float32), min=min_non_zero).to(compute_dtype)
+
+    # Perform predictions
+    dict_key, output_image_list, count_map_list = None, [], []
+    _initialized_ss = -1
+    is_tensor_output = True  # whether the predictor's output is a tensor (instead of dict/tuple)
+
+    # for each patch
+    for slice_g in tqdm(range(0, total_slices, sw_batch_size)) if progress else range(0, total_slices, sw_batch_size):
+        slice_range = range(slice_g, min(slice_g + sw_batch_size, total_slices))
+        unravel_slice = [
+            [slice(int(idx / num_win), int(idx / num_win) + 1), slice(None)] + list(slices[idx % num_win])
+            for idx in slice_range
+        ]
+        window_data = torch.cat(
+            [convert_data_type(inputs[win_slice], torch.Tensor)[0] for win_slice in unravel_slice]
+        ).to(sw_device)
+        seg_prob_out = predictor(window_data, *args, **kwargs)  # batched patch segmentation
+        #print('seg_prob_out',seg_prob_out[0].shape)
+        # convert seg_prob_out to tuple seg_prob_tuple, this does not allocate new memory.
+        seg_prob_tuple: Tuple[torch.Tensor, ...]
+        if isinstance(seg_prob_out, torch.Tensor):
+            seg_prob_tuple = (seg_prob_out,)
+        elif isinstance(seg_prob_out, Mapping):
+            if dict_key is None:
+                dict_key = sorted(seg_prob_out.keys())  # track predictor's output keys
+            seg_prob_tuple = tuple(seg_prob_out[k] for k in dict_key)
+            is_tensor_output = False
+        else:
+            seg_prob_tuple = ensure_tuple(seg_prob_out)
+            is_tensor_output = False
+
+        # for each output in multi-output list
+        for ss, seg_prob in enumerate(seg_prob_tuple):
+            seg_prob = seg_prob.to(device)  # BxCxMxNxP or BxCxMxN
+
+            # compute zoom scale: out_roi_size/in_roi_size
+            zoom_scale = []
+            for axis, (img_s_i, out_w_i, in_w_i) in enumerate(
+                zip(image_size, seg_prob.shape[2:], window_data.shape[2:])
+            ):
+                _scale = out_w_i / float(in_w_i)
+                if not (img_s_i * _scale).is_integer():
+                    warnings.warn(
+                        f"For spatial axis: {axis}, output[{ss}] will have non-integer shape. Spatial "
+                        f"zoom_scale between output[{ss}] and input is {_scale}. Please pad inputs."
+                    )
+                zoom_scale.append(_scale)
+
+            if _initialized_ss < ss:  # init. the ss-th buffer at the first iteration
+                # construct multi-resolution outputs
+                output_classes = seg_prob.shape[1]
+                output_shape = [batch_size, output_classes] + [
+                    int(image_size_d * zoom_scale_d) for image_size_d, zoom_scale_d in zip(image_size, zoom_scale)
+                ]
+                # allocate memory to store the full output and the count for overlapping parts
+                output_image_list.append(torch.zeros(output_shape, dtype=compute_dtype, device=device))
+                count_map_list.append(torch.zeros([1, 1] + output_shape[2:], dtype=compute_dtype, device=device))
+                _initialized_ss += 1
+
+            # resizing the importance_map
+            resizer = Resize(spatial_size=seg_prob.shape[2:], mode="nearest", anti_aliasing=False)
+
+            # store the result in the proper location of the full output. Apply weights from importance map.
+            for idx, original_idx in zip(slice_range, unravel_slice):
+                # zoom roi
+                original_idx_zoom = list(original_idx)  # 4D for 2D image, 5D for 3D image
+                for axis in range(2, len(original_idx_zoom)):
+                    zoomed_start = original_idx[axis].start * zoom_scale[axis - 2]
+                    zoomed_end = original_idx[axis].stop * zoom_scale[axis - 2]
+                    if not zoomed_start.is_integer() or (not zoomed_end.is_integer()):
+                        warnings.warn(
+                            f"For axis-{axis-2} of output[{ss}], the output roi range is not int. "
+                            f"Input roi range is ({original_idx[axis].start}, {original_idx[axis].stop}). "
+                            f"Spatial zoom_scale between output[{ss}] and input is {zoom_scale[axis - 2]}. "
+                            f"Corresponding output roi range is ({zoomed_start}, {zoomed_end}).\n"
+                            f"Please change overlap ({overlap}) or roi_size ({roi_size[axis-2]}) for axis-{axis-2}. "
+                            "Tips: if overlap*roi_size*zoom_scale is an integer, it usually works."
+                        )
+                    original_idx_zoom[axis] = slice(int(zoomed_start), int(zoomed_end), None)
+                importance_map_zoom = resizer(importance_map.unsqueeze(0))[0].to(compute_dtype)
+                # store results and weights
+                output_image_list[ss][original_idx_zoom] += importance_map_zoom * seg_prob[idx - slice_g]
+                count_map_list[ss][original_idx_zoom] += (
+                    importance_map_zoom.unsqueeze(0).unsqueeze(0).expand(count_map_list[ss][original_idx_zoom].shape)
+                )
+
+    # account for any overlapping sections
+    for ss in range(len(output_image_list)):
+        output_image_list[ss] = (output_image_list[ss] / count_map_list.pop(0)).to(compute_dtype)
+
+    # remove padding if image_size smaller than roi_size
+    for ss, output_i in enumerate(output_image_list):
+        if torch.isnan(output_i).any() or torch.isinf(output_i).any():
+            warnings.warn("Sliding window inference results contain NaN or Inf.")
+
+        zoom_scale = [
+            seg_prob_map_shape_d / roi_size_d for seg_prob_map_shape_d, roi_size_d in zip(output_i.shape[2:], roi_size)
+        ]
+
+        final_slicing: List[slice] = []
+        for sp in range(num_spatial_dims):
+            slice_dim = slice(pad_size[sp * 2], image_size_[num_spatial_dims - sp - 1] + pad_size[sp * 2])
+            slice_dim = slice(
+                int(round(slice_dim.start * zoom_scale[num_spatial_dims - sp - 1])),
+                int(round(slice_dim.stop * zoom_scale[num_spatial_dims - sp - 1])),
+            )
+            final_slicing.insert(0, slice_dim)
+        while len(final_slicing) < len(output_i.shape):
+            final_slicing.insert(0, slice(None))
+        output_image_list[ss] = output_i[final_slicing]
+
+    if dict_key is not None:  # if output of predictor is a dict
+        final_output = dict(zip(dict_key, output_image_list))
+    else:
+        final_output = tuple(output_image_list)  # type: ignore
+    final_output = final_output[0] if is_tensor_output else final_output
+
+    if isinstance(inputs, MetaTensor):
+        final_output = convert_to_dst_type(final_output, inputs, device=device)[0]  # type: ignore
+    return final_output
+
+
+def _get_scan_interval(
+    image_size: Sequence[int], roi_size: Sequence[int], num_spatial_dims: int, overlap: float
+) -> Tuple[int, ...]:
+    """
+    Compute scan interval according to the image size, roi size and overlap.
+    Scan interval will be `int((1 - overlap) * roi_size)`, if interval is 0,
+    use 1 instead to make sure sliding window works.
+
+    """
+    if len(image_size) != num_spatial_dims:
+        raise ValueError("image coord different from spatial dims.")
+    if len(roi_size) != num_spatial_dims:
+        raise ValueError("roi coord different from spatial dims.")
+
+    scan_interval = []
+    for i in range(num_spatial_dims):
+        if roi_size[i] == image_size[i]:
+            scan_interval.append(int(roi_size[i]))
+        else:
+            interval = int(roi_size[i] * (1 - overlap))
+            scan_interval.append(interval if interval > 0 else 1)
+    return tuple(scan_interval)