Init

app.py

@@ -0,0 +1,7 @@
+import gradio as gr
+
+def greet(name):
+    return "Hello " + name + "!!"
+
+iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+iface.launch()

main_model.pt

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

predict.py

@@ -0,0 +1,1256 @@
+import torch
+from torch.nn import (
+    Module,
+    Conv2d,
+    BatchNorm2d,
+    Identity,
+    UpsamplingBilinear2d,
+    Mish,
+    ReLU,
+    Sequential,
+)
+from torch.nn.functional import interpolate, grid_sample, pad
+import numpy as np
+from copy import deepcopy
+import os, argparse, math
+import tifffile as tif
+from typing import Tuple, List, Mapping
+
+from monai.utils import (
+    BlendMode,
+    PytorchPadMode,
+    convert_data_type,
+    ensure_tuple,
+    fall_back_tuple,
+    look_up_option,
+    convert_to_dst_type,
+)
+from monai.utils.misc import ensure_tuple_size, ensure_tuple_rep, issequenceiterable
+from monai.networks.layers.convutils import gaussian_1d
+from monai.networks.layers.simplelayers import separable_filtering
+
+from segmentation_models_pytorch import MAnet
+
+from skimage.io import imread as io_imread
+from skimage.util.dtype import dtype_range
+from skimage._shared.utils import _supported_float_type
+from scipy.ndimage import find_objects, binary_fill_holes
+
+
+########################### Data Loading Modules #########################################################
+DTYPE_RANGE = dtype_range.copy()
+DTYPE_RANGE.update((d.__name__, limits) for d, limits in dtype_range.items())
+DTYPE_RANGE.update(
+    {
+        "uint10": (0, 2 ** 10 - 1),
+        "uint12": (0, 2 ** 12 - 1),
+        "uint14": (0, 2 ** 14 - 1),
+        "bool": dtype_range[bool],
+        "float": dtype_range[np.float64],
+    }
+)
+
+
+def _output_dtype(dtype_or_range, image_dtype):
+    if type(dtype_or_range) in [list, tuple, np.ndarray]:
+        # pair of values: always return float.
+        return _supported_float_type(image_dtype)
+    if type(dtype_or_range) == type:
+        # already a type: return it
+        return dtype_or_range
+    if dtype_or_range in DTYPE_RANGE:
+        # string key in DTYPE_RANGE dictionary
+        try:
+            # if it's a canonical numpy dtype, convert
+            return np.dtype(dtype_or_range).type
+        except TypeError:  # uint10, uint12, uint14
+            # otherwise, return uint16
+            return np.uint16
+    else:
+        raise ValueError(
+            "Incorrect value for out_range, should be a valid image data "
+            f"type or a pair of values, got {dtype_or_range}."
+        )
+
+
+def intensity_range(image, range_values="image", clip_negative=False):
+    if range_values == "dtype":
+        range_values = image.dtype.type
+
+    if range_values == "image":
+        i_min = np.min(image)
+        i_max = np.max(image)
+    elif range_values in DTYPE_RANGE:
+        i_min, i_max = DTYPE_RANGE[range_values]
+        if clip_negative:
+            i_min = 0
+    else:
+        i_min, i_max = range_values
+    return i_min, i_max
+
+
+def rescale_intensity(image, in_range="image", out_range="dtype"):
+    out_dtype = _output_dtype(out_range, image.dtype)
+
+    imin, imax = map(float, intensity_range(image, in_range))
+    omin, omax = map(
+        float, intensity_range(image, out_range, clip_negative=(imin >= 0))
+    )
+    image = np.clip(image, imin, imax)
+
+    if imin != imax:
+        image = (image - imin) / (imax - imin)
+        return np.asarray(image * (omax - omin) + omin, dtype=out_dtype)
+    else:
+        return np.clip(image, omin, omax).astype(out_dtype)
+
+
+def _normalize(img):
+    non_zero_vals = img[np.nonzero(img)]
+    percentiles = np.percentile(non_zero_vals, [0, 99.5])
+    img_norm = rescale_intensity(
+        img, in_range=(percentiles[0], percentiles[1]), out_range="uint8"
+    )
+
+    return img_norm.astype(np.uint8)
+
+
+def pred_transforms(filename):
+    # LoadImage
+    img = (
+        tif.imread(filename)
+        if filename.endswith(".tif") or filename.endswith(".tiff")
+        else io_imread(filename)
+    )
+
+    if len(img.shape) == 2:
+        img = np.repeat(np.expand_dims(img, axis=-1), 3, axis=-1)
+    elif len(img.shape) == 3 and img.shape[-1] > 3:
+        img = img[:, :, :3]
+
+    img = img.astype(np.float32)
+    img = _normalize(img)
+    img = np.moveaxis(img, -1, 0)
+    img = (img - img.min()) / (img.max() - img.min())
+
+    return torch.FloatTensor(img).unsqueeze(0)
+
+
+################################################################################
+
+########################### MODEL Architecture #################################
+class SegformerGH(MAnet):
+    def __init__(
+        self,
+        encoder_name: str = "mit_b5",
+        encoder_weights="imagenet",
+        decoder_channels=(256, 128, 64, 32, 32),
+        decoder_pab_channels=256,
+        in_channels: int = 3,
+        classes: int = 3,
+    ):
+        super(SegformerGH, self).__init__(
+            encoder_name=encoder_name,
+            encoder_weights=encoder_weights,
+            decoder_channels=decoder_channels,
+            decoder_pab_channels=decoder_pab_channels,
+            in_channels=in_channels,
+            classes=classes,
+        )
+
+        convert_relu_to_mish(self.encoder)
+        convert_relu_to_mish(self.decoder)
+
+        self.cellprob_head = DeepSegmantationHead(
+            in_channels=decoder_channels[-1], out_channels=1, kernel_size=3,
+        )
+        self.gradflow_head = DeepSegmantationHead(
+            in_channels=decoder_channels[-1], out_channels=2, kernel_size=3,
+        )
+
+    def forward(self, x):
+        """Sequentially pass `x` trough model`s encoder, decoder and heads"""
+        self.check_input_shape(x)
+
+        features = self.encoder(x)
+        decoder_output = self.decoder(*features)
+
+        gradflow_mask = self.gradflow_head(decoder_output)
+        cellprob_mask = self.cellprob_head(decoder_output)
+
+        masks = torch.cat([gradflow_mask, cellprob_mask], dim=1)
+
+        return masks
+
+
+class DeepSegmantationHead(Sequential):
+    def __init__(self, in_channels, out_channels, kernel_size=3, upsampling=1):
+        conv2d_1 = Conv2d(
+            in_channels,
+            in_channels // 2,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+        )
+        bn = BatchNorm2d(in_channels // 2)
+        conv2d_2 = Conv2d(
+            in_channels // 2,
+            out_channels,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+        )
+        mish = Mish(inplace=True)
+
+        upsampling = (
+            UpsamplingBilinear2d(scale_factor=upsampling)
+            if upsampling > 1
+            else Identity()
+        )
+        activation = Identity()
+        super().__init__(conv2d_1, mish, bn, conv2d_2, upsampling, activation)
+
+
+def convert_relu_to_mish(model):
+    for child_name, child in model.named_children():
+        if isinstance(child, ReLU):
+            setattr(model, child_name, Mish(inplace=True))
+        else:
+            convert_relu_to_mish(child)
+
+
+#####################################################################################
+
+########################### Sliding Window Inference #################################
+class GaussianFilter(Module):
+    def __init__(
+        self, spatial_dims, sigma, truncated=4.0, approx="erf", requires_grad=False,
+    ) -> None:
+        if issequenceiterable(sigma):
+            if len(sigma) != spatial_dims:  # type: ignore
+                raise ValueError
+        else:
+            sigma = [deepcopy(sigma) for _ in range(spatial_dims)]  # type: ignore
+        super().__init__()
+        self.sigma = [
+            torch.nn.Parameter(
+                torch.as_tensor(
+                    s,
+                    dtype=torch.float,
+                    device=s.device if isinstance(s, torch.Tensor) else None,
+                ),
+                requires_grad=requires_grad,
+            )
+            for s in sigma  # type: ignore
+        ]
+        self.truncated = truncated
+        self.approx = approx
+        for idx, param in enumerate(self.sigma):
+            self.register_parameter(f"kernel_sigma_{idx}", param)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _kernel = [
+            gaussian_1d(s, truncated=self.truncated, approx=self.approx)
+            for s in self.sigma
+        ]
+        return separable_filtering(x=x, kernels=_kernel)
+
+
+def compute_importance_map(
+    patch_size, mode=BlendMode.CONSTANT, sigma_scale=0.125, device="cpu"
+):
+    mode = look_up_option(mode, BlendMode)
+    device = torch.device(device)
+
+    center_coords = [i // 2 for i in patch_size]
+    sigma_scale = ensure_tuple_rep(sigma_scale, len(patch_size))
+    sigmas = [i * sigma_s for i, sigma_s in zip(patch_size, sigma_scale)]
+
+    importance_map = torch.zeros(patch_size, device=device)
+    importance_map[tuple(center_coords)] = 1
+    pt_gaussian = GaussianFilter(len(patch_size), sigmas).to(
+        device=device, dtype=torch.float
+    )
+    importance_map = pt_gaussian(importance_map.unsqueeze(0).unsqueeze(0))
+    importance_map = importance_map.squeeze(0).squeeze(0)
+    importance_map = importance_map / torch.max(importance_map)
+    importance_map = importance_map.float()
+
+    return importance_map
+
+
+def first(iterable, default=None):
+    for i in iterable:
+        return i
+
+    return default
+
+
+def dense_patch_slices(image_size, patch_size, scan_interval):
+    num_spatial_dims = len(image_size)
+    patch_size = get_valid_patch_size(image_size, patch_size)
+    scan_interval = ensure_tuple_size(scan_interval, num_spatial_dims)
+
+    scan_num = []
+    for i in range(num_spatial_dims):
+        if scan_interval[i] == 0:
+            scan_num.append(1)
+        else:
+            num = int(math.ceil(float(image_size[i]) / scan_interval[i]))
+            scan_dim = first(
+                d
+                for d in range(num)
+                if d * scan_interval[i] + patch_size[i] >= image_size[i]
+            )
+            scan_num.append(scan_dim + 1 if scan_dim is not None else 1)
+
+    starts = []
+    for dim in range(num_spatial_dims):
+        dim_starts = []
+        for idx in range(scan_num[dim]):
+            start_idx = idx * scan_interval[dim]
+            start_idx -= max(start_idx + patch_size[dim] - image_size[dim], 0)
+            dim_starts.append(start_idx)
+        starts.append(dim_starts)
+    out = np.asarray([x.flatten() for x in np.meshgrid(*starts, indexing="ij")]).T
+    return [tuple(slice(s, s + patch_size[d]) for d, s in enumerate(x)) for x in out]
+
+
+def get_valid_patch_size(image_size, patch_size):
+    ndim = len(image_size)
+    patch_size_ = ensure_tuple_size(patch_size, ndim)
+
+    # ensure patch size dimensions are not larger than image dimension, if a dimension is None or 0 use whole dimension
+    return tuple(min(ms, ps or ms) for ms, ps in zip(image_size, patch_size_))
+
+
+class Resize:
+    def __init__(self, spatial_size):
+        self.size_mode = "all"
+        self.spatial_size = spatial_size
+
+    def __call__(self, img):
+        input_ndim = img.ndim - 1  # spatial ndim
+        output_ndim = len(ensure_tuple(self.spatial_size))
+
+        if output_ndim > input_ndim:
+            input_shape = ensure_tuple_size(img.shape, output_ndim + 1, 1)
+            img = img.reshape(input_shape)
+
+        spatial_size_ = fall_back_tuple(self.spatial_size, img.shape[1:])
+
+        if (
+            tuple(img.shape[1:]) == spatial_size_
+        ):  # spatial shape is already the desired
+            return img
+
+        img_, *_ = convert_data_type(img, torch.Tensor, dtype=torch.float)
+
+        resized = interpolate(
+            input=img_.unsqueeze(0), size=spatial_size_, mode="nearest",
+        )
+        out, *_ = convert_to_dst_type(resized.squeeze(0), img)
+        return out
+
+
+def sliding_window_inference(
+    inputs,
+    roi_size,
+    sw_batch_size,
+    predictor,
+    overlap,
+    mode=BlendMode.CONSTANT,
+    sigma_scale=0.125,
+    padding_mode=PytorchPadMode.CONSTANT,
+    cval=0.0,
+    sw_device=None,
+    device=None,
+    roi_weight_map=None,
+):
+    compute_dtype = inputs.dtype
+    num_spatial_dims = len(inputs.shape) - 2
+    batch_size, _, *image_size_ = inputs.shape
+
+    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 = pad(
+        inputs,
+        pad=pad_size,
+        mode=look_up_option(padding_mode, PytorchPadMode).value,
+        value=cval,
+    )
+
+    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:
+        importance_map = compute_importance_map(
+            valid_patch_size, mode=mode, sigma_scale=sigma_scale, device=device
+        )
+
+    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 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([inputs[win_slice] for win_slice in unravel_slice]).to(
+            sw_device
+        )
+        seg_prob_out = predictor(window_data)  # batched patch segmentation
+
+        # 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)
+
+                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:])
+
+            # 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]
+
+                    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):
+        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
+
+    return final_output[0] if is_tensor_output else final_output  # type: ignore
+
+
+def _get_scan_interval(
+    image_size, roi_size, num_spatial_dims: int, overlap: float
+) -> Tuple[int, ...]:
+    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)
+
+
+#####################################################################################
+
+########################### Main Inference Functions #################################
+def post_process(pred_mask, device):
+    dP, cellprob = pred_mask[:2], 1 / (1 + np.exp(-pred_mask[-1]))
+    H, W = pred_mask.shape[-2], pred_mask.shape[-1]
+
+    if np.prod(H * W) < (5000 * 5000):
+        pred_mask = compute_masks(
+            dP,
+            cellprob,
+            use_gpu=True,
+            flow_threshold=0.4,
+            device=device,
+            cellprob_threshold=0.4,
+        )[0]
+
+    else:
+        print("\n[Whole Slide] Grid Prediction starting...")
+        roi_size = 2000
+
+        # Get patch grid by roi_size
+        if H % roi_size != 0:
+            n_H = H // roi_size + 1
+            new_H = roi_size * n_H
+        else:
+            n_H = H // roi_size
+            new_H = H
+
+        if W % roi_size != 0:
+            n_W = W // roi_size + 1
+            new_W = roi_size * n_W
+        else:
+            n_W = W // roi_size
+            new_W = W
+
+        # Allocate values on the grid
+        pred_pad = np.zeros((new_H, new_W), dtype=np.uint32)
+        dP_pad = np.zeros((2, new_H, new_W), dtype=np.float32)
+        cellprob_pad = np.zeros((new_H, new_W), dtype=np.float32)
+
+        dP_pad[:, :H, :W], cellprob_pad[:H, :W] = dP, cellprob
+
+        for i in range(n_H):
+            for j in range(n_W):
+                print("Pred on Grid (%d, %d) processing..." % (i, j))
+                dP_roi = dP_pad[
+                    :,
+                    roi_size * i : roi_size * (i + 1),
+                    roi_size * j : roi_size * (j + 1),
+                ]
+                cellprob_roi = cellprob_pad[
+                    roi_size * i : roi_size * (i + 1),
+                    roi_size * j : roi_size * (j + 1),
+                ]
+
+                pred_mask = compute_masks(
+                    dP_roi,
+                    cellprob_roi,
+                    use_gpu=True,
+                    flow_threshold=0.4,
+                    device=device,
+                    cellprob_threshold=0.4,
+                )[0]
+
+                pred_pad[
+                    roi_size * i : roi_size * (i + 1),
+                    roi_size * j : roi_size * (j + 1),
+                ] = pred_mask
+
+        pred_mask = pred_pad[:H, :W]
+
+    cell_idx, cell_sizes = np.unique(pred_mask, return_counts=True)
+    cell_idx, cell_sizes = cell_idx[1:], cell_sizes[1:]
+    cell_drop = np.where(cell_sizes < np.mean(cell_sizes) - 2.7 * np.std(cell_sizes))
+
+    for drop_cell in cell_idx[cell_drop]:
+        pred_mask[pred_mask == drop_cell] = 0
+
+    return pred_mask
+
+
+def hflip(x):
+    """flip batch of images horizontally"""
+    return x.flip(3)
+
+
+def vflip(x):
+    """flip batch of images vertically"""
+    return x.flip(2)
+
+
+class DualTransform:
+    identity_param = None
+
+    def __init__(
+        self, name: str, params,
+    ):
+        self.params = params
+        self.pname = name
+
+    def apply_aug_image(self, image, *args, **params):
+        raise NotImplementedError
+
+    def apply_deaug_mask(self, mask, *args, **params):
+        raise NotImplementedError
+
+
+class HorizontalFlip(DualTransform):
+    """Flip images horizontally (left->right)"""
+
+    identity_param = False
+
+    def __init__(self):
+        super().__init__("apply", [False, True])
+
+    def apply_aug_image(self, image, apply=False, **kwargs):
+        if apply:
+            image = hflip(image)
+        return image
+
+    def apply_deaug_mask(self, mask, apply=False, **kwargs):
+        if apply:
+            mask = hflip(mask)
+        return mask
+
+
+class VerticalFlip(DualTransform):
+    """Flip images vertically (up->down)"""
+
+    identity_param = False
+
+    def __init__(self):
+        super().__init__("apply", [False, True])
+
+    def apply_aug_image(self, image, apply=False, **kwargs):
+        if apply:
+            image = vflip(image)
+        return image
+
+    def apply_deaug_mask(self, mask, apply=False, **kwargs):
+        if apply:
+            mask = vflip(mask)
+        return mask
+
+
+#################### GradFlow Modules ##################################################
+from scipy.ndimage.filters import maximum_filter1d
+import scipy.ndimage
+import fastremap
+from skimage import morphology
+
+from scipy.ndimage import mean
+
+torch_GPU = torch.device("cuda")
+torch_CPU = torch.device("cpu")
+
+
+def _extend_centers_gpu(
+    neighbors, centers, isneighbor, Ly, Lx, n_iter=200, device=torch.device("cuda")
+):
+    if device is not None:
+        device = device
+    nimg = neighbors.shape[0] // 9
+    pt = torch.from_numpy(neighbors).to(device)
+
+    T = torch.zeros((nimg, Ly, Lx), dtype=torch.double, device=device)
+    meds = torch.from_numpy(centers.astype(int)).to(device).long()
+    isneigh = torch.from_numpy(isneighbor).to(device)
+    for i in range(n_iter):
+        T[:, meds[:, 0], meds[:, 1]] += 1
+        Tneigh = T[:, pt[:, :, 0], pt[:, :, 1]]
+        Tneigh *= isneigh
+        T[:, pt[0, :, 0], pt[0, :, 1]] = Tneigh.mean(axis=1)
+    del meds, isneigh, Tneigh
+    T = torch.log(1.0 + T)
+    # gradient positions
+    grads = T[:, pt[[2, 1, 4, 3], :, 0], pt[[2, 1, 4, 3], :, 1]]
+    del pt
+    dy = grads[:, 0] - grads[:, 1]
+    dx = grads[:, 2] - grads[:, 3]
+    del grads
+    mu_torch = np.stack((dy.cpu().squeeze(), dx.cpu().squeeze()), axis=-2)
+    return mu_torch
+
+
+def diameters(masks):
+    _, counts = np.unique(np.int32(masks), return_counts=True)
+    counts = counts[1:]
+    md = np.median(counts ** 0.5)
+    if np.isnan(md):
+        md = 0
+    md /= (np.pi ** 0.5) / 2
+    return md, counts ** 0.5
+
+
+def masks_to_flows_gpu(masks, device=None):
+    if device is None:
+        device = torch.device("cuda")
+
+    Ly0, Lx0 = masks.shape
+    Ly, Lx = Ly0 + 2, Lx0 + 2
+
+    masks_padded = np.zeros((Ly, Lx), np.int64)
+    masks_padded[1:-1, 1:-1] = masks
+
+    # get mask pixel neighbors
+    y, x = np.nonzero(masks_padded)
+    neighborsY = np.stack((y, y - 1, y + 1, y, y, y - 1, y - 1, y + 1, y + 1), axis=0)
+    neighborsX = np.stack((x, x, x, x - 1, x + 1, x - 1, x + 1, x - 1, x + 1), axis=0)
+    neighbors = np.stack((neighborsY, neighborsX), axis=-1)
+
+    # get mask centers
+    slices = scipy.ndimage.find_objects(masks)
+
+    centers = np.zeros((masks.max(), 2), "int")
+    for i, si in enumerate(slices):
+        if si is not None:
+            sr, sc = si
+
+            ly, lx = sr.stop - sr.start + 1, sc.stop - sc.start + 1
+            yi, xi = np.nonzero(masks[sr, sc] == (i + 1))
+            yi = yi.astype(np.int32) + 1  # add padding
+            xi = xi.astype(np.int32) + 1  # add padding
+            ymed = np.median(yi)
+            xmed = np.median(xi)
+            imin = np.argmin((xi - xmed) ** 2 + (yi - ymed) ** 2)
+            xmed = xi[imin]
+            ymed = yi[imin]
+            centers[i, 0] = ymed + sr.start
+            centers[i, 1] = xmed + sc.start
+
+    # get neighbor validator (not all neighbors are in same mask)
+    neighbor_masks = masks_padded[neighbors[:, :, 0], neighbors[:, :, 1]]
+    isneighbor = neighbor_masks == neighbor_masks[0]
+    ext = np.array(
+        [[sr.stop - sr.start + 1, sc.stop - sc.start + 1] for sr, sc in slices]
+    )
+    n_iter = 2 * (ext.sum(axis=1)).max()
+    # run diffusion
+    mu = _extend_centers_gpu(
+        neighbors, centers, isneighbor, Ly, Lx, n_iter=n_iter, device=device
+    )
+
+    # normalize
+    mu /= 1e-20 + (mu ** 2).sum(axis=0) ** 0.5
+
+    # put into original image
+    mu0 = np.zeros((2, Ly0, Lx0))
+    mu0[:, y - 1, x - 1] = mu
+    mu_c = np.zeros_like(mu0)
+    return mu0, mu_c
+
+
+def masks_to_flows(masks, use_gpu=False, device=None):
+    if masks.max() == 0 or (masks != 0).sum() == 1:
+        # dynamics_logger.warning('empty masks!')
+        return np.zeros((2, *masks.shape), "float32")
+
+    if use_gpu:
+        if use_gpu and device is None:
+            device = torch_GPU
+        elif device is None:
+            device = torch_CPU
+        masks_to_flows_device = masks_to_flows_gpu
+
+    if masks.ndim == 3:
+        Lz, Ly, Lx = masks.shape
+        mu = np.zeros((3, Lz, Ly, Lx), np.float32)
+        for z in range(Lz):
+            mu0 = masks_to_flows_device(masks[z], device=device)[0]
+            mu[[1, 2], z] += mu0
+        for y in range(Ly):
+            mu0 = masks_to_flows_device(masks[:, y], device=device)[0]
+            mu[[0, 2], :, y] += mu0
+        for x in range(Lx):
+            mu0 = masks_to_flows_device(masks[:, :, x], device=device)[0]
+            mu[[0, 1], :, :, x] += mu0
+        return mu
+    elif masks.ndim == 2:
+        mu, mu_c = masks_to_flows_device(masks, device=device)
+        return mu
+
+    else:
+        raise ValueError("masks_to_flows only takes 2D or 3D arrays")
+
+
+def steps2D_interp(p, dP, niter, use_gpu=False, device=None):
+    shape = dP.shape[1:]
+    if use_gpu:
+        if device is None:
+            device = torch_GPU
+        shape = (
+            np.array(shape)[[1, 0]].astype("float") - 1
+        )  # Y and X dimensions (dP is 2.Ly.Lx), flipped X-1, Y-1
+        pt = (
+            torch.from_numpy(p[[1, 0]].T).float().to(device).unsqueeze(0).unsqueeze(0)
+        )  # p is n_points by 2, so pt is [1 1 2 n_points]
+        im = (
+            torch.from_numpy(dP[[1, 0]]).float().to(device).unsqueeze(0)
+        )  # covert flow numpy array to tensor on GPU, add dimension
+        # normalize pt between  0 and  1, normalize the flow
+        for k in range(2):
+            im[:, k, :, :] *= 2.0 / shape[k]
+            pt[:, :, :, k] /= shape[k]
+
+        # normalize to between -1 and 1
+        pt = pt * 2 - 1
+
+        # here is where the stepping happens
+        for t in range(niter):
+            # align_corners default is False, just added to suppress warning
+            dPt = grid_sample(im, pt, align_corners=False)
+
+            for k in range(2):  # clamp the final pixel locations
+                pt[:, :, :, k] = torch.clamp(
+                    pt[:, :, :, k] + dPt[:, k, :, :], -1.0, 1.0
+                )
+
+        # undo the normalization from before, reverse order of operations
+        pt = (pt + 1) * 0.5
+        for k in range(2):
+            pt[:, :, :, k] *= shape[k]
+
+        p = pt[:, :, :, [1, 0]].cpu().numpy().squeeze().T
+        return p
+
+    else:
+        assert print("ho")
+
+
+def follow_flows(dP, mask=None, niter=200, interp=True, use_gpu=True, device=None):
+    shape = np.array(dP.shape[1:]).astype(np.int32)
+    niter = np.uint32(niter)
+
+    p = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing="ij")
+    p = np.array(p).astype(np.float32)
+
+    inds = np.array(np.nonzero(np.abs(dP[0]) > 1e-3)).astype(np.int32).T
+
+    if inds.ndim < 2 or inds.shape[0] < 5:
+        return p, None
+
+    if not interp:
+        assert print("woo")
+
+    else:
+        p_interp = steps2D_interp(
+            p[:, inds[:, 0], inds[:, 1]], dP, niter, use_gpu=use_gpu, device=device
+        )
+        p[:, inds[:, 0], inds[:, 1]] = p_interp
+
+    return p, inds
+
+
+def flow_error(maski, dP_net, use_gpu=False, device=None):
+    if dP_net.shape[1:] != maski.shape:
+        print("ERROR: net flow is not same size as predicted masks")
+        return
+
+    # flows predicted from estimated masks
+    dP_masks = masks_to_flows(maski, use_gpu=use_gpu, device=device)
+    # difference between predicted flows vs mask flows
+    flow_errors = np.zeros(maski.max())
+    for i in range(dP_masks.shape[0]):
+        flow_errors += mean(
+            (dP_masks[i] - dP_net[i] / 5.0) ** 2,
+            maski,
+            index=np.arange(1, maski.max() + 1),
+        )
+
+    return flow_errors, dP_masks
+
+
+def remove_bad_flow_masks(masks, flows, threshold=0.4, use_gpu=False, device=None):
+    merrors, _ = flow_error(masks, flows, use_gpu, device)
+    badi = 1 + (merrors > threshold).nonzero()[0]
+    masks[np.isin(masks, badi)] = 0
+    return masks
+
+
+def get_masks(p, iscell=None, rpad=20):
+    pflows = []
+    edges = []
+    shape0 = p.shape[1:]
+    dims = len(p)
+
+    for i in range(dims):
+        pflows.append(p[i].flatten().astype("int32"))
+        edges.append(np.arange(-0.5 - rpad, shape0[i] + 0.5 + rpad, 1))
+
+    h, _ = np.histogramdd(tuple(pflows), bins=edges)
+    hmax = h.copy()
+    for i in range(dims):
+        hmax = maximum_filter1d(hmax, 5, axis=i)
+
+    seeds = np.nonzero(np.logical_and(h - hmax > -1e-6, h > 10))
+    Nmax = h[seeds]
+    isort = np.argsort(Nmax)[::-1]
+    for s in seeds:
+        s = s[isort]
+
+    pix = list(np.array(seeds).T)
+
+    shape = h.shape
+    if dims == 3:
+        expand = np.nonzero(np.ones((3, 3, 3)))
+    else:
+        expand = np.nonzero(np.ones((3, 3)))
+    for e in expand:
+        e = np.expand_dims(e, 1)
+
+    for iter in range(5):
+        for k in range(len(pix)):
+            if iter == 0:
+                pix[k] = list(pix[k])
+            newpix = []
+            iin = []
+            for i, e in enumerate(expand):
+                epix = e[:, np.newaxis] + np.expand_dims(pix[k][i], 0) - 1
+                epix = epix.flatten()
+                iin.append(np.logical_and(epix >= 0, epix < shape[i]))
+                newpix.append(epix)
+            iin = np.all(tuple(iin), axis=0)
+            for p in newpix:
+                p = p[iin]
+            newpix = tuple(newpix)
+            igood = h[newpix] > 2
+            for i in range(dims):
+                pix[k][i] = newpix[i][igood]
+            if iter == 4:
+                pix[k] = tuple(pix[k])
+
+    M = np.zeros(h.shape, np.uint32)
+    for k in range(len(pix)):
+        M[pix[k]] = 1 + k
+
+    for i in range(dims):
+        pflows[i] = pflows[i] + rpad
+    M0 = M[tuple(pflows)]
+
+    # remove big masks
+    uniq, counts = fastremap.unique(M0, return_counts=True)
+    big = np.prod(shape0) * 0.9
+    bigc = uniq[counts > big]
+    if len(bigc) > 0 and (len(bigc) > 1 or bigc[0] != 0):
+        M0 = fastremap.mask(M0, bigc)
+    fastremap.renumber(M0, in_place=True)  # convenient to guarantee non-skipped labels
+    M0 = np.reshape(M0, shape0)
+    return M0
+
+def fill_holes_and_remove_small_masks(masks, min_size=15):
+    """ fill holes in masks (2D/3D) and discard masks smaller than min_size (2D)
+    
+    fill holes in each mask using scipy.ndimage.morphology.binary_fill_holes
+    (might have issues at borders between cells, todo: check and fix)
+    
+    Parameters
+    ----------------
+    masks: int, 2D or 3D array
+        labelled masks, 0=NO masks; 1,2,...=mask labels,
+        size [Ly x Lx] or [Lz x Ly x Lx]
+    min_size: int (optional, default 15)
+        minimum number of pixels per mask, can turn off with -1
+    Returns
+    ---------------
+    masks: int, 2D or 3D array
+        masks with holes filled and masks smaller than min_size removed, 
+        0=NO masks; 1,2,...=mask labels,
+        size [Ly x Lx] or [Lz x Ly x Lx]
+    
+    """
+    
+    slices = find_objects(masks)
+    j = 0
+    for i,slc in enumerate(slices):
+        if slc is not None:
+            msk = masks[slc] == (i+1)
+            npix = msk.sum()
+            if min_size > 0 and npix < min_size:
+                masks[slc][msk] = 0
+            elif npix > 0:   
+                if msk.ndim==3:
+                    for k in range(msk.shape[0]):
+                        msk[k] = binary_fill_holes(msk[k])
+                else:          
+                    msk = binary_fill_holes(msk)
+                masks[slc][msk] = (j+1)
+                j+=1
+    return masks
+
+def compute_masks(
+    dP,
+    cellprob,
+    p=None,
+    niter=200,
+    cellprob_threshold=0.4,
+    flow_threshold=0.4,
+    interp=True,
+    resize=None,
+    use_gpu=False,
+    device=None,
+):
+    """compute masks using dynamics from dP, cellprob, and boundary"""
+
+    cp_mask = cellprob > cellprob_threshold
+    cp_mask = morphology.remove_small_holes(cp_mask, area_threshold=16)
+    cp_mask = morphology.remove_small_objects(cp_mask, min_size=16)
+
+    if np.any(cp_mask):  # mask at this point is a cell cluster binary map, not labels
+        # follow flows
+        if p is None:
+            p, inds = follow_flows(
+                dP * cp_mask / 5.0,
+                niter=niter,
+                interp=interp,
+                use_gpu=use_gpu,
+                device=device,
+            )
+            if inds is None:
+                shape = resize if resize is not None else cellprob.shape
+                mask = np.zeros(shape, np.uint16)
+                p = np.zeros((len(shape), *shape), np.uint16)
+                return mask, p
+
+        # calculate masks
+        mask = get_masks(p, iscell=cp_mask)
+
+        # flow thresholding factored out of get_masks
+        shape0 = p.shape[1:]
+        if mask.max() > 0 and flow_threshold is not None and flow_threshold > 0:
+            # make sure labels are unique at output of get_masks
+            mask = remove_bad_flow_masks(
+                mask, dP, threshold=flow_threshold, use_gpu=use_gpu, device=device
+            )
+        
+        mask = fill_holes_and_remove_small_masks(mask, min_size=15)
+        
+    else:  # nothing to compute, just make it compatible
+        shape = resize if resize is not None else cellprob.shape
+        mask = np.zeros(shape, np.uint16)
+        p = np.zeros((len(shape), *shape), np.uint16)
+        return mask, p
+    
+    return mask, p
+
+def main(args):
+    model = torch.load(args.model_path, map_location=args.device)
+    model.eval()
+    hflip_tta = HorizontalFlip()
+    vflip_tta = VerticalFlip()
+
+    img_names = sorted(os.listdir(args.input_path))
+    os.makedirs(args.output_path, exist_ok=True)
+
+    for img_name in img_names:
+        print(f"Segmenting {img_name}")
+        img_path = os.path.join(args.input_path, img_name)
+        img_data = pred_transforms(img_path)
+        img_data = img_data.to(args.device)
+        img_size = img_data.shape[-1] * img_data.shape[-2]
+        
+        if img_size < 1150000 and 900000 < img_size:
+            overlap = 0.5
+        else:
+            overlap = 0.6
+
+        with torch.no_grad():
+            img0 = img_data
+            outputs0 = sliding_window_inference(
+                img0,
+                512,
+                4,
+                model,
+                padding_mode="reflect",
+                mode="gaussian",
+                overlap=overlap,
+                device="cpu",
+            )
+            outputs0 = outputs0.cpu().squeeze()
+
+            if img_size < 2000 * 2000:
+                
+                model.load_state_dict(torch.load(args.model_path2, map_location=args.device))
+                model.eval()
+                
+                img2 = hflip_tta.apply_aug_image(img_data, apply=True)
+                outputs2 = sliding_window_inference(
+                    img2,
+                    512,
+                    4,
+                    model,
+                    padding_mode="reflect",
+                    mode="gauusian",
+                    overlap=overlap,
+                    device="cpu",
+                )
+                outputs2 = hflip_tta.apply_deaug_mask(outputs2, apply=True)
+                outputs2 = outputs2.cpu().squeeze()
+
+                outputs = torch.zeros_like(outputs0)
+                outputs[0] = (outputs0[0] + outputs2[0]) / 2
+                outputs[1] = (outputs0[1] - outputs2[1]) / 2
+                outputs[2] = (outputs0[2] + outputs2[2]) / 2
+                
+            elif img_size < 5000*5000:
+                # Hflip TTA
+                img2 = hflip_tta.apply_aug_image(img_data, apply=True)
+                outputs2 = sliding_window_inference(
+                    img2,
+                    512,
+                    4,
+                    model,
+                    padding_mode="reflect",
+                    mode="gaussian",
+                    overlap=overlap,
+                    device="cpu",
+                )
+                outputs2 = hflip_tta.apply_deaug_mask(outputs2, apply=True)
+                outputs2 = outputs2.cpu().squeeze()
+                img2 = img2.cpu()
+                
+                ##################
+                #                #
+                #    ensemble    #
+                #                #
+                ##################
+                
+                model.load_state_dict(torch.load(args.model_path2, map_location=args.device))
+                model.eval()
+                
+                img1 = img_data
+                outputs1 = sliding_window_inference(
+                    img1,
+                    512,
+                    4,
+                    model,
+                    padding_mode="reflect",
+                    mode="gaussian",
+                    overlap=overlap,
+                    device="cpu",
+                )
+                outputs1 = outputs1.cpu().squeeze()
+                
+                # Vflip TTA
+                img3 = vflip_tta.apply_aug_image(img_data, apply=True)
+                outputs3 = sliding_window_inference(
+                    img3,
+                    512,
+                    4,
+                    model,
+                    padding_mode="reflect",
+                    mode="gaussian",
+                    overlap=overlap,
+                    device="cpu",
+                )
+                outputs3 = vflip_tta.apply_deaug_mask(outputs3, apply=True)
+                outputs3 = outputs3.cpu().squeeze()
+                img3 = img3.cpu()
+
+                # Merge Results
+                outputs = torch.zeros_like(outputs0)
+                outputs[0] = (outputs0[0] + outputs1[0] + outputs2[0] - outputs3[0]) / 4
+                outputs[1] = (outputs0[1] + outputs1[1] - outputs2[1] + outputs3[1]) / 4
+                outputs[2] = (outputs0[2] + outputs1[2] + outputs2[2] + outputs3[2]) / 4
+            else:
+                outputs = outputs0
+
+            pred_mask = post_process(outputs.squeeze(0).cpu().numpy(), args.device)
+
+        file_path = os.path.join(
+            args.output_path, img_name.split(".")[0] + "_label.tiff"
+        )
+
+        tif.imwrite(file_path, pred_mask, compression="zlib")
+
+
+parser = argparse.ArgumentParser("Submission for Challenge", add_help=False)
+parser.add_argument("--model_path", default="./model.pt", type=str)
+parser.add_argument("--model_path2", default="./model_sec.pth", type=str)
+
+# Dataset parameters
+parser.add_argument(
+    "-i",
+    "--input_path",
+    default="/workspace/inputs/",
+    type=str,
+    help="training data path; subfolders: images, labels",
+)
+parser.add_argument(
+    "-o", "--output_path", default="/workspace/outputs/", type=str, help="output path",
+)
+parser.add_argument("--device", default="cuda:0", type=str)
+
+args = parser.parse_args()
+
+if __name__ == "__main__":
+    print("Starting")
+    main(args)

predict.sh

@@ -0,0 +1 @@
+python predict.py -i "./inputs"  -o "./outputs" --device "cuda:0" --model_path="./main_model.pt" --model_path2="./sub_model.pth"

requirements.txt

@@ -0,0 +1,83 @@
+backcall @ file:///home/ktietz/src/ci/backcall_1611930011877/work
+beautifulsoup4 @ file:///opt/conda/conda-bld/beautifulsoup4_1650462163268/work
+brotlipy==0.7.0
+certifi @ file:///opt/conda/conda-bld/certifi_1655968806487/work/certifi
+cffi @ file:///opt/conda/conda-bld/cffi_1642701102775/work
+chardet @ file:///tmp/build/80754af9/chardet_1607706768982/work
+charset-normalizer @ file:///tmp/build/80754af9/charset-normalizer_1630003229654/work
+colorama @ file:///tmp/build/80754af9/colorama_1607707115595/work
+coloredlogs==15.0.1
+conda==4.13.0
+conda-build==3.21.9
+conda-content-trust @ file:///tmp/build/80754af9/conda-content-trust_1617045594566/work
+conda-package-handling @ file:///tmp/build/80754af9/conda-package-handling_1649105789509/work
+cryptography @ file:///tmp/build/80754af9/cryptography_1652083456434/work
+decorator @ file:///opt/conda/conda-bld/decorator_1643638310831/work
+fastremap==1.13.3
+filelock @ file:///opt/conda/conda-bld/filelock_1647002191454/work
+flatbuffers==22.9.24
+glob2 @ file:///home/linux1/recipes/ci/glob2_1610991677669/work
+huggingface-hub==0.10.1
+humanfriendly==10.0
+idna @ file:///tmp/build/80754af9/idna_1637925883363/work
+imagecodecs==2021.11.20
+imageio==2.22.2
+importlib-metadata==5.0.0
+itk==5.2.1.post1
+itk-core==5.2.1.post1
+itk-filtering==5.2.1.post1
+itk-io==5.2.1.post1
+itk-numerics==5.2.1.post1
+itk-registration==5.2.1.post1
+itk-segmentation==5.2.1.post1
+jedi @ file:///tmp/build/80754af9/jedi_1644299024593/work
+Jinja2==2.10.1
+libarchive-c @ file:///tmp/build/80754af9/python-libarchive-c_1617780486945/work
+MarkupSafe @ file:///tmp/build/80754af9/markupsafe_1621528142364/work
+matplotlib-inline @ file:///tmp/build/80754af9/matplotlib-inline_1628242447089/work
+mkl-fft==1.3.1
+mkl-random @ file:///tmp/build/80754af9/mkl_random_1626179032232/work
+mkl-service==2.4.0
+monai==0.9.0
+mpmath==1.2.1
+networkx==2.6.3
+numpy @ file:///opt/conda/conda-bld/numpy_and_numpy_base_1651563629415/work
+onnxruntime-gpu==1.12.1
+opencv-python==4.6.0.66
+packaging==21.3
+parso @ file:///opt/conda/conda-bld/parso_1641458642106/work
+pexpect @ file:///tmp/build/80754af9/pexpect_1605563209008/work
+pickleshare @ file:///tmp/build/80754af9/pickleshare_1606932040724/work
+Pillow==9.0.1
+pkginfo @ file:///tmp/build/80754af9/pkginfo_1643162084911/work
+prompt-toolkit @ file:///tmp/build/80754af9/prompt-toolkit_1633440160888/work
+protobuf==4.21.8
+psutil @ file:///tmp/build/80754af9/psutil_1612298016854/work
+ptyprocess @ file:///tmp/build/80754af9/ptyprocess_1609355006118/work/dist/ptyprocess-0.7.0-py2.py3-none-any.whl
+pycosat==0.6.3
+pycparser @ file:///tmp/build/80754af9/pycparser_1636541352034/work
+Pygments @ file:///opt/conda/conda-bld/pygments_1644249106324/work
+pyOpenSSL @ file:///opt/conda/conda-bld/pyopenssl_1643788558760/work
+pyparsing==3.0.9
+PySocks @ file:///tmp/build/80754af9/pysocks_1594394576006/work
+pytz==2022.2.1
+PyWavelets==1.3.0
+PyYAML==6.0
+requests @ file:///opt/conda/conda-bld/requests_1641824580448/work
+ruamel-yaml-conda @ file:///tmp/build/80754af9/ruamel_yaml_1616016701961/work
+scikit-image==0.19.3
+scipy==1.7.2
+six @ file:///tmp/build/80754af9/six_1644875935023/work
+soupsieve @ file:///tmp/build/80754af9/soupsieve_1636706018808/work
+sympy==1.10.1
+tifffile==2021.11.2
+timm==0.6.11
+torch==1.12.1
+torchtext==0.13.1
+torchvision==0.13.1
+tqdm==4.64.1
+traitlets @ file:///tmp/build/80754af9/traitlets_1636710298902/work
+typing_extensions @ file:///tmp/abs_ben9emwtky/croots/recipe/typing_extensions_1659638822008/work
+urllib3 @ file:///opt/conda/conda-bld/urllib3_1643638302206/work
+wcwidth @ file:///Users/ktietz/demo/mc3/conda-bld/wcwidth_1629357192024/work
+zipp==3.9.0

save_model.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn as nn
+
+from segmentation_models_pytorch import MAnet
+from segmentation_models_pytorch.base.modules import Activation
+
+
+class SegformerGH(MAnet):
+    def __init__(
+        self,
+        encoder_name: str = "mit_b5",
+        encoder_weights="imagenet",
+        decoder_channels=(256, 128, 64, 32, 32),
+        decoder_pab_channels=256,
+        in_channels: int = 3,
+        classes: int = 3,
+    ):
+        super(SegformerGH, self).__init__(
+            encoder_name=encoder_name,
+            encoder_weights=encoder_weights,
+            decoder_channels=decoder_channels,
+            decoder_pab_channels=decoder_pab_channels,
+            in_channels=in_channels,
+            classes=classes,
+        )
+
+        convert_relu_to_mish(self.encoder)
+        convert_relu_to_mish(self.decoder)
+
+        self.cellprob_head = DeepSegmantationHead(
+            in_channels=decoder_channels[-1], out_channels=1, kernel_size=3,
+        )
+        self.gradflow_head = DeepSegmantationHead(
+            in_channels=decoder_channels[-1], out_channels=2, kernel_size=3,
+        )
+
+    def forward(self, x):
+        """Sequentially pass `x` trough model`s encoder, decoder and heads"""
+        self.check_input_shape(x)
+
+        features = self.encoder(x)
+        decoder_output = self.decoder(*features)
+
+        gradflow_mask = self.gradflow_head(decoder_output)
+        cellprob_mask = self.cellprob_head(decoder_output)
+
+        masks = torch.cat([gradflow_mask, cellprob_mask], dim=1)
+
+        return masks
+
+
+class DeepSegmantationHead(nn.Sequential):
+    def __init__(
+        self, in_channels, out_channels, kernel_size=3, activation=None, upsampling=1
+    ):
+        conv2d_1 = nn.Conv2d(
+            in_channels,
+            in_channels // 2,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+        )
+        bn = nn.BatchNorm2d(in_channels // 2)
+        conv2d_2 = nn.Conv2d(
+            in_channels // 2,
+            out_channels,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+        )
+        mish = nn.Mish(inplace=True)
+
+        upsampling = (
+            nn.UpsamplingBilinear2d(scale_factor=upsampling)
+            if upsampling > 1
+            else nn.Identity()
+        )
+        activation = Activation(activation)
+        super().__init__(conv2d_1, mish, bn, conv2d_2, upsampling, activation)
+
+
+def convert_relu_to_mish(model):
+    for child_name, child in model.named_children():
+        if isinstance(child, nn.ReLU):
+            setattr(model, child_name, nn.Mish(inplace=True))
+        else:
+            convert_relu_to_mish(child)
+
+
+if __name__ == "__main__":
+    model = SegformerGH(
+        encoder_name="mit_b5",
+        encoder_weights=None,
+        decoder_channels=(1024, 512, 256, 128, 64),
+        decoder_pab_channels=256,
+        in_channels=3,
+        classes=3,
+    )
+
+    model.load_state_dict(torch.load("./main_model.pth",map_location="cpu"))
+    torch.save(model, "main_model.pt")

segmentation_models_pytorch/__init__.py

@@ -0,0 +1,61 @@
+from . import datasets
+from . import encoders
+from . import decoders
+from . import losses
+from . import metrics
+
+from .decoders.unet import Unet
+from .decoders.unetplusplus import UnetPlusPlus
+from .decoders.manet import MAnet
+from .decoders.linknet import Linknet
+from .decoders.fpn import FPN
+from .decoders.pspnet import PSPNet
+from .decoders.deeplabv3 import DeepLabV3, DeepLabV3Plus
+from .decoders.pan import PAN
+
+from .__version__ import __version__
+
+# some private imports for create_model function
+from typing import Optional as _Optional
+import torch as _torch
+
+
+def create_model(
+    arch: str,
+    encoder_name: str = "resnet34",
+    encoder_weights: _Optional[str] = "imagenet",
+    in_channels: int = 3,
+    classes: int = 1,
+    **kwargs,
+) -> _torch.nn.Module:
+    """Models entrypoint, allows to create any model architecture just with
+    parameters, without using its class
+    """
+
+    archs = [
+        Unet,
+        UnetPlusPlus,
+        MAnet,
+        Linknet,
+        FPN,
+        PSPNet,
+        DeepLabV3,
+        DeepLabV3Plus,
+        PAN,
+    ]
+    archs_dict = {a.__name__.lower(): a for a in archs}
+    try:
+        model_class = archs_dict[arch.lower()]
+    except KeyError:
+        raise KeyError(
+            "Wrong architecture type `{}`. Available options are: {}".format(
+                arch, list(archs_dict.keys()),
+            )
+        )
+    return model_class(
+        encoder_name=encoder_name,
+        encoder_weights=encoder_weights,
+        in_channels=in_channels,
+        classes=classes,
+        **kwargs,
+    )

segmentation_models_pytorch/__version__.py

@@ -0,0 +1,3 @@
+VERSION = (0, 3, 0)
+
+__version__ = ".".join(map(str, VERSION))

segmentation_models_pytorch/base/__init__.py

@@ -0,0 +1,11 @@
+from .model import SegmentationModel
+
+from .modules import (
+    Conv2dReLU,
+    Attention,
+)
+
+from .heads import (
+    SegmentationHead,
+    ClassificationHead,
+)

segmentation_models_pytorch/base/heads.py

@@ -0,0 +1,34 @@
+import torch.nn as nn
+from .modules import Activation
+
+
+class SegmentationHead(nn.Sequential):
+    def __init__(
+        self, in_channels, out_channels, kernel_size=3, activation=None, upsampling=1
+    ):
+        conv2d = nn.Conv2d(
+            in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2
+        )
+        upsampling = (
+            nn.UpsamplingBilinear2d(scale_factor=upsampling)
+            if upsampling > 1
+            else nn.Identity()
+        )
+        activation = Activation(activation)
+        super().__init__(conv2d, upsampling, activation)
+
+
+class ClassificationHead(nn.Sequential):
+    def __init__(
+        self, in_channels, classes, pooling="avg", dropout=0.2, activation=None
+    ):
+        if pooling not in ("max", "avg"):
+            raise ValueError(
+                "Pooling should be one of ('max', 'avg'), got {}.".format(pooling)
+            )
+        pool = nn.AdaptiveAvgPool2d(1) if pooling == "avg" else nn.AdaptiveMaxPool2d(1)
+        flatten = nn.Flatten()
+        dropout = nn.Dropout(p=dropout, inplace=True) if dropout else nn.Identity()
+        linear = nn.Linear(in_channels, classes, bias=True)
+        activation = Activation(activation)
+        super().__init__(pool, flatten, dropout, linear, activation)

segmentation_models_pytorch/base/initialization.py

@@ -0,0 +1,27 @@
+import torch.nn as nn
+
+
+def initialize_decoder(module):
+    for m in module.modules():
+
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_uniform_(m.weight, mode="fan_in", nonlinearity="relu")
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+        elif isinstance(m, nn.BatchNorm2d):
+            nn.init.constant_(m.weight, 1)
+            nn.init.constant_(m.bias, 0)
+
+        elif isinstance(m, nn.Linear):
+            nn.init.xavier_uniform_(m.weight)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+
+def initialize_head(module):
+    for m in module.modules():
+        if isinstance(m, (nn.Linear, nn.Conv2d)):
+            nn.init.xavier_uniform_(m.weight)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)

segmentation_models_pytorch/base/model.py

@@ -0,0 +1,64 @@
+import torch
+from . import initialization as init
+
+
+class SegmentationModel(torch.nn.Module):
+    def initialize(self):
+        init.initialize_decoder(self.decoder)
+        init.initialize_head(self.segmentation_head)
+        if self.classification_head is not None:
+            init.initialize_head(self.classification_head)
+
+    def check_input_shape(self, x):
+
+        h, w = x.shape[-2:]
+        output_stride = self.encoder.output_stride
+        if h % output_stride != 0 or w % output_stride != 0:
+            new_h = (
+                (h // output_stride + 1) * output_stride
+                if h % output_stride != 0
+                else h
+            )
+            new_w = (
+                (w // output_stride + 1) * output_stride
+                if w % output_stride != 0
+                else w
+            )
+            raise RuntimeError(
+                f"Wrong input shape height={h}, width={w}. Expected image height and width "
+                f"divisible by {output_stride}. Consider pad your images to shape ({new_h}, {new_w})."
+            )
+
+    def forward(self, x):
+        """Sequentially pass `x` trough model`s encoder, decoder and heads"""
+
+        self.check_input_shape(x)
+
+        features = self.encoder(x)
+        decoder_output = self.decoder(*features)
+
+        masks = self.segmentation_head(decoder_output)
+
+        if self.classification_head is not None:
+            labels = self.classification_head(features[-1])
+            return masks, labels
+
+        return masks
+
+    @torch.no_grad()
+    def predict(self, x):
+        """Inference method. Switch model to `eval` mode, call `.forward(x)` with `torch.no_grad()`
+
+        Args:
+            x: 4D torch tensor with shape (batch_size, channels, height, width)
+
+        Return:
+            prediction: 4D torch tensor with shape (batch_size, classes, height, width)
+
+        """
+        if self.training:
+            self.eval()
+
+        x = self.forward(x)
+
+        return x

segmentation_models_pytorch/base/modules.py

@@ -0,0 +1,131 @@
+import torch
+import torch.nn as nn
+
+try:
+    from inplace_abn import InPlaceABN
+except ImportError:
+    InPlaceABN = None
+
+
+class Conv2dReLU(nn.Sequential):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        padding=0,
+        stride=1,
+        use_batchnorm=True,
+    ):
+
+        if use_batchnorm == "inplace" and InPlaceABN is None:
+            raise RuntimeError(
+                "In order to use `use_batchnorm='inplace'` inplace_abn package must be installed. "
+                + "To install see: https://github.com/mapillary/inplace_abn"
+            )
+
+        conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=not (use_batchnorm),
+        )
+        relu = nn.ReLU(inplace=True)
+
+        if use_batchnorm == "inplace":
+            bn = InPlaceABN(out_channels, activation="leaky_relu", activation_param=0.0)
+            relu = nn.Identity()
+
+        elif use_batchnorm and use_batchnorm != "inplace":
+            bn = nn.BatchNorm2d(out_channels)
+
+        else:
+            bn = nn.Identity()
+
+        super(Conv2dReLU, self).__init__(conv, bn, relu)
+
+
+class SCSEModule(nn.Module):
+    def __init__(self, in_channels, reduction=16):
+        super().__init__()
+        self.cSE = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(in_channels, in_channels // reduction, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_channels // reduction, in_channels, 1),
+            nn.Sigmoid(),
+        )
+        self.sSE = nn.Sequential(nn.Conv2d(in_channels, 1, 1), nn.Sigmoid())
+
+    def forward(self, x):
+        return x * self.cSE(x) + x * self.sSE(x)
+
+
+class ArgMax(nn.Module):
+    def __init__(self, dim=None):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        return torch.argmax(x, dim=self.dim)
+
+
+class Clamp(nn.Module):
+    def __init__(self, min=0, max=1):
+        super().__init__()
+        self.min, self.max = min, max
+
+    def forward(self, x):
+        return torch.clamp(x, self.min, self.max)
+
+
+class Activation(nn.Module):
+    def __init__(self, name, **params):
+
+        super().__init__()
+
+        if name is None or name == "identity":
+            self.activation = nn.Identity(**params)
+        elif name == "sigmoid":
+            self.activation = nn.Sigmoid()
+        elif name == "softmax2d":
+            self.activation = nn.Softmax(dim=1, **params)
+        elif name == "softmax":
+            self.activation = nn.Softmax(**params)
+        elif name == "logsoftmax":
+            self.activation = nn.LogSoftmax(**params)
+        elif name == "tanh":
+            self.activation = nn.Tanh()
+        elif name == "argmax":
+            self.activation = ArgMax(**params)
+        elif name == "argmax2d":
+            self.activation = ArgMax(dim=1, **params)
+        elif name == "clamp":
+            self.activation = Clamp(**params)
+        elif callable(name):
+            self.activation = name(**params)
+        else:
+            raise ValueError(
+                f"Activation should be callable/sigmoid/softmax/logsoftmax/tanh/"
+                f"argmax/argmax2d/clamp/None; got {name}"
+            )
+
+    def forward(self, x):
+        return self.activation(x)
+
+
+class Attention(nn.Module):
+    def __init__(self, name, **params):
+        super().__init__()
+
+        if name is None:
+            self.attention = nn.Identity(**params)
+        elif name == "scse":
+            self.attention = SCSEModule(**params)
+        else:
+            raise ValueError("Attention {} is not implemented".format(name))
+
+    def forward(self, x):
+        return self.attention(x)

segmentation_models_pytorch/datasets/__init__.py

@@ -0,0 +1 @@
+from .oxford_pet import OxfordPetDataset, SimpleOxfordPetDataset

segmentation_models_pytorch/datasets/oxford_pet.py

@@ -0,0 +1,136 @@
+import os
+import torch
+import shutil
+import numpy as np
+
+from PIL import Image
+from tqdm import tqdm
+from urllib.request import urlretrieve
+
+
+class OxfordPetDataset(torch.utils.data.Dataset):
+    def __init__(self, root, mode="train", transform=None):
+
+        assert mode in {"train", "valid", "test"}
+
+        self.root = root
+        self.mode = mode
+        self.transform = transform
+
+        self.images_directory = os.path.join(self.root, "images")
+        self.masks_directory = os.path.join(self.root, "annotations", "trimaps")
+
+        self.filenames = self._read_split()  # read train/valid/test splits
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, idx):
+
+        filename = self.filenames[idx]
+        image_path = os.path.join(self.images_directory, filename + ".jpg")
+        mask_path = os.path.join(self.masks_directory, filename + ".png")
+
+        image = np.array(Image.open(image_path).convert("RGB"))
+
+        trimap = np.array(Image.open(mask_path))
+        mask = self._preprocess_mask(trimap)
+
+        sample = dict(image=image, mask=mask, trimap=trimap)
+        if self.transform is not None:
+            sample = self.transform(**sample)
+
+        return sample
+
+    @staticmethod
+    def _preprocess_mask(mask):
+        mask = mask.astype(np.float32)
+        mask[mask == 2.0] = 0.0
+        mask[(mask == 1.0) | (mask == 3.0)] = 1.0
+        return mask
+
+    def _read_split(self):
+        split_filename = "test.txt" if self.mode == "test" else "trainval.txt"
+        split_filepath = os.path.join(self.root, "annotations", split_filename)
+        with open(split_filepath) as f:
+            split_data = f.read().strip("\n").split("\n")
+        filenames = [x.split(" ")[0] for x in split_data]
+        if self.mode == "train":  # 90% for train
+            filenames = [x for i, x in enumerate(filenames) if i % 10 != 0]
+        elif self.mode == "valid":  # 10% for validation
+            filenames = [x for i, x in enumerate(filenames) if i % 10 == 0]
+        return filenames
+
+    @staticmethod
+    def download(root):
+
+        # load images
+        filepath = os.path.join(root, "images.tar.gz")
+        download_url(
+            url="https://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz",
+            filepath=filepath,
+        )
+        extract_archive(filepath)
+
+        # load annotations
+        filepath = os.path.join(root, "annotations.tar.gz")
+        download_url(
+            url="https://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz",
+            filepath=filepath,
+        )
+        extract_archive(filepath)
+
+
+class SimpleOxfordPetDataset(OxfordPetDataset):
+    def __getitem__(self, *args, **kwargs):
+
+        sample = super().__getitem__(*args, **kwargs)
+
+        # resize images
+        image = np.array(
+            Image.fromarray(sample["image"]).resize((256, 256), Image.LINEAR)
+        )
+        mask = np.array(
+            Image.fromarray(sample["mask"]).resize((256, 256), Image.NEAREST)
+        )
+        trimap = np.array(
+            Image.fromarray(sample["trimap"]).resize((256, 256), Image.NEAREST)
+        )
+
+        # convert to other format HWC -> CHW
+        sample["image"] = np.moveaxis(image, -1, 0)
+        sample["mask"] = np.expand_dims(mask, 0)
+        sample["trimap"] = np.expand_dims(trimap, 0)
+
+        return sample
+
+
+class TqdmUpTo(tqdm):
+    def update_to(self, b=1, bsize=1, tsize=None):
+        if tsize is not None:
+            self.total = tsize
+        self.update(b * bsize - self.n)
+
+
+def download_url(url, filepath):
+    directory = os.path.dirname(os.path.abspath(filepath))
+    os.makedirs(directory, exist_ok=True)
+    if os.path.exists(filepath):
+        return
+
+    with TqdmUpTo(
+        unit="B",
+        unit_scale=True,
+        unit_divisor=1024,
+        miniters=1,
+        desc=os.path.basename(filepath),
+    ) as t:
+        urlretrieve(url, filename=filepath, reporthook=t.update_to, data=None)
+        t.total = t.n
+
+
+def extract_archive(filepath):
+    extract_dir = os.path.dirname(os.path.abspath(filepath))
+    dst_dir = os.path.splitext(filepath)[0]
+    if not os.path.exists(dst_dir):
+        shutil.unpack_archive(filepath, extract_dir)

segmentation_models_pytorch/decoders/deeplabv3/__init__.py

@@ -0,0 +1 @@
+from .model import DeepLabV3, DeepLabV3Plus

segmentation_models_pytorch/decoders/deeplabv3/decoder.py

@@ -0,0 +1,220 @@
+"""
+BSD 3-Clause License
+
+Copyright (c) Soumith Chintala 2016,
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+__all__ = ["DeepLabV3Decoder"]
+
+
+class DeepLabV3Decoder(nn.Sequential):
+    def __init__(self, in_channels, out_channels=256, atrous_rates=(12, 24, 36)):
+        super().__init__(
+            ASPP(in_channels, out_channels, atrous_rates),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+        )
+        self.out_channels = out_channels
+
+    def forward(self, *features):
+        return super().forward(features[-1])
+
+
+class DeepLabV3PlusDecoder(nn.Module):
+    def __init__(
+        self,
+        encoder_channels,
+        out_channels=256,
+        atrous_rates=(12, 24, 36),
+        output_stride=16,
+    ):
+        super().__init__()
+        if output_stride not in {8, 16}:
+            raise ValueError(
+                "Output stride should be 8 or 16, got {}.".format(output_stride)
+            )
+
+        self.out_channels = out_channels
+        self.output_stride = output_stride
+
+        self.aspp = nn.Sequential(
+            ASPP(encoder_channels[-1], out_channels, atrous_rates, separable=True),
+            SeparableConv2d(
+                out_channels, out_channels, kernel_size=3, padding=1, bias=False
+            ),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+        )
+
+        scale_factor = 2 if output_stride == 8 else 4
+        self.up = nn.UpsamplingBilinear2d(scale_factor=scale_factor)
+
+        highres_in_channels = encoder_channels[-4]
+        highres_out_channels = 48  # proposed by authors of paper
+        self.block1 = nn.Sequential(
+            nn.Conv2d(
+                highres_in_channels, highres_out_channels, kernel_size=1, bias=False
+            ),
+            nn.BatchNorm2d(highres_out_channels),
+            nn.ReLU(),
+        )
+        self.block2 = nn.Sequential(
+            SeparableConv2d(
+                highres_out_channels + out_channels,
+                out_channels,
+                kernel_size=3,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+        )
+
+    def forward(self, *features):
+        aspp_features = self.aspp(features[-1])
+        aspp_features = self.up(aspp_features)
+        high_res_features = self.block1(features[-4])
+        concat_features = torch.cat([aspp_features, high_res_features], dim=1)
+        fused_features = self.block2(concat_features)
+        return fused_features
+
+
+class ASPPConv(nn.Sequential):
+    def __init__(self, in_channels, out_channels, dilation):
+        super().__init__(
+            nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size=3,
+                padding=dilation,
+                dilation=dilation,
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+        )
+
+
+class ASPPSeparableConv(nn.Sequential):
+    def __init__(self, in_channels, out_channels, dilation):
+        super().__init__(
+            SeparableConv2d(
+                in_channels,
+                out_channels,
+                kernel_size=3,
+                padding=dilation,
+                dilation=dilation,
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+        )
+
+
+class ASPPPooling(nn.Sequential):
+    def __init__(self, in_channels, out_channels):
+        super().__init__(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+        )
+
+    def forward(self, x):
+        size = x.shape[-2:]
+        for mod in self:
+            x = mod(x)
+        return F.interpolate(x, size=size, mode="bilinear", align_corners=False)
+
+
+class ASPP(nn.Module):
+    def __init__(self, in_channels, out_channels, atrous_rates, separable=False):
+        super(ASPP, self).__init__()
+        modules = []
+        modules.append(
+            nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, 1, bias=False),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(),
+            )
+        )
+
+        rate1, rate2, rate3 = tuple(atrous_rates)
+        ASPPConvModule = ASPPConv if not separable else ASPPSeparableConv
+
+        modules.append(ASPPConvModule(in_channels, out_channels, rate1))
+        modules.append(ASPPConvModule(in_channels, out_channels, rate2))
+        modules.append(ASPPConvModule(in_channels, out_channels, rate3))
+        modules.append(ASPPPooling(in_channels, out_channels))
+
+        self.convs = nn.ModuleList(modules)
+
+        self.project = nn.Sequential(
+            nn.Conv2d(5 * out_channels, out_channels, kernel_size=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+        )
+
+    def forward(self, x):
+        res = []
+        for conv in self.convs:
+            res.append(conv(x))
+        res = torch.cat(res, dim=1)
+        return self.project(res)
+
+
+class SeparableConv2d(nn.Sequential):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        bias=True,
+    ):
+        dephtwise_conv = nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=in_channels,
+            bias=False,
+        )
+        pointwise_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias,)
+        super().__init__(dephtwise_conv, pointwise_conv)

segmentation_models_pytorch/decoders/deeplabv3/model.py

@@ -0,0 +1,179 @@
+from torch import nn
+from typing import Optional
+
+from segmentation_models_pytorch.base import (
+    SegmentationModel,
+    SegmentationHead,
+    ClassificationHead,
+)
+from segmentation_models_pytorch.encoders import get_encoder
+from .decoder import DeepLabV3Decoder, DeepLabV3PlusDecoder
+
+
+class DeepLabV3(SegmentationModel):
+    """DeepLabV3_ implementation from "Rethinking Atrous Convolution for Semantic Image Segmentation"
+
+    Args:
+        encoder_name: Name of the classification model that will be used as an encoder (a.k.a backbone)
+            to extract features of different spatial resolution
+        encoder_depth: A number of stages used in encoder in range [3, 5]. Each stage generate features
+            two times smaller in spatial dimensions than previous one (e.g. for depth 0 we will have features
+            with shapes [(N, C, H, W),], for depth 1 - [(N, C, H, W), (N, C, H // 2, W // 2)] and so on).
+            Default is 5
+        encoder_weights: One of **None** (random initialization), **"imagenet"** (pre-training on ImageNet) and
+            other pretrained weights (see table with available weights for each encoder_name)
+        decoder_channels: A number of convolution filters in ASPP module. Default is 256
+        in_channels: A number of input channels for the model, default is 3 (RGB images)
+        classes: A number of classes for output mask (or you can think as a number of channels of output mask)
+        activation: An activation function to apply after the final convolution layer.
+            Available options are **"sigmoid"**, **"softmax"**, **"logsoftmax"**, **"tanh"**, **"identity"**,
+                **callable** and **None**.
+            Default is **None**
+        upsampling: Final upsampling factor. Default is 8 to preserve input-output spatial shape identity
+        aux_params: Dictionary with parameters of the auxiliary output (classification head). Auxiliary output is build
+            on top of encoder if **aux_params** is not **None** (default). Supported params:
+                - classes (int): A number of classes
+                - pooling (str): One of "max", "avg". Default is "avg"
+                - dropout (float): Dropout factor in [0, 1)
+                - activation (str): An activation function to apply "sigmoid"/"softmax"
+                    (could be **None** to return logits)
+    Returns:
+        ``torch.nn.Module``: **DeepLabV3**
+
+    .. _DeeplabV3:
+        https://arxiv.org/abs/1706.05587
+
+    """
+
+    def __init__(
+        self,
+        encoder_name: str = "resnet34",
+        encoder_depth: int = 5,
+        encoder_weights: Optional[str] = "imagenet",
+        decoder_channels: int = 256,
+        in_channels: int = 3,
+        classes: int = 1,
+        activation: Optional[str] = None,
+        upsampling: int = 8,
+        aux_params: Optional[dict] = None,
+    ):
+        super().__init__()
+
+        self.encoder = get_encoder(
+            encoder_name,
+            in_channels=in_channels,
+            depth=encoder_depth,
+            weights=encoder_weights,
+            output_stride=8,
+        )
+
+        self.decoder = DeepLabV3Decoder(
+            in_channels=self.encoder.out_channels[-1], out_channels=decoder_channels,
+        )
+
+        self.segmentation_head = SegmentationHead(
+            in_channels=self.decoder.out_channels,
+            out_channels=classes,
+            activation=activation,
+            kernel_size=1,
+            upsampling=upsampling,
+        )
+
+        if aux_params is not None:
+            self.classification_head = ClassificationHead(
+                in_channels=self.encoder.out_channels[-1], **aux_params
+            )
+        else:
+            self.classification_head = None
+
+
+class DeepLabV3Plus(SegmentationModel):
+    """DeepLabV3+ implementation from "Encoder-Decoder with Atrous Separable
+    Convolution for Semantic Image Segmentation"
+
+    Args:
+        encoder_name: Name of the classification model that will be used as an encoder (a.k.a backbone)
+            to extract features of different spatial resolution
+        encoder_depth: A number of stages used in encoder in range [3, 5]. Each stage generate features
+            two times smaller in spatial dimensions than previous one (e.g. for depth 0 we will have features
+            with shapes [(N, C, H, W),], for depth 1 - [(N, C, H, W), (N, C, H // 2, W // 2)] and so on).
+            Default is 5
+        encoder_weights: One of **None** (random initialization), **"imagenet"** (pre-training on ImageNet) and
+            other pretrained weights (see table with available weights for each encoder_name)
+        encoder_output_stride: Downsampling factor for last encoder features (see original paper for explanation)
+        decoder_atrous_rates: Dilation rates for ASPP module (should be a tuple of 3 integer values)
+        decoder_channels: A number of convolution filters in ASPP module. Default is 256
+        in_channels: A number of input channels for the model, default is 3 (RGB images)
+        classes: A number of classes for output mask (or you can think as a number of channels of output mask)
+        activation: An activation function to apply after the final convolution layer.
+            Available options are **"sigmoid"**, **"softmax"**, **"logsoftmax"**, **"tanh"**, **"identity"**,
+                **callable** and **None**.
+            Default is **None**
+        upsampling: Final upsampling factor. Default is 4 to preserve input-output spatial shape identity
+        aux_params: Dictionary with parameters of the auxiliary output (classification head). Auxiliary output is build
+            on top of encoder if **aux_params** is not **None** (default). Supported params:
+                - classes (int): A number of classes
+                - pooling (str): One of "max", "avg". Default is "avg"
+                - dropout (float): Dropout factor in [0, 1)
+                - activation (str): An activation function to apply "sigmoid"/"softmax"
+                    (could be **None** to return logits)
+    Returns:
+        ``torch.nn.Module``: **DeepLabV3Plus**
+
+    Reference:
+        https://arxiv.org/abs/1802.02611v3
+
+    """
+
+    def __init__(
+        self,
+        encoder_name: str = "resnet34",
+        encoder_depth: int = 5,
+        encoder_weights: Optional[str] = "imagenet",
+        encoder_output_stride: int = 16,
+        decoder_channels: int = 256,
+        decoder_atrous_rates: tuple = (12, 24, 36),
+        in_channels: int = 3,
+        classes: int = 1,
+        activation: Optional[str] = None,
+        upsampling: int = 4,
+        aux_params: Optional[dict] = None,
+    ):
+        super().__init__()
+
+        if encoder_output_stride not in [8, 16]:
+            raise ValueError(
+                "Encoder output stride should be 8 or 16, got {}".format(
+                    encoder_output_stride
+                )
+            )
+
+        self.encoder = get_encoder(
+            encoder_name,
+            in_channels=in_channels,
+            depth=encoder_depth,
+            weights=encoder_weights,
+            output_stride=encoder_output_stride,
+        )
+
+        self.decoder = DeepLabV3PlusDecoder(
+            encoder_channels=self.encoder.out_channels,
+            out_channels=decoder_channels,
+            atrous_rates=decoder_atrous_rates,
+            output_stride=encoder_output_stride,
+        )
+
+        self.segmentation_head = SegmentationHead(
+            in_channels=self.decoder.out_channels,
+            out_channels=classes,
+            activation=activation,
+            kernel_size=1,
+            upsampling=upsampling,
+        )
+
+        if aux_params is not None:
+            self.classification_head = ClassificationHead(
+                in_channels=self.encoder.out_channels[-1], **aux_params
+            )
+        else:
+            self.classification_head = None

segmentation_models_pytorch/decoders/fpn/__init__.py

@@ -0,0 +1 @@
+from .model import FPN