Upload 24 files

README.md

@@ -1,3 +1,72 @@
----
-license: apache-2.0
----
+# Solution of Team Sribd-med for NeurIPS-CellSeg Challenge
+This repository provides the solution of team Sribd-med for [NeurIPS-CellSeg](https://neurips22-cellseg.grand-challenge.org/) 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](https://github.com/JunMa11/NeurIPS-CellSeg) repository,
+
+You can reproduce our method as follows step by step:
+
+## Environments and Requirements:
+Install requirements by
+
+```shell
+python -m pip install -r requirements.txt
+```
+
+## Dataset
+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 
+
+## Automatic cell classification
+You can classify the cells into four classes in this step.
+Put all the images (competition + Cellpose + Omnipose + Sartorius) in one folder (data/allimages).
+Run classification code:
+
+```shell
+python classification/unsup_classification.py
+```
+The results can be stored in data/classification_results/
+
+## CNN-base classification model training
+Using the classified images in data/classification_results/. A resnet18 is trained:
+```shell
+python classification/train_classification.py
+```
+## Segmentation Training
+Pre-training convnext-stardist using all the images (data/allimages).
+```shell
+python train_convnext_stardist.py
+```
+For class 0,2,3 finetune on the classified data (Take class1 as a example):
+```shell
+python finetune_convnext_stardist.py model_dir=(The pretrained convnext-stardist model) data_dir='data/classification_results/class1'
+```
+For class 1 train the convnext-hover from scratch using classified class 3 data.
+```shell
+python train_convnext_hover.py data_dir='data/classification_results/class3'
+```
+
+Finally, four segmentation models will be trained.
+
+## Trained models
+The models can be downloaded from this link:
+https://drive.google.com/drive/folders/1MkEOpgmdkg5Yqw6Ng5PoOhtmo9xPPwIj?usp=sharing
+
+## Inference
+The inference process includes classification and segmentation.
+```shell
+python predict.py -i input_path -o output_path --model_path './models' 
+```
+
+## Evaluation
+Calculate the F-score for evaluation:
+```shell
+python compute_metric.py --gt_path path_to_labels --seg_path output_path
+```
+
+## Results
+The tuning set F1 score of our method is 0.8795. The rank running time of our method on all the 101 cases in the tuning set is zero in our local
+workstation. 
+## Acknowledgement
+We thank for the contributors of public datasets.
+

cellseg_time_eval.py

@@ -0,0 +1,97 @@
+"""
+The code was adapted from the MICCAI FLARE Challenge
+https://flare22.grand-challenge.org/
+
+The testing images will be evaluated one by one. 
+To compensate for the Docker container startup time, we give a time tolerance for the running time. 
+https://neurips22-cellseg.grand-challenge.org/metrics/
+"""
+
+import os
+join = os.path.join
+import sys
+import shutil
+import time
+import torch
+import argparse
+from collections import OrderedDict
+from skimage import io 
+import tifffile as tif 
+import numpy as np
+import pandas as pd
+
+parser = argparse.ArgumentParser('Segmentation efficiency eavluation for docker containers', add_help=False)
+parser.add_argument('-i', '--test_img_path', default='./val-imgs-30/', type=str, help='testing data path')
+parser.add_argument('-o','--save_path', default='./val_team_seg', type=str, help='segmentation output path')
+parser.add_argument('-d','--docker_folder_path', default='./team_docker', type=str, help='team docker path')
+args = parser.parse_args()
+
+test_img_path = args.test_img_path
+save_path = args.save_path
+docker_path = args.docker_folder_path
+
+input_temp = './inputs/'
+output_temp = './outputs'
+os.makedirs(save_path, exist_ok=True)
+
+dockers = sorted(os.listdir(docker_path))
+test_cases = sorted(os.listdir(test_img_path))
+
+for docker in dockers:
+    try:
+        # create temp folers for inference one-by-one
+        if os.path.exists(input_temp):
+            shutil.rmtree(input_temp)
+        if os.path.exists(output_temp):
+            shutil.rmtree(output_temp)
+        os.makedirs(input_temp)
+        os.makedirs(output_temp)
+        # load docker and create a new folder to save segmentation results
+        teamname = docker.split('.')[0].lower()
+        print('teamname docker: ', docker)
+        # os.system('docker image load < {}'.format(join(docker_path, docker)))
+        team_outpath = join(save_path, teamname)
+        if os.path.exists(team_outpath):
+            shutil.rmtree(team_outpath)
+        os.mkdir(team_outpath)
+        metric = OrderedDict()
+        metric['Img Name'] = []
+        metric['Real Running Time'] = []
+        metric['Rank Running Time'] = []
+        # To obtain the running time for each case, we inference the testing case one-by-one
+        for case in test_cases:
+            shutil.copy(join(test_img_path, case), input_temp)
+            if case.endswith('.tif') or case.endswith('.tiff'):
+                img = tif.imread(join(input_temp, case))
+            else:
+                img = io.imread(join(input_temp, case))
+            pix_num = img.shape[0] * img.shape[1]
+            cmd = 'docker container run --gpus="device=0" -m 28g --name {} --rm -v $PWD/inputs/:/workspace/inputs/ -v $PWD/outputs/:/workspace/outputs/ {}:latest /bin/bash -c "sh predict.sh" '.format(teamname, teamname)
+            print(teamname, ' docker command:', cmd, '\n', 'testing image name:', case)
+            start_time = time.time()
+            os.system(cmd)
+            real_running_time = time.time() - start_time
+            print(f"{case} finished! Inference time: {real_running_time}")
+            # save metrics
+            metric['Img Name'].append(case)
+            metric['Real Running Time'].append(real_running_time)
+            if pix_num <= 1000000:
+                rank_running_time = np.max([0, real_running_time-10])
+            else:
+                rank_running_time = np.max([0, real_running_time-10*(pix_num/1000000)])
+            metric['Rank Running Time'].append(rank_running_time)
+            os.remove(join(input_temp, case))  
+            seg_name = case.split('.')[0] + '_label.tiff'
+            try:
+                os.rename(join(output_temp, seg_name), join(team_outpath, seg_name))
+            except:
+                print(f"{join(output_temp, seg_name)}, {join(team_outpath, seg_name)}")
+                print("Wrong segmentation name!!! It should be image_name.split(\'.\')[0] + \'_label.tiff\' ")
+        metric_df = pd.DataFrame(metric)
+        metric_df.to_csv(join(team_outpath, teamname + '_running_time.csv'), index=False)
+        torch.cuda.empty_cache()
+        # os.system("docker rmi {}:latest".format(teamname))
+        shutil.rmtree(input_temp)
+        shutil.rmtree(output_temp)
+    except Exception as e:
+        print(e)

classification/train_classification.py

@@ -0,0 +1,181 @@
+import os, sys, glob, time, random, shutil, copy
+from tqdm import tqdm
+import numpy as np
+import pandas as pd
+import torch
+import torchvision
+from torchvision import datasets, models, transforms
+import torch.utils.data as data
+import torch.nn as nn
+import torch.optim as optim
+from torch.optim import lr_scheduler
+import torch.nn.functional as F
+from torchsummary import summary
+from matplotlib import pyplot as plt
+from torchvision.models import resnet18, ResNet18_Weights # do not import
+from PIL import Image, ImageFile
+from skimage import io
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+# Set the train and validation directory paths
+train_directory = 'dataset/train'
+valid_directory = 'dataset/val'
+
+# Batch size
+bs = 64 
+# Number of epochs
+num_epochs = 20
+# Number of classes
+num_classes = 4
+# Number of workers
+num_cpu = 8 
+
+# Applying transforms to the data
+image_transforms = { 
+    'train': transforms.Compose([
+        transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),
+        transforms.RandomRotation(degrees=15),
+        transforms.RandomHorizontalFlip(),
+        transforms.CenterCrop(size=224),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406],
+                             [0.229, 0.224, 0.225])
+    ]),
+    'valid': 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])
+    ])
+}
+ 
+# Load data from folders
+dataset = {
+    'train': datasets.ImageFolder(root=train_directory, transform=image_transforms['train']),
+    'valid': datasets.ImageFolder(root=valid_directory, transform=image_transforms['valid'])
+}
+ 
+# Size of train and validation data
+dataset_sizes = {
+    'train':len(dataset['train']),
+    'valid':len(dataset['valid'])
+}
+
+# Create iterators for data loading
+dataloaders = {
+    'train':data.DataLoader(dataset['train'], batch_size=bs, shuffle=True,
+                            num_workers=num_cpu, pin_memory=True, drop_last=False),
+    'valid':data.DataLoader(dataset['valid'], batch_size=bs, shuffle=False,
+                            num_workers=num_cpu, pin_memory=True, drop_last=False)
+}
+
+# Class names or target labels
+class_names = dataset['train'].classes
+print("Classes:", class_names)
+ 
+# Print the train and validation data sizes
+print("Training-set size:",dataset_sizes['train'],
+      "\nValidation-set size:", dataset_sizes['valid'])
+
+modelname = 'resnet18'
+
+# Set default device as gpu, if available
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+weights = ResNet18_Weights.DEFAULT
+model = resnet18(weights=None)    
+num_ftrs = model.fc.in_features
+model.fc = nn.Linear(num_ftrs, num_classes)
+
+
+# Transfer the model to GPU
+model = model.to(device)
+
+# Print model summary
+print('Model Summary:-\n')
+for num, (name, param) in enumerate(model.named_parameters()):
+    print(num, name, param.requires_grad )
+summary(model, input_size=(3, 224, 224))
+
+# Loss function
+criterion = nn.CrossEntropyLoss()
+
+# Optimizer 
+optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
+
+# Learning rate decay
+scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
+
+since = time.time()
+
+best_model_wts = copy.deepcopy(model.state_dict())
+best_acc = 0.0
+
+for epoch in range(1, num_epochs+1):
+    print('Epoch {}/{}'.format(epoch, num_epochs))
+    print('-' * 10)
+
+    # Each epoch has a training and validation phase
+    for phase in ['train', 'valid']:
+        if phase == 'train':
+            model.train()  # Set model to training mode
+        else:
+            model.eval()   # Set model to evaluate mode
+
+        running_loss = 0.0
+        running_corrects = 0
+
+        # Iterate over data.
+        n = 0
+        stream = tqdm(dataloaders[phase])
+        for i, (inputs, labels) in enumerate(stream, start=1):
+            inputs = inputs.to(device)
+            labels = labels.to(device)
+
+            # zero the parameter gradients
+            optimizer.zero_grad()
+
+            # forward
+            # track history if only in train
+            with torch.set_grad_enabled(phase == 'train'):
+                outputs = model(inputs)
+                _, preds = torch.max(outputs, 1)
+                loss = criterion(outputs, labels)
+
+                # backward + optimize only if in training phase
+                if phase == 'train':
+                    loss.backward()
+                    optimizer.step()
+
+            # statistics
+            n += inputs.shape[0]
+            running_loss += loss.item() * inputs.size(0)
+            running_corrects += torch.sum(preds == labels.data)
+            
+            stream.set_description(f'Batch {i}/{len(dataloaders[phase])} | Loss: {running_loss/n:.4f}, Acc: {running_corrects/n:.4f}')
+
+        if phase == 'train':
+            scheduler.step()
+
+        epoch_loss = running_loss / dataset_sizes[phase]
+        epoch_acc = running_corrects.double() / dataset_sizes[phase]
+
+        print('Epoch {}-{} Loss: {:.4f} Acc: {:.4f}'.format(
+            epoch, phase, epoch_loss, epoch_acc))
+
+        # deep copy the model
+        if phase == 'valid' and epoch_acc >= best_acc:
+            best_acc = epoch_acc
+            best_model_wts = copy.deepcopy(model.state_dict())
+            print('Update best model!')
+
+time_elapsed = time.time() - since
+print('Training complete in {:.0f}m {:.0f}s'.format(
+    time_elapsed // 60, time_elapsed % 60))
+print('Best val Acc: {:4f}'.format(best_acc))
+
+# load best model weights
+model.load_state_dict(best_model_wts)
+torch.save(model, 'logs/resnet18_4class.pth')
+torch.save(model.state_dict(), 'logs/resnet18_4class.tar')

classification/unsup_classification.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+import os
+import numpy as np
+import shutil
+import torch
+import torch.nn
+import torchvision.models as models
+from torch.autograd import Variable 
+import torch.cuda
+import torchvision.transforms as transforms
+from PIL import Image
+import torch.nn as nn
+import torch.nn.functional as F
+from sklearn.datasets import make_blobs
+from sklearn.cluster import KMeans
+from sklearn.metrics import silhouette_score
+from sklearn.preprocessing import StandardScaler
+from sklearn.metrics import pairwise_distances_argmin_min
+from scipy.spatial.distance import pdist, squareform
+from skimage import io, segmentation, morphology, exposure
+from skimage.color import rgb2hsv
+img_to_tensor = transforms.ToTensor()
+import random
+import tifffile as tif
+path = '/data1/partitionA/CUHKSZ/histopath_2022/grand_competition/Train_Labeled/images/'
+files = os.listdir(path)
+binary_path = '0/'
+gray_path = '1/'
+colored_path = 'colored/'
+os.makedirs(binary_path, exist_ok=True)
+os.makedirs(colored_path, exist_ok=True)
+os.makedirs(gray_path, exist_ok=True)
+for img_name in files:
+    img_path = path + str(img_name)
+    if img_name.endswith('.tif') or img_name.endswith('.tiff'):
+        img_data = tif.imread(img_path)
+    else:
+        img_data = io.imread(img_path)
+    if len(img_data.shape) == 2 or (len(img_data.shape) == 3 and img_data.shape[-1] == 1):
+        shutil.copyfile(path + img_name, binary_path + img_name)
+    elif len(img_data.shape) == 3 and img_data.shape[-1] > 3:
+        shutil.copyfile(path + img_name, colored_path + img_name)
+    else:
+        hsv_img = rgb2hsv(img_data)
+        s = hsv_img[:,:,1]
+        v = hsv_img[:,:,2]
+        print(img_name,s.mean(),v.mean())
+        if s.mean() > 0.1 or (v.mean()<0.1 or v.mean() > 0.6):
+            shutil.copyfile(path + img_name, colored_path + img_name)
+        else:
+            shutil.copyfile(path + img_name, gray_path + img_name)
+
+
+
+# In[3]:
+
+
+####Phrase 2 clustering by cell size
+from skimage import measure
+colored_path = 'colored/'
+label_path = 'allimages/tif/'
+big_path = '2/'
+small_path = '3/'
+files = os.listdir(colored_path)
+os.makedirs(big_path, exist_ok=True)
+os.makedirs(small_path, exist_ok=True)
+for img_name in files:
+    label =  tif.imread(label_path + img_name.split('.')[0]+'.tif')
+    props = measure.regionprops(label)
+    num_pix = []
+    for idx in range(len(props)):
+        num_pix.append(props[idx].area)
+    max_area = max(num_pix)
+    print(max_area)
+    if max_area > 30000:
+        shutil.copyfile(path + img_name, big_path + img_name)
+    else:
+        shutil.copyfile(path + img_name, small_path + img_name)
+
+
+
+
+
+
+
+

compute_metric.py

@@ -0,0 +1,201 @@
+"""
+Created on Thu Mar 31 18:10:52 2022
+adapted form https://github.com/stardist/stardist/blob/master/stardist/matching.py
+Thanks the authors of Stardist for sharing the great code
+
+"""
+
+import argparse
+import numpy as np
+from numba import jit
+from scipy.optimize import linear_sum_assignment
+from collections import OrderedDict
+import pandas as pd
+from skimage import segmentation
+import tifffile as tif
+import os
+join = os.path.join
+from tqdm import tqdm
+
+def _intersection_over_union(masks_true, masks_pred):
+    """ intersection over union of all mask pairs
+    
+    Parameters
+    ------------
+    
+    masks_true: ND-array, int 
+        ground truth masks, where 0=NO masks; 1,2... are mask labels
+    masks_pred: ND-array, int
+        predicted masks, where 0=NO masks; 1,2... are mask labels
+    """
+    overlap = _label_overlap(masks_true, masks_pred)
+    n_pixels_pred = np.sum(overlap, axis=0, keepdims=True)
+    n_pixels_true = np.sum(overlap, axis=1, keepdims=True)
+    iou = overlap / (n_pixels_pred + n_pixels_true - overlap)
+    iou[np.isnan(iou)] = 0.0
+    return iou
+
+@jit(nopython=True)
+def _label_overlap(x, y):
+    """ fast function to get pixel overlaps between masks in x and y 
+    
+    Parameters
+    ------------
+
+    x: ND-array, int
+        where 0=NO masks; 1,2... are mask labels
+    y: ND-array, int
+        where 0=NO masks; 1,2... are mask labels
+
+    Returns
+    ------------
+
+    overlap: ND-array, int
+        matrix of pixel overlaps of size [x.max()+1, y.max()+1]
+    
+    """
+    x = x.ravel()
+    y = y.ravel()
+    
+    # preallocate a 'contact map' matrix
+    overlap = np.zeros((1+x.max(),1+y.max()), dtype=np.uint)
+    
+    # loop over the labels in x and add to the corresponding
+    # overlap entry. If label A in x and label B in y share P
+    # pixels, then the resulting overlap is P
+    # len(x)=len(y), the number of pixels in the whole image 
+    for i in range(len(x)):
+        overlap[x[i],y[i]] += 1
+    return overlap
+
+def _true_positive(iou, th):
+    """ true positive at threshold th
+    
+    Parameters
+    ------------
+
+    iou: float, ND-array
+        array of IOU pairs
+    th: float
+        threshold on IOU for positive label
+
+    Returns
+    ------------
+
+    tp: float
+        number of true positives at threshold
+    """
+    n_min = min(iou.shape[0], iou.shape[1])
+    costs = -(iou >= th).astype(float) - iou / (2*n_min)
+    true_ind, pred_ind = linear_sum_assignment(costs)
+    match_ok = iou[true_ind, pred_ind] >= th
+    tp = match_ok.sum()
+    return tp
+
+def eval_tp_fp_fn(masks_true, masks_pred, threshold=0.5):
+    num_inst_gt = np.max(masks_true)
+    num_inst_seg = np.max(masks_pred)
+    if num_inst_seg>0:
+        iou = _intersection_over_union(masks_true, masks_pred)[1:, 1:]
+            # for k,th in enumerate(threshold):
+        tp = _true_positive(iou, threshold)
+        fp = num_inst_seg - tp
+        fn = num_inst_gt - tp
+    else:
+        print('No segmentation results!')
+        tp = 0
+        fp = 0
+        fn = 0
+        
+    return tp, fp, fn
+
+def remove_boundary_cells(mask):
+    W, H = mask.shape
+    bd = np.ones((W, H))
+    bd[2:W-2, 2:H-2] = 0
+    bd_cells = np.unique(mask*bd)
+    for i in bd_cells[1:]:
+        mask[mask==i] = 0
+    new_label,_,_ = segmentation.relabel_sequential(mask)
+    return new_label
+
+def main():
+    parser = argparse.ArgumentParser('Compute F1 score for cell segmentation results', add_help=False)
+    # Dataset parameters
+    parser.add_argument('--gt_path', type=str, help='path to ground truth; file names end with _label.tiff', required=True)
+    parser.add_argument('--seg_path', type=str, help='path to segmentation results; file names are the same as ground truth', required=True)
+    parser.add_argument('--save_path', default='./', help='path where to save metrics')
+    args = parser.parse_args()
+
+    gt_path = args.gt_path
+    seg_path = args.seg_path
+    names = sorted(os.listdir(seg_path))
+    seg_metric = OrderedDict()
+    seg_metric['Names'] = []
+    seg_metric['F1_Score'] = []
+    for name in tqdm(names):
+        assert name.endswith('_label.tiff'), 'The suffix of label name should be _label.tiff'
+        
+        # Load the images for this case
+        gt = tif.imread(join(gt_path, name))
+        seg = tif.imread(join(seg_path, name))
+        
+        # Score the cases
+        # do not consider cells on the boundaries during evaluation
+        if np.prod(gt.shape)<25000000:
+            gt = remove_boundary_cells(gt.astype(np.int32)) 
+            seg = remove_boundary_cells(seg.astype(np.int32))           
+            tp, fp, fn = eval_tp_fp_fn(gt, seg, threshold=0.5)
+        else: # for large images (>5000x5000), the F1 score is computed by a patch-based way
+            H, W = gt.shape
+            roi_size = 2000
+        
+            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    
+        
+            gt_pad = np.zeros((new_H, new_W), dtype=gt.dtype)
+            seg_pad = np.zeros((new_H, new_W), dtype=gt.dtype)
+            gt_pad[:H, :W] = gt
+            seg_pad[:H, :W] = seg
+              
+            tp = 0
+            fp = 0
+            fn = 0
+            for i in range(n_H):
+                for j in range(n_W):
+                    gt_roi  = remove_boundary_cells(gt_pad[roi_size*i:roi_size*(i+1), roi_size*j:roi_size*(j+1)])
+                    seg_roi = remove_boundary_cells(seg_pad[roi_size*i:roi_size*(i+1), roi_size*j:roi_size*(j+1)])
+                    tp_i, fp_i, fn_i = eval_tp_fp_fn(gt_roi, seg_roi, threshold=0.5)
+                    tp += tp_i
+                    fp += fp_i
+                    fn += fn_i            
+        
+        if tp == 0:
+            precision = 0
+            recall = 0
+            f1 = 0
+        else:
+            precision = tp / (tp + fp)
+            recall = tp / (tp + fn)
+            f1 = 2*(precision * recall)/ (precision + recall)
+        seg_metric['Names'].append(name)
+        seg_metric['F1_Score'].append(np.round(f1, 4))
+    
+    
+    seg_metric_df = pd.DataFrame(seg_metric)
+    seg_metric_df.to_csv(join(args.save_path, 'seg_metric.csv'), index=False)
+    print('mean F1 Score:', np.mean(seg_metric['F1_Score']))
+
+if __name__ == '__main__':
+    main()

finetune_convnext_stardist.py

@@ -0,0 +1,388 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Adapted form MONAI Tutorial: https://github.com/Project-MONAI/tutorials/tree/main/2d_segmentation/torch
+"""
+
+import argparse
+import os
+
+join = os.path.join
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from stardist import star_dist,edt_prob
+from stardist import dist_to_coord, non_maximum_suppression, polygons_to_label
+from stardist import random_label_cmap,ray_angles
+import monai
+from collections import OrderedDict
+from compute_metric import eval_tp_fp_fn,remove_boundary_cells
+from monai.data import decollate_batch, PILReader
+from monai.inferers import sliding_window_inference
+from monai.metrics import DiceMetric
+from monai.transforms import (
+    Activations,
+    AsChannelFirstd,
+    AddChanneld,
+    AsDiscrete,
+    Compose,
+    LoadImaged,
+    SpatialPadd,
+    RandSpatialCropd,
+    RandRotate90d,
+    ScaleIntensityd,
+    RandAxisFlipd,
+    RandZoomd,
+    RandGaussianNoised,
+    RandAdjustContrastd,
+    RandGaussianSmoothd,
+    RandHistogramShiftd,
+    EnsureTyped,
+    EnsureType,
+)
+from monai.visualize import plot_2d_or_3d_image
+import matplotlib.pyplot as plt
+from datetime import datetime
+import shutil
+import tqdm
+from models.unetr2d import UNETR2D
+from models.swin_unetr import SwinUNETR
+from models.flexible_unet import FlexibleUNet 
+from models.flexible_unet_convext import FlexibleUNetConvext
+print("Successfully imported all requirements!")
+torch.backends.cudnn.enabled =False
+#os.environ["OMP_NUM_THREADS"] = "1" 
+#os.environ["MKL_NUM_THREADS"] = "1" 
+def main():
+    parser = argparse.ArgumentParser("Baseline for Microscopy image segmentation")
+    # Dataset parameters
+    parser.add_argument(
+        "--data_path",
+        default="",
+        type=str,
+        help="training data path; subfolders: images, labels",
+    )
+    parser.add_argument(
+        "--work_dir", default="/mntnfs/med_data5/louwei/nips_comp/stardist_finetune1/", help="path where to save models and logs"
+    )
+    parser.add_argument(
+        "--model_dir", default="/", help="path where to load pretrained model"
+    )
+    parser.add_argument("--seed", default=2022, type=int)
+    # parser.add_argument("--resume", default=False, help="resume from checkpoint")
+    parser.add_argument("--num_workers", default=4, type=int)
+    #parser.add_argument("--local_rank", type=int)
+    # Model parameters
+    parser.add_argument(
+        "--model_name", default="efficientunet", help="select mode: unet, unetr, swinunetr"
+    )
+    parser.add_argument("--num_class", default=3, type=int, help="segmentation classes")
+    parser.add_argument(
+        "--input_size", default=512, type=int, help="segmentation classes"
+    )
+    # Training parameters
+    parser.add_argument("--batch_size", default=16, type=int, help="Batch size per GPU")
+    parser.add_argument("--max_epochs", default=2000, type=int)
+    parser.add_argument("--val_interval", default=10, type=int)
+    parser.add_argument("--epoch_tolerance", default=100, type=int)
+    parser.add_argument("--initial_lr", type=float, default=1e-4, help="learning rate")
+
+    args = parser.parse_args()
+    #torch.cuda.set_device(args.local_rank)
+    #torch.distributed.init_process_group(backend='nccl')
+    monai.config.print_config()
+    n_rays = 32
+    pre_trained = True
+    #%% set training/validation split
+    np.random.seed(args.seed)
+	pre_trained_path = args.model_dir
+    model_path = join(args.work_dir, args.model_name + "_3class")
+    os.makedirs(model_path, exist_ok=True)
+    run_id = datetime.now().strftime("%Y%m%d-%H%M")
+    # This must be change every runing time ! ! ! ! ! ! ! ! ! ! !
+    model_file = "models/flexible_unet_convext.py"
+    shutil.copyfile(
+        __file__, join(model_path, os.path.basename(__file__))
+    )
+    shutil.copyfile(
+        model_file, join(model_path, os.path.basename(model_file))
+    )
+    img_path = join(args.data_path, "train/images")
+    gt_path = join(args.data_path, "train/tif")
+    val_img_path = join(args.data_path, "valid/images")
+    val_gt_path = join(args.data_path, "valid/tif")
+    img_names = sorted(os.listdir(img_path))
+    gt_names = [img_name.split(".")[0] + ".tif" for img_name in img_names]
+    img_num = len(img_names)
+    val_frac = 0.1
+    val_img_names = sorted(os.listdir(val_img_path))
+    val_gt_names = [img_name.split(".")[0] + ".tif" for img_name in val_img_names]
+    #indices = np.arange(img_num)
+    #np.random.shuffle(indices)
+    #val_split = int(img_num * val_frac)
+    #train_indices = indices[val_split:]
+    #val_indices = indices[:val_split]
+
+    train_files = [
+        {"img": join(img_path, img_names[i]), "label": join(gt_path, gt_names[i])}
+        for i in range(len(img_names))
+    ]
+    val_files = [
+        {"img": join(val_img_path, val_img_names[i]), "label": join(val_gt_path, val_gt_names[i])}
+        for i in range(len(val_img_names))
+    ]
+    print(
+        f"training image num: {len(train_files)}, validation image num: {len(val_files)}"
+    )
+    #%% define transforms for image and segmentation
+    train_transforms = Compose(
+        [
+            LoadImaged(
+                keys=["img", "label"], reader=PILReader, dtype=np.float32
+            ),  # image three channels (H, W, 3); label: (H, W)
+            AddChanneld(keys=["label"], allow_missing_keys=True),  # label: (1, H, W)
+            AsChannelFirstd(
+                keys=["img"], channel_dim=-1, allow_missing_keys=True
+            ),  # image: (3, H, W)
+            #ScaleIntensityd(
+                #keys=["img"], allow_missing_keys=True
+            #),  # Do not scale label
+            SpatialPadd(keys=["img", "label"], spatial_size=args.input_size),
+            RandSpatialCropd(
+                keys=["img", "label"], roi_size=args.input_size, random_size=False
+            ),
+            RandAxisFlipd(keys=["img", "label"], prob=0.5),
+            RandRotate90d(keys=["img", "label"], prob=0.5, spatial_axes=[0, 1]),
+            # # intensity transform
+            RandGaussianNoised(keys=["img"], prob=0.25, mean=0, std=0.1),
+            RandAdjustContrastd(keys=["img"], prob=0.25, gamma=(1, 2)),
+            RandGaussianSmoothd(keys=["img"], prob=0.25, sigma_x=(1, 2)),
+            RandHistogramShiftd(keys=["img"], prob=0.25, num_control_points=3),
+            RandZoomd(
+                keys=["img", "label"],
+                prob=0.15,
+                min_zoom=0.5,
+                max_zoom=2,
+                mode=["area", "nearest"],
+            ),
+            EnsureTyped(keys=["img", "label"]),
+        ]
+    )
+
+    val_transforms = Compose(
+        [
+            LoadImaged(keys=["img", "label"], reader=PILReader, dtype=np.float32),
+            AddChanneld(keys=["label"], allow_missing_keys=True),
+            AsChannelFirstd(keys=["img"], channel_dim=-1, allow_missing_keys=True),
+            #ScaleIntensityd(keys=["img"], allow_missing_keys=True),
+            # AsDiscreted(keys=['label'], to_onehot=3),
+            EnsureTyped(keys=["img", "label"]),
+        ]
+    )
+
+    #% define dataset, data loader
+    check_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
+    check_loader = DataLoader(check_ds, batch_size=1, num_workers=4)
+    check_data = monai.utils.misc.first(check_loader)
+    print(
+        "sanity check:",
+        check_data["img"].shape,
+        torch.max(check_data["img"]),
+        check_data["label"].shape,
+        torch.max(check_data["label"]),
+    )
+
+    #%% create a training data loader
+    train_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
+    # use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory =True,
+    )
+    # create a validation data loader
+    val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
+    val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=1)
+
+    dice_metric = DiceMetric(
+        include_background=False, reduction="mean", get_not_nans=False
+    )
+
+    post_pred = Compose(
+        [EnsureType(), Activations(softmax=True), AsDiscrete(threshold=0.5)]
+    )
+    post_gt = Compose([EnsureType(), AsDiscrete(to_onehot=None)])
+    # create UNet, DiceLoss and Adam optimizer
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    if args.model_name.lower() == "efficientunet":
+        model = FlexibleUNetConvext(
+            in_channels=3,
+            out_channels=n_rays+1,
+            backbone='convnext_small',
+            pretrained=False,
+        ).to(device)
+  
+    #loss_masked_dice = monai.losses.DiceCELoss(softmax=True)
+    loss_dice = monai.losses.DiceLoss(squared_pred=True,jaccard=True)
+    loss_bce = nn.BCELoss()
+    loss_dist_mae = nn.L1Loss()
+    activatation = nn.ReLU()
+    sigmoid = nn.Sigmoid()
+    #loss_dist_mae = monai.losses.DiceCELoss(softmax=True)
+    initial_lr = args.initial_lr
+    encoder = list(map(id, model.encoder.parameters()))
+    base_params = filter(lambda p: id(p) not in encoder, model.parameters())
+    params = [
+        {"params": base_params, "lr":initial_lr},
+        {"params": model.encoder.parameters(), "lr": initial_lr * 0.1},
+    ]
+    optimizer = torch.optim.AdamW(params, initial_lr)
+    if pre_trained == True:
+
+        checkpoint = torch.load(pre_trained_path, map_location=torch.device(device))
+        model.load_state_dict(checkpoint['model_state_dict'])
+        print('Load pretrained weights...')
+    max_epochs = args.max_epochs
+    epoch_tolerance = args.epoch_tolerance
+    val_interval = args.val_interval
+    best_metric = -1
+    best_metric_epoch = -1
+    epoch_loss_values = list()
+    metric_values = list()
+    writer = SummaryWriter(model_path)
+    max_f1 = 0
+    for epoch in range(0, max_epochs):
+        model.train()
+        epoch_loss = 0
+        epoch_loss_prob = 0
+        epoch_loss_dist_2 = 0
+        epoch_loss_dist_1 = 0
+        for step, batch_data in enumerate(train_loader, 1):
+            print(step)
+            inputs, labels = batch_data["img"],batch_data["label"]
+
+            processes_labels = []
+            
+            for i in range(labels.shape[0]):
+                label = labels[i][0]
+                distances = star_dist(label,n_rays,mode='opencl')
+                distances = np.transpose(distances,(2,0,1))
+                #print(distances.shape)
+                obj_probabilities = edt_prob(label.astype(int))
+                obj_probabilities = np.expand_dims(obj_probabilities,0)
+                #print(obj_probabilities.shape)
+                final_label = np.concatenate((distances,obj_probabilities),axis=0)
+                #print(final_label.shape)
+                processes_labels.append(final_label)
+            
+            labels = np.stack(processes_labels)
+
+            #print(inputs.shape,labels.shape)
+            inputs, labels = torch.tensor(inputs).to(device), torch.tensor(labels).to(device)
+            #print(inputs.shape,labels.shape)
+            optimizer.zero_grad()
+            output_dist,output_prob = model(inputs)
+            #print(outputs.shape)
+            dist_output = output_dist
+            prob_output = output_prob
+            dist_label = labels[:,:n_rays,:,:]
+            prob_label = torch.unsqueeze(labels[:,-1,:,:], 1)
+            #print(dist_output.shape,prob_output.shape,dist_label.shape)
+            #labels_onehot = monai.networks.one_hot(
+                #labels, args.num_class
+            #)  # (b,cls,256,256)
+            #print(prob_label.max(),prob_label.min())
+            loss_dist_1 = loss_dice(dist_output*prob_label,dist_label*prob_label)
+            #print(loss_dist_1)
+            loss_prob = loss_bce(prob_output,prob_label)
+            #print(prob_label.shape,dist_output.shape)
+            loss_dist_2 = loss_dist_mae(dist_output*prob_label,dist_label*prob_label)
+            #print(loss_dist_2)
+            loss = loss_prob + loss_dist_2*0.3 + loss_dist_1
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item()
+            epoch_loss_prob += loss_prob.item()
+            epoch_loss_dist_2 += loss_dist_2.item()
+            epoch_loss_dist_1 += loss_dist_1.item()
+            epoch_len = len(train_ds) // train_loader.batch_size
+            
+        epoch_loss /= step
+        epoch_loss_prob /= step
+        epoch_loss_dist_2 /= step
+        epoch_loss_dist_1 /= step
+        epoch_loss_values.append(epoch_loss)
+        print(f"epoch {epoch} average loss: {epoch_loss:.4f}")
+        writer.add_scalar("train_loss", epoch_loss, epoch)
+        print('dist dice: '+str(epoch_loss_dist_1)+' dist mae: '+str(epoch_loss_dist_2)+' prob bce: '+str(epoch_loss_prob))
+        checkpoint = {	
+            "epoch": epoch,
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "loss": epoch_loss_values,
+        }
+        if epoch < 40:
+            continue
+        if epoch > 1 and epoch % val_interval == 0:
+            torch.save(checkpoint, join(model_path, str(epoch) + ".pth"))
+            model.eval()
+            with torch.no_grad():
+                val_images = None
+                val_labels = None
+                val_outputs = None
+                seg_metric = OrderedDict()
+                seg_metric['F1_Score'] = []
+                for val_data in tqdm.tqdm(val_loader):
+                    val_images, val_labels = val_data["img"].to(device), val_data[
+                        "label"
+                    ].to(device)
+                    roi_size = (512, 512)
+                    sw_batch_size = 4
+                    output_dist,output_prob = sliding_window_inference(
+                        val_images, roi_size, sw_batch_size, model
+                        )
+                    val_labels = val_labels[0][0].cpu().numpy()
+                    prob = output_prob[0][0].cpu().numpy()
+                    dist = output_dist[0].cpu().numpy()
+                    #print(val_labels.shape,prob.shape,dist.shape)
+                    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)
+            
+                    star_label = polygons_to_label(disti, points, prob=probi,shape=prob.shape)
+                    gt = remove_boundary_cells(val_labels.astype(np.int32)) 
+                    seg = remove_boundary_cells(star_label.astype(np.int32))           
+                    tp, fp, fn = eval_tp_fp_fn(gt, seg, threshold=0.5)
+                    if tp == 0:
+                        precision = 0
+                        recall = 0
+                        f1 = 0
+                    else:
+                        precision = tp / (tp + fp)
+                        recall = tp / (tp + fn)
+                        f1 = 2*(precision * recall)/ (precision + recall)
+                    f1 = np.round(f1, 4)
+                    seg_metric['F1_Score'].append(np.round(f1, 4))
+                avg_f1 = np.mean(seg_metric['F1_Score'])
+                writer.add_scalar("val_f1score", avg_f1, epoch)
+                if avg_f1 > max_f1:
+                    max_f1 = avg_f1
+                    print(str(epoch) + 'f1 score: ' + str(max_f1))
+                    torch.save(checkpoint, join(model_path, "best_model.pth"))
+    np.savez_compressed(
+        join(model_path, "train_log.npz"),
+        val_dice=metric_values,
+        epoch_loss=epoch_loss_values,
+    )
+
+
+if __name__ == "__main__":
+    main()

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_convext.py

@@ -0,0 +1,451 @@
+# 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 FlexibleUNetConvext(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,
+    ) -> 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=True,in_22k=True)
+        # self.encoder = VAN(embed_dims=[64, 128, 320, 512],
+        #  depths=[3, 3, 12, 3],
+        #  init_cfg=dict(type='Pretrained', checkpoint='pretrained/van_b2.pth'),
+        #  norm_cfg=dict(type='BN', requires_grad=True)
+        #  )
+        # self.encoder = VAN(embed_dims=[64, 128, 320, 512],
+        #  depths=[2, 2, 4, 2],
+        #  init_cfg=dict(type='Pretrained', checkpoint='pretrained/van_b1.pth'),
+        #  norm_cfg=dict(type='BN', requires_grad=True)
+        #  )
+        # self.encoder.init_weights()
+        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=64,
+            kernel_size=1,
+            act='relu',
+            scale_factor = 2,
+        )
+        self.prob_head = SegmentationHead(
+            spatial_dims=spatial_dims,
+            in_channels=decoder_channels[-1],
+            out_channels=1,
+            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=True,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[ ]:
+
+
+
+

predict.py

@@ -0,0 +1,157 @@
+
+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 overlay import visualize_instances_map
+
+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)
+#torch.cuda.synchronize()
+parser = argparse.ArgumentParser('Baseline for Microscopy image segmentation', add_help=False)
+# Dataset parameters
+parser.add_argument('-i', '--input_path', default='./inputs', type=str, help='training data path; subfolders: images, labels')
+parser.add_argument("-o", '--output_path', default='./outputs', type=str, help='output path')
+parser.add_argument('--model_path', default='./models', help='path where to save models and segmentation results')
+parser.add_argument('--show_overlay', required=False, default=False, action="store_true", help='save segmentation overlay')
+
+# Model parameters
+parser.add_argument('--model_name', default='efficientunet', help='select mode: unet, unetr, swinunetr')
+parser.add_argument('--input_size', default=512, type=int, help='segmentation classes')
+args = parser.parse_args()
+input_path = args.input_path
+output_path = args.output_path
+model_path = args.model_path
+os.makedirs(output_path, exist_ok=True)
+#overlay_path = 'overlays/'
+#print(input_path)
+
+img_names = sorted(os.listdir(join(input_path)))
+#print(img_names)
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+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])
+])
+roi_size = (512, 512)
+overlap = 0.5
+np_thres, ksize, overall_thres, obj_size_thres = 0.6, 15, 0.4, 100
+n_rays = 32
+sw_batch_size = 4
+num_classes= 4
+block_size = 2048
+min_overlap = 128
+context = 128
+with torch.no_grad():
+    for img_name in img_names:
+        #print(img_name)
+        if img_name.endswith('.tif') or img_name.endswith('.tiff'):
+            img_data = tif.imread(join(input_path, img_name))
+        else:
+            img_data = io.imread(join(input_path, 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)
+        inputs=preprocess(Image.fromarray(pre_img_data)).unsqueeze(0).to(device)
+        cls_MODEL = model_path + '/cls/resnet18_4class_all_modified.tar'
+        model = resnet18().to(device)
+        model.load_state_dict(torch.load(cls_MODEL))
+        model.eval()
+        outputs = model(inputs)
+        _, preds = torch.max(outputs, 1)    
+        label=preds[0].cpu().numpy()
+        #print(label)
+        test_npy01 = pre_img_data
+        if label in [0,1,2] or img_data.shape[0] > 4000:
+            if label == 0:
+                model = FlexibleUNet_star(in_channels=3,out_channels=n_rays+1,backbone='convnext_small',pretrained=False,n_rays=n_rays,prob_out_channels=1,).to(device)
+                checkpoint = torch.load(model_path+'/0/best_model.pth', map_location=torch.device(device))        
+                model.load_state_dict(checkpoint['model_state_dict'])
+                model.eval()
+                
+                output_label = sliding_window_inference_large(test_npy01,block_size,min_overlap,context, roi_size,sw_batch_size,predictor=model,device=device)
+                tif.imwrite(join(output_path, img_name.split('.')[0]+'_label.tiff'), output_label)
+
+            elif label == 1:
+                model = FlexibleUNet_star(in_channels=3,out_channels=n_rays+1,backbone='convnext_small',pretrained=False,n_rays=n_rays,prob_out_channels=1,).to(device)
+                checkpoint = torch.load(model_path+'/1/best_model.pth', map_location=torch.device(device))
+                model.load_state_dict(checkpoint['model_state_dict'])
+                model.eval()
+                
+                output_label = sliding_window_inference_large(test_npy01,block_size,min_overlap,context, roi_size,sw_batch_size,predictor=model,device=device)
+                tif.imwrite(join(output_path, img_name.split('.')[0]+'_label.tiff'), output_label) 
+            elif label == 2:
+                model = FlexibleUNet_star(in_channels=3,out_channels=n_rays+1,backbone='convnext_small',pretrained=False,n_rays=n_rays,prob_out_channels=1,).to(device)
+                checkpoint = torch.load(model_path+'/2/best_model.pth', map_location=torch.device(device))
+                model.load_state_dict(checkpoint['model_state_dict'])
+                model.eval()
+                
+                output_label = sliding_window_inference_large(test_npy01,block_size,min_overlap,context, roi_size,sw_batch_size,predictor=model,device=device)
+                tif.imwrite(join(output_path, img_name.split('.')[0]+'_label.tiff'), output_label)
+
+
+        else:
+            model = FlexibleUNet_hv(in_channels=3,out_channels=2+2,backbone='convnext_small',pretrained=False,n_rays=2,prob_out_channels=2,).to(device)
+            checkpoint = torch.load(model_path+'/3/best_model_converted.pth', map_location=torch.device(device))
+            #model.load_state_dict(checkpoint['model_state_dict'])
+            #od = OrderedDict()
+            #for k, v in checkpoint['model_state_dict'].items():
+                #od[k.replace('module.', '')] = v
+            model.load_state_dict(checkpoint)
+            model.to(device)
+            model.eval()
+            test_tensor = torch.from_numpy(np.expand_dims(test_npy01, 0)).permute(0, 3, 1, 2).type(torch.FloatTensor).to(device) 
+            if isinstance(roi_size, tuple):
+                roi = roi_size
+
+            output_hv, output_np = sliding_window_inference(test_tensor, roi, sw_batch_size, model, overlap=overlap)
+            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, np_thres, ksize, overall_thres, obj_size_thres)
+            raw_pred_shape = pred_inst.shape[:2]
+            output_label = pred_inst
+
+            tif.imwrite(join(output_path, img_name.split('.')[0]+'_label.tiff'), output_label)
+
+
+
+

predict_unet_convnext.py

@@ -0,0 +1,123 @@
+
+import os
+join = os.path.join
+import argparse
+import numpy as np
+import torch
+import monai
+import torch.nn as nn
+
+from utils import sliding_window_inference
+#from baseline.models.unetr2d import UNETR2D
+import time
+from stardist import dist_to_coord, non_maximum_suppression, polygons_to_label
+from stardist import random_label_cmap,ray_angles
+from stardist import star_dist,edt_prob
+from skimage import io, segmentation, morphology, measure, exposure
+import tifffile as tif
+import cv2
+from overlay import visualize_instances_map
+from models.flexible_unet import FlexibleUNet 
+from models.flexible_unet_convext import FlexibleUNetConvext
+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)
+
+def main():
+    parser = argparse.ArgumentParser('Baseline for Microscopy image segmentation', add_help=False)
+    # Dataset parameters
+    #parser.add_argument('-i', '--input_path', default='./inputs', type=str, help='training data path; subfolders: images, labels')
+    #parser.add_argument("-o", '--output_path', default='./outputs', type=str, help='output path')
+    parser.add_argument('--model_path', default='./work_dir/swinunetr_3class', help='path where to save models and segmentation results')
+    parser.add_argument('--show_overlay', required=False, default=False, action="store_true", help='save segmentation overlay')
+
+    # Model parameters
+    parser.add_argument('--model_name', default='efficientunet', help='select mode: unet, unetr, swinunetr')
+    parser.add_argument('--num_class', default=3, type=int, help='segmentation classes')
+    parser.add_argument('--input_size', default=512, type=int, help='segmentation classes')
+    args = parser.parse_args()
+    
+    input_path = '/home/data/TuningSet/'
+    output_path = '/home/data/output/'
+    overlay_path = '/home/data/overlay/'
+
+
+    img_names = sorted(os.listdir(join(input_path)))
+    n_rays = 32
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+
+    if args.model_name.lower() == "efficientunet":
+        model = FlexibleUNetConvext(
+            in_channels=3,
+            out_channels=n_rays+1,
+            backbone='convnext_small',
+            pretrained=True,
+        ).to(device)
+    
+
+
+    sigmoid = nn.Sigmoid()
+    checkpoint = torch.load('/home/louwei/stardist_convnext/efficientunet_3class/best_model.pth', map_location=torch.device(device))
+    model.load_state_dict(checkpoint['model_state_dict'])
+    #%%
+    roi_size = (args.input_size, args.input_size)
+    sw_batch_size = 4
+    model.eval()
+    with torch.no_grad():
+        for img_name in img_names:
+            print(img_name)
+            if img_name.endswith('.tif') or img_name.endswith('.tiff'):
+                img_data = tif.imread(join(input_path, img_name))
+            else:
+                img_data = io.imread(join(input_path, 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)
+            
+            t0 = time.time()
+            #test_npy01 = pre_img_data/np.max(pre_img_data)
+            test_npy01 = pre_img_data
+            test_tensor = torch.from_numpy(np.expand_dims(test_npy01, 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, model)
+            #test_pred_out = torch.nn.functional.softmax(test_pred_out, dim=1) # (B, C, H, W)
+            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)
+            
+            star_label = polygons_to_label(disti, points, prob=probi,shape=prob.shape)
+            tif.imwrite(join(output_path, img_name.split('.')[0]+'_label.tiff'), star_label)
+            overlay = visualize_instances_map(pre_img_data,star_label)
+            cv2.imwrite(join(overlay_path, img_name.split('.')[0]+'.png'), cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
+            
+            
+        
+if __name__ == "__main__":
+    main()
+
+
+
+
+

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
+

train_convnext_hover..py

@@ -0,0 +1,516 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Adapted form MONAI Tutorial: https://github.com/Project-MONAI/tutorials/tree/main/2d_segmentation/torch
+"""
+
+import argparse
+import os, sys
+
+join = os.path.join
+#sys.path.append('/data2/yuxinyi/stardist_pytorch')
+
+from tqdm import tqdm
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import DataParallel
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from torch.optim.lr_scheduler import ReduceLROnPlateau, StepLR
+from stardist import star_dist, edt_prob
+from stardist import dist_to_coord, non_maximum_suppression, polygons_to_label
+from stardist import random_label_cmap, ray_angles
+import monai
+from collections import OrderedDict
+from compute_metric import eval_tp_fp_fn, remove_boundary_cells
+from monai.data import decollate_batch, PILReader
+from monai.inferers import sliding_window_inference
+from monai.metrics import DiceMetric
+from monai.transforms import (
+    Activations,
+    AsChannelFirstd,
+    AddChanneld,
+    AsDiscrete,
+    CenterSpatialCropd,
+    Compose,
+    Lambdad,
+    LoadImaged,
+    # LoadImaged_modified,
+    SpatialPadd,
+    RandSpatialCropd,
+    RandRotate90d,
+    ScaleIntensityd,
+    RandAxisFlipd,
+    RandZoomd,
+    RandGaussianNoised,
+    RandAdjustContrastd,
+    RandGaussianSmoothd,
+    RandHistogramShiftd,
+    EnsureTyped,
+    EnsureType,
+    apply_transform,
+)
+from monai.visualize import plot_2d_or_3d_image
+import matplotlib.pyplot as plt
+from datetime import datetime
+import shutil
+from skimage import io
+from skimage.color import gray2rgb
+
+from models.unetr2d import UNETR2D
+from models.swin_unetr import SwinUNETR
+from models.flexible_unet_convext import FlexibleUNet_hv
+
+from utils import cropping_center, gen_targets, xentropy_loss, dice_loss, mse_loss, msge_loss
+
+import warnings
+warnings.filterwarnings("ignore")
+
+print("Successfully imported all requirements!")
+torch.backends.cudnn.enabled = False
+
+def rm_n_mkdir(dir_path):
+    """Remove and make directory."""
+    if os.path.isdir(dir_path):
+        shutil.rmtree(dir_path)
+    os.makedirs(dir_path)
+
+class HoverDataset(Dataset):
+    def __init__(self, data, transform, mask_shape):
+        self.data = data
+        self.transform = transform
+        self.mask_shape = mask_shape
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def _transform(self, index):
+        data_i = self.data[index]
+        return apply_transform(self.transform, data_i) if self.transform is not None else data_i
+
+    def __getitem__(self, index):
+        ret = self._transform(index)
+        # print(target_dict['img'].dtype, target_dict['label'].dtype)
+        # gen targets
+        inst_map = np.squeeze(ret['label'].numpy()).astype('int32') # 1HW -> HW
+        target_dict = gen_targets(inst_map, inst_map.shape[:2])  # original code: self.mask_shape -> current code: aug_size
+        np_map, hv_map = target_dict['np_map'], target_dict['hv_map']
+        np_map = cropping_center(np_map, self.mask_shape)  # HW
+        hv_map = cropping_center(hv_map, self.mask_shape) # HW2
+        target_dict['np_map'] = torch.tensor(np_map)
+        target_dict['hv_map'] = torch.tensor(hv_map)
+        # centercrop img
+        img = cropping_center(ret['img'].permute(1,2,0), self.mask_shape).permute(2,0,1) # CHW -> HWC -> CHW
+        ret['img'] = img
+        ret.update(target_dict)
+        return ret
+
+def valid_step(model, batch_data):
+
+    model.eval()  # infer mode
+
+    ####
+    imgs = batch_data["img"]
+    true_np = batch_data["np_map"]
+    true_hv = batch_data["hv_map"]
+
+    imgs_gpu = imgs.to("cuda").type(torch.float32)  # NCHW
+
+    # HWC
+    true_np = torch.squeeze(true_np).type(torch.int64)
+    true_hv = torch.squeeze(true_hv).type(torch.float32)
+
+    true_dict = {
+        "np": true_np,
+        "hv": true_hv,
+    }
+
+    # --------------------------------------------------------------
+    with torch.no_grad():  # dont compute gradient
+        preds = model(imgs_gpu)
+        pred_dict = {'np': preds[1], 'hv': preds[0]}
+        pred_dict = OrderedDict(
+            [[k, v.permute(0, 2, 3, 1).contiguous()] for k, v in pred_dict.items()]
+        )
+        pred_dict["np"] = F.softmax(pred_dict["np"], dim=-1)[..., 1]
+
+    # * Its up to user to define the protocol to process the raw output per step!
+    result_dict = {  # protocol for contents exchange within `raw`
+        "raw": {
+            "imgs": imgs.numpy(),
+            "true_np": true_dict["np"].numpy(),
+            "true_hv": true_dict["hv"].numpy(),
+            "prob_np": pred_dict["np"].cpu().numpy(),
+            "pred_hv": pred_dict["hv"].cpu().numpy(),
+        }
+    }
+
+    return result_dict
+
+def proc_valid_step_output(raw_data, nr_types=None):
+
+    track_dict = {}
+    
+    def _dice_info(true, pred, label):
+        true = np.array(true == label, np.int32)
+        pred = np.array(pred == label, np.int32)
+        inter = (pred * true).sum()
+        total = (pred + true).sum()
+        return inter, total
+
+    over_inter = 0
+    over_total = 0
+    over_correct = 0
+    prob_np = raw_data["prob_np"]
+    true_np = raw_data["true_np"]
+    for idx in range(len(raw_data["true_np"])):
+        patch_prob_np = prob_np[idx]
+        patch_true_np = true_np[idx]
+        patch_pred_np = np.array(patch_prob_np > 0.5, dtype=np.int32)
+        inter, total = _dice_info(patch_true_np, patch_pred_np, 1)
+        correct = (patch_pred_np == patch_true_np).sum()
+        over_inter += inter
+        over_total += total
+        over_correct += correct
+    nr_pixels = len(true_np) * np.size(true_np[0])
+    acc_np = over_correct / nr_pixels
+    dice_np = 2 * over_inter / (over_total + 1.0e-8)
+    track_dict['np_acc'] = acc_np
+    track_dict['np_dice'] = dice_np
+
+    # * HV regression statistic
+    pred_hv = raw_data["pred_hv"]
+    true_hv = raw_data["true_hv"]
+
+    over_squared_error = 0
+    for idx in range(len(raw_data["true_np"])):
+        patch_pred_hv = pred_hv[idx]
+        patch_true_hv = true_hv[idx]
+        squared_error = patch_pred_hv - patch_true_hv
+        squared_error = squared_error * squared_error
+        over_squared_error += squared_error.sum()
+    mse = over_squared_error / nr_pixels
+    track_dict['hv_mse'] = mse
+
+    return track_dict
+
+def main():
+
+    # class Args:
+    #     def __init__(self, data_path, seed, num_workers, model_name, input_size, mask_size, batch_size, max_epochs,
+    #                  val_interval, save_interval, initial_lr, gpu_id, n_rays):
+    #         self.data_path = data_path
+    #         self.seed = seed
+    #         self.num_workers = num_workers
+    #         self.model_name = model_name
+    #         self.input_size = input_size
+    #         self.mask_size = mask_size
+    #         self.batch_size = batch_size
+    #         self.max_epochs = max_epochs
+    #         self.val_interval = val_interval
+    #         self.save_interval = save_interval
+    #         self.initial_lr = initial_lr
+    #         self.gpu_id = gpu_id
+    #         self.n_rays = n_rays
+
+    # args = Args('/data2/yuxinyi/stardist_pytorch/dataset/class3_seed2', 2022, 4, 'efficientunet', 512, 256, 16, 600,
+    #             1, 10, 1e-4, '4', 32)
+    modelname = 'star-hover'
+    strategy = 'aug256_out256'
+    parser = argparse.ArgumentParser("Baseline for Microscopy image segmentation")
+    # Dataset parameters
+    parser.add_argument(
+        "--data_path",
+        default=f"/mntnfs/med_data5/louwei/consep/",
+        type=str,
+        help="training data path; subfolders: images, labels",
+    )
+    parser.add_argument("--seed", default=10, type=int)
+    # parser.add_argument("--resume", default=False, help="resume from checkpoint")
+    parser.add_argument("--num_workers", default=4, type=int)
+    
+    # Model parameters
+    parser.add_argument(
+        "--model_name", default="efficientunet", help="select mode: unet, unetr, swinunetr"
+    )
+    parser.add_argument("--input_size", default=512, type=int, help="after rand crop")
+    parser.add_argument("--mask_size", default=256, type=int, help="after gen target")
+    # Training parameters
+    parser.add_argument("--batch_size", default=12, type=int, help="Batch size per GPU")
+    parser.add_argument("--max_epochs", default=800, type=int)
+    parser.add_argument("--val_interval", default=1, type=int)
+    parser.add_argument("--save_interval", default=10, type=int)
+    parser.add_argument("--initial_lr", type=float, default=1e-4, help="learning rate")
+    parser.add_argument('--gpu_id', type=str, default='0', help='gpu id')
+
+    args = parser.parse_args()
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)
+
+    work_dir = f'/mntnfs/med_data5/louwei/hover_stardist/class_{modelname}_{strategy}'
+
+    # monai.config.print_config()
+    pre_trained = False
+    # %% set training/validation split
+    np.random.seed(args.seed)
+    model_path = join(work_dir)
+    rm_n_mkdir(model_path)
+    run_id = datetime.now().strftime("%Y%m%d-%H%M")
+    shutil.copyfile(
+        __file__, join(model_path, run_id + "_" + os.path.basename(__file__))
+    )
+    img_path = join(args.data_path, "Train/Images_3channels")
+    gt_path = join(args.data_path, "Train/tif")
+    val_img_path = join(args.data_path, "Test/Images_3channels")
+    val_gt_path = join(args.data_path, "Test/tif")
+    img_names = sorted(os.listdir(img_path))
+    gt_names = [img_name.replace('.png', '.tif') for img_name in img_names]
+    img_num = len(img_names)
+    val_frac = 0.1
+    val_img_names = sorted(os.listdir(val_img_path))
+    val_gt_names = [img_name.replace('.png', '.tif') for img_name in val_img_names]
+
+    train_files = [
+        {"img": join(img_path, img_names[i]), "label": join(gt_path, gt_names[i]), 'name': img_names[i]}
+        for i in range(len(img_names))
+    ]
+    val_files = [
+        {"img": join(val_img_path, val_img_names[i]), "label": join(val_gt_path, val_gt_names[i]),
+         'name': val_img_names[i]}
+        for i in range(len(val_img_names))
+    ]
+    print(
+        f"training image num: {len(train_files)}, validation image num: {len(val_files)}"
+    )
+    
+    def load_img(img):
+        ret = io.imread(img)
+        if len(ret.shape) == 2:
+            ret = gray2rgb(ret)
+        return ret.astype('float32')
+    
+    def load_ann(ann):
+        ret = np.squeeze(io.imread(ann)).astype('float32')
+        return ret
+    
+    # %% define transforms for image and segmentation
+    train_transforms = Compose(
+        [
+            Lambdad(('img',), load_img),
+            Lambdad(('label',), load_ann),
+            # LoadImaged(
+            #     keys=["img", "label"], reader=PILReader, dtype=np.float32
+            # ),  # image three channels (H, W, 3); label: (H, W)
+            AddChanneld(keys=["label"], allow_missing_keys=True),  # label: (1, H, W)
+            AsChannelFirstd(
+                keys=["img"], channel_dim=-1, allow_missing_keys=True
+            ),  # image: (3, H, W)
+            # ScaleIntensityd(
+            # keys=["img"], allow_missing_keys=True
+            # ),  # Do not scale label
+            # SpatialPadd(keys=["img", "label"], spatial_size=args.input_size),
+            # RandSpatialCropd(
+            #     keys=["img", "label"], roi_size=args.input_size, random_size=False
+            # ),
+            RandAxisFlipd(keys=["img", "label"], prob=0.5),
+            RandRotate90d(keys=["img", "label"], prob=0.5, spatial_axes=[0, 1]),
+            # # intensity transform
+            RandGaussianNoised(keys=["img"], prob=0.25, mean=0, std=0.1),
+            RandAdjustContrastd(keys=["img"], prob=0.25, gamma=(1, 2)),
+            RandGaussianSmoothd(keys=["img"], prob=0.25, sigma_x=(1, 2)),
+            RandHistogramShiftd(keys=["img"], prob=0.25, num_control_points=3),
+            RandZoomd(
+                keys=["img", "label"],
+                prob=0.15,
+                min_zoom=0.5,
+                max_zoom=2.0,
+                mode=["area", "nearest"],
+            ),
+            EnsureTyped(keys=["img", "label"]),
+        ]
+    )
+
+    val_transforms = Compose(
+        [
+            Lambdad(('img',), load_img),
+            Lambdad(('label',), load_ann),
+            # LoadImaged(keys=["img", "label"], reader=PILReader, dtype=np.float32),
+            AddChanneld(keys=["label"], allow_missing_keys=True),
+            AsChannelFirstd(keys=["img"], channel_dim=-1, allow_missing_keys=True),
+            # ScaleIntensityd(keys=["img"], allow_missing_keys=True),
+            # AsDiscreted(keys=['label'], to_onehot=3),
+            # CenterSpatialCropd(
+            #     keys=["img", "label"], roi_size=args.input_size
+            # ),
+            EnsureTyped(keys=["img", "label"]),
+        ]
+    )
+
+    # % define dataset, data loader
+    # check_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
+    check_ds = HoverDataset(data=train_files, transform=train_transforms, mask_shape=(args.mask_size, args.mask_size))
+    print(len(check_ds))
+    tmp = check_ds[0]
+    print(tmp['img'].shape, tmp['label'].shape, tmp['hv_map'].shape, tmp['np_map'].shape)
+    check_loader = DataLoader(check_ds, batch_size=1, num_workers=4)
+    check_data = monai.utils.misc.first(check_loader)
+    print(
+        "sanity check:",
+        check_data["img"].shape,
+        torch.max(check_data["img"]),
+        check_data["label"].shape,
+        torch.max(check_data["label"]),
+        check_data["hv_map"].shape,
+        torch.max(check_data["hv_map"]),
+        check_data["np_map"].shape,
+        torch.max(check_data["np_map"]),
+    )
+
+    # %% create a training data loader
+    # train_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
+    train_ds = HoverDataset(data=train_files, transform=train_transforms, mask_shape=(args.mask_size, args.mask_size))
+    print(len(train_ds))
+    # example = train_ds[0]
+    # plt.imshow(np.array(example['img']).transpose(1,2,0).astype('uint8'))
+    # plt.imshow(np.squeeze(example['np_map'].numpy()).astype('uint8'), 'gray')
+    # plt.imshow(example['hv_map'].numpy()[...,0])
+    # plt.imshow(example['hv_map'].numpy()[..., 1])
+    # plt.show()
+    # use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=torch.cuda.is_available(),
+    )
+    # create a validation data loader
+    # val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
+    val_ds = HoverDataset(data=val_files, transform=val_transforms, mask_shape=(args.mask_size, args.mask_size))
+    val_loader = DataLoader(val_ds, batch_size=16, shuffle=False, num_workers=4)
+
+    model = FlexibleUNet_hv(
+        in_channels=3,
+        out_channels=2+2,
+        backbone='convnext_small',
+        pretrained=True,
+        n_rays=2,
+        prob_out_channels=2,
+    )
+
+    activatation = nn.ReLU()
+    sigmoid = nn.Sigmoid()
+    initial_lr = args.initial_lr
+    optimizer = torch.optim.AdamW(model.parameters(), initial_lr)
+    scheduler = StepLR(optimizer, 100, 0.1)
+    #if pre_trained == True:
+        #print('Load pretrained weights...')
+        #checkpoint = torch.load('/data2/yuxinyi/stardist_pytorch/pretrained/overall/330.pth')
+        #model.load_state_dict(checkpoint['model_state_dict'])
+    # model = DataParallel(model)
+    model = model.to('cuda')
+    # start a typical PyTorch training
+    max_epochs = args.max_epochs
+    val_interval = args.val_interval
+    save_interval = args.save_interval
+    epoch_loss_values = []
+    writer = SummaryWriter(model_path)
+
+    #*# record loss and f1
+    loss_file = f'{work_dir}/train_loss.txt'
+    f1_file = f'{work_dir}/train_loss.txt'
+    if os.path.exists(loss_file):
+        os.remove(loss_file)
+    if os.path.exists(f1_file):
+        os.remove(f1_file)
+    #*#
+
+    for epoch in range(1, args.max_epochs):
+        model.train()
+        epoch_loss = 0
+        running_np_1, running_np_2, running_hv_1, running_hv_2 = 0.0, 0.0, 0.0, 0.0
+        stream = tqdm(train_loader)
+        for step, batch_data in enumerate(stream, start=1):
+            
+            #*# hv map
+            inputs, true_np, true_hv = batch_data["img"], batch_data["np_map"], batch_data['hv_map']
+            true_np = true_np.to("cuda").type(torch.int64) # NHW
+            true_hv = true_hv.to("cuda").type(torch.float32) # NHWC
+            true_np_onehot = (F.one_hot(true_np, num_classes=2)).type(torch.float32) # NHWC
+            inputs = torch.tensor(inputs).to('cuda')
+            # print(inputs.shape, true_np.shape, true_hv.shape)
+            
+            optimizer.zero_grad()
+            pred_hv, pred_np = model(inputs) # NCHW
+            pred_hv = pred_hv.permute(0, 2, 3, 1).contiguous() # NHWC
+            pred_np = pred_np.permute(0, 2, 3, 1).contiguous() # NHWC
+            pred_np = F.softmax(pred_np, dim=-1)
+            
+            # losses
+            loss_np_1 = xentropy_loss(true_np_onehot, pred_np) # bce
+            loss_np_2 = dice_loss(true_np_onehot, pred_np) # dice
+            loss_hv_1 = mse_loss(true_hv, pred_hv) # mse
+            loss_hv_2 = msge_loss(true_hv, pred_hv, true_np_onehot[...,1]) # msge
+            loss = loss_np_1 + loss_np_2 + loss_hv_1 + loss_hv_2
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item()
+            epoch_len = len(train_ds) // train_loader.batch_size
+    
+            running_np_1 += loss_np_1.item()
+            running_np_2 += loss_np_2.item()
+            running_hv_1 += loss_hv_1.item()
+            running_hv_2 += loss_hv_2.item()
+            #*#
+            
+            stream.set_description(
+                f'Epoch {epoch} | np bce: {running_np_1 / step:.4f}, np dice: {running_np_2 / step:.4f}, hv mse: {running_hv_1 / step:.4f}, hv msge: {running_hv_2 / step:.4f}')
+        
+        epoch_loss /= step
+        epoch_loss_values.append(epoch_loss)
+        writer.add_scalar("train_loss", epoch_loss, epoch)
+        writer.add_scalar("np_bce", running_np_1 / step, epoch)
+        writer.add_scalar("np_dice", running_np_2 / step, epoch)
+        writer.add_scalar("hv_mse", running_hv_1 / step, epoch)
+        writer.add_scalar("hv_msge", running_hv_2 / step, epoch)
+        print(f"epoch {epoch} average loss: {epoch_loss:.4f}, lr: {optimizer.param_groups[0]['lr']}")
+    
+        #*# record
+        with open(loss_file, 'a') as f:
+            f.write(f'Epoch{epoch}\tloss:{epoch_loss:.4f}\tnp_bce:{running_np_1/step:.4f}\tnp_dice:{running_np_2/step:.4f}\thv_mse:{running_hv_1/step:.4f}\thv_msge:{running_hv_2/step:.4f}\n')
+        #*#
+    
+        checkpoint = {
+            "epoch": epoch,
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "loss": epoch_loss_values,
+        }
+        if epoch % save_interval == 0:
+            torch.save(checkpoint, join(model_path, str(epoch) + ".pth"))
+    
+        running_np_acc, running_np_dice, running_hv_mse = 0.0, 0.0, 0.0
+        stream_val = tqdm(val_loader)
+        for step, batch_data in enumerate(stream_val, start=1):
+            raw_data = valid_step(model, batch_data)['raw']
+            track_dict = proc_valid_step_output(raw_data)
+            running_np_acc += track_dict['np_acc']
+            running_np_dice += track_dict['np_dice']
+            running_hv_mse += track_dict['hv_mse']
+            stream.set_description(f'Epoch {epoch} | np acc: {running_np_acc / step:.4f}, np dice: {running_np_dice / step:.4f}, hv mse: {running_hv_mse / step:.4f}')
+        writer.add_scalar("np_acc", running_np_acc / step, epoch)
+        writer.add_scalar("np_dice", running_np_dice / step, epoch)
+        writer.add_scalar("hv_mse", running_hv_mse / step, epoch)
+        print(f'Epoch {epoch} | np acc: {running_np_acc / step:.4f}, np dice: {running_np_dice / step:.4f}, hv mse: {running_hv_mse / step:.4f}')
+        
+        #*# record
+        with open(loss_file, 'a') as f:
+            f.write(f'Validation | Epoch{epoch}\tloss:{epoch_loss:.4f}\tnp_acc:{running_np_acc/step:.4f}\tnp_dice:{running_np_dice/step:.4f}\thv_mse:{running_hv_mse/step:.4f}\n')
+        #*#
+        
+        scheduler.step()
+                  
+if __name__ == "__main__":
+    main()

train_convnext_stardist.py

@@ -0,0 +1,417 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Adapted form MONAI Tutorial: https://github.com/Project-MONAI/tutorials/tree/main/2d_segmentation/torch
+"""
+
+import argparse
+import os
+
+join = os.path.join
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from stardist import star_dist,edt_prob
+from stardist import dist_to_coord, non_maximum_suppression, polygons_to_label
+from stardist import random_label_cmap,ray_angles
+import monai
+from collections import OrderedDict
+from compute_metric import eval_tp_fp_fn,remove_boundary_cells
+from monai.data import decollate_batch, PILReader
+from monai.inferers import sliding_window_inference
+from monai.metrics import DiceMetric
+from monai.transforms import (
+    Activations,
+    AsChannelFirstd,
+    AddChanneld,
+    AsDiscrete,
+    Compose,
+    LoadImaged,
+    SpatialPadd,
+    RandSpatialCropd,
+    RandRotate90d,
+    ScaleIntensityd,
+    RandAxisFlipd,
+    RandZoomd,
+    RandGaussianNoised,
+    RandAdjustContrastd,
+    RandGaussianSmoothd,
+    RandHistogramShiftd,
+    EnsureTyped,
+    EnsureType,
+)
+from monai.visualize import plot_2d_or_3d_image
+import matplotlib.pyplot as plt
+from datetime import datetime
+import shutil
+import tqdm
+from models.unetr2d import UNETR2D
+from models.swin_unetr import SwinUNETR
+from models.flexible_unet import FlexibleUNet 
+from models.flexible_unet_convext import FlexibleUNetConvext
+print("Successfully imported all requirements!")
+torch.backends.cudnn.enabled =False
+
+def main():
+    parser = argparse.ArgumentParser("Baseline for Microscopy image segmentation")
+    # Dataset parameters
+    parser.add_argument(
+        "--data_path",
+        default="/data2/liuchenyu/external_processed/split",
+        type=str,
+        help="training data path; subfolders: images, labels",
+    )
+    parser.add_argument(
+        "--work_dir", default="/data/louwei/nips_comp/convnext_fold0", help="path where to save models and logs"
+    )
+    parser.add_argument("--seed", default=2022, type=int)
+    # parser.add_argument("--resume", default=False, help="resume from checkpoint")
+    parser.add_argument("--num_workers", default=8, type=int)
+    parser.add_argument("--local_rank", type=int)
+    # Model parameters
+    parser.add_argument(
+        "--model_name", default="efficientunet", help="select mode: unet, unetr, swinunetr"
+    )
+    parser.add_argument("--num_class", default=3, type=int, help="segmentation classes")
+    parser.add_argument(
+        "--input_size", default=512, type=int, help="segmentation classes"
+    )
+    # Training parameters
+    parser.add_argument("--batch_size", default=16, type=int, help="Batch size per GPU")
+    parser.add_argument("--max_epochs", default=2000, type=int)
+    parser.add_argument("--val_interval", default=5, type=int)
+    parser.add_argument("--epoch_tolerance", default=100, type=int)
+    parser.add_argument("--initial_lr", type=float, default=1e-4, help="learning rate")
+
+    args = parser.parse_args()
+    torch.cuda.set_device(args.local_rank)
+    torch.distributed.init_process_group(backend='nccl')
+    monai.config.print_config()
+    n_rays = 32
+    pre_trained = True
+    #%% set training/validation split
+    np.random.seed(args.seed)
+    model_path = join(args.work_dir, args.model_name + "_3class")
+    os.makedirs(model_path, exist_ok=True)
+    run_id = datetime.now().strftime("%Y%m%d-%H%M")
+    # This must be change every runing time ! ! ! ! ! ! ! ! ! ! !
+    model_file = "models/flexible_unet_convext.py"
+    shutil.copyfile(
+        __file__, join(model_path, os.path.basename(__file__))
+    )
+    shutil.copyfile(
+        model_file, join(model_path, os.path.basename(model_file))
+    )
+    all_image_path = '/data/louwei/nips_comp/train_cellpose_multi0/'
+    all_img_path = join(all_image_path, "train/images")
+    all_gt_path = join(all_image_path, "train/tif")    
+    
+    all_img_names = sorted(os.listdir(all_img_path))
+    all_gt_names = [img_name.split(".")[0] + ".tif" for img_name in all_img_names]
+    all_img_files = [join(all_img_path, all_img_names[i]) for i in range(len(all_img_names))]
+    all_gt_files = [join(all_gt_path, all_gt_names[i]) for i in range(len(all_img_names))]    
+    img_path = join(args.data_path, "train/images")
+    gt_path = join(args.data_path, "train/tif")
+    val_img_path = join(args.data_path, "test/images")
+    val_gt_path = join(args.data_path, "test/tif")
+    img_names = sorted(os.listdir(img_path))
+    gt_names = [img_name.split(".")[0] + ".tif" for img_name in img_names]
+    train_img_files = [join(img_path, img_names[i]) for i in range(len(img_names))]
+    train_gt_files = [join(gt_path, gt_names[i]) for i in range(len(img_names))]
+    cat_img_files = train_img_files + all_img_files
+    cat_gt_files = train_gt_files + all_gt_files
+    img_num = len(img_names)
+    val_frac = 0.1
+    val_img_names = sorted(os.listdir(val_img_path))
+    val_gt_names = [img_name.split(".")[0] + ".tif" for img_name in val_img_names]
+    #indices = np.arange(img_num)
+    #np.random.shuffle(indices)
+    #val_split = int(img_num * val_frac)
+    #train_indices = indices[val_split:]
+    #val_indices = indices[:val_split]
+
+    train_files = [
+        {"img": cat_img_files[i], "label": cat_gt_files[i]}
+        for i in range(len(cat_img_files))
+    ]
+    val_files = [
+        {"img": join(val_img_path, val_img_names[i]), "label": join(val_gt_path, val_gt_names[i])}
+        for i in range(len(val_img_names))
+    ]
+    print(
+        f"training image num: {len(train_files)}, validation image num: {len(val_files)}"
+    )
+    #%% define transforms for image and segmentation
+    train_transforms = Compose(
+        [
+            LoadImaged(
+                keys=["img", "label"], reader=PILReader, dtype=np.float32
+            ),  # image three channels (H, W, 3); label: (H, W)
+            AddChanneld(keys=["label"], allow_missing_keys=True),  # label: (1, H, W)
+            AsChannelFirstd(
+                keys=["img"], channel_dim=-1, allow_missing_keys=True
+            ),  # image: (3, H, W)
+            #ScaleIntensityd(
+                #keys=["img"], allow_missing_keys=True
+            #),  # Do not scale label
+            SpatialPadd(keys=["img", "label"], spatial_size=args.input_size),
+            RandSpatialCropd(
+                keys=["img", "label"], roi_size=args.input_size, random_size=False
+            ),
+            RandAxisFlipd(keys=["img", "label"], prob=0.5),
+            RandRotate90d(keys=["img", "label"], prob=0.5, spatial_axes=[0, 1]),
+            # # intensity transform
+            RandGaussianNoised(keys=["img"], prob=0.25, mean=0, std=0.1),
+            RandAdjustContrastd(keys=["img"], prob=0.25, gamma=(1, 2)),
+            RandGaussianSmoothd(keys=["img"], prob=0.25, sigma_x=(1, 2)),
+            RandHistogramShiftd(keys=["img"], prob=0.25, num_control_points=3),
+            RandZoomd(
+                keys=["img", "label"],
+                prob=0.15,
+                min_zoom=0.5,
+                max_zoom=2,
+                mode=["area", "nearest"],
+            ),
+            EnsureTyped(keys=["img", "label"]),
+        ]
+    )
+
+    val_transforms = Compose(
+        [
+            LoadImaged(keys=["img", "label"], reader=PILReader, dtype=np.float32),
+            AddChanneld(keys=["label"], allow_missing_keys=True),
+            AsChannelFirstd(keys=["img"], channel_dim=-1, allow_missing_keys=True),
+            #ScaleIntensityd(keys=["img"], allow_missing_keys=True),
+            # AsDiscreted(keys=['label'], to_onehot=3),
+            EnsureTyped(keys=["img", "label"]),
+        ]
+    )
+
+    #% define dataset, data loader
+    check_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
+    check_loader = DataLoader(check_ds, batch_size=1, num_workers=4)
+    check_data = monai.utils.misc.first(check_loader)
+    print(
+        "sanity check:",
+        check_data["img"].shape,
+        torch.max(check_data["img"]),
+        check_data["label"].shape,
+        torch.max(check_data["label"]),
+    )
+
+    #%% create a training data loader
+    train_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
+    # use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=torch.cuda.is_available(),
+    )
+    # create a validation data loader
+    val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
+    val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=1)
+
+    dice_metric = DiceMetric(
+        include_background=False, reduction="mean", get_not_nans=False
+    )
+
+    post_pred = Compose(
+        [EnsureType(), Activations(softmax=True), AsDiscrete(threshold=0.5)]
+    )
+    post_gt = Compose([EnsureType(), AsDiscrete(to_onehot=None)])
+    # create UNet, DiceLoss and Adam optimizer
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if args.model_name.lower() == "unet":
+        model = monai.networks.nets.UNet(
+            spatial_dims=2,
+            in_channels=3,
+            out_channels=args.num_class,
+            channels=(16, 32, 64, 128, 256),
+            strides=(2, 2, 2, 2),
+            num_res_units=2,
+        ).to(device)
+
+    if args.model_name.lower() == "efficientunet":
+        model = FlexibleUNetConvext(
+            in_channels=3,
+            out_channels=n_rays+1,
+            backbone='convnext_small',
+            pretrained=True,
+        ).to(device)
+    
+    if args.model_name.lower() == "swinunetr":
+        model = SwinUNETR(
+            img_size=(args.input_size, args.input_size),
+            in_channels=3,
+            out_channels=n_rays+1,
+            feature_size=24,  # should be divisible by 12
+            spatial_dims=2,
+        ).to(device)
+  
+    #loss_masked_dice = monai.losses.DiceCELoss(softmax=True)
+    loss_dice = monai.losses.DiceLoss(squared_pred=True,jaccard=True)
+    loss_bce = nn.BCELoss()
+    loss_dist_mae = nn.L1Loss()
+    activatation = nn.ReLU()
+    sigmoid = nn.Sigmoid()
+    #loss_dist_mae = monai.losses.DiceCELoss(softmax=True)
+    initial_lr = args.initial_lr
+    encoder = list(map(id, model.encoder.parameters()))
+    base_params = filter(lambda p: id(p) not in encoder, model.parameters())
+    params = [
+        {"params": base_params, "lr":initial_lr},
+        {"params": model.encoder.parameters(), "lr": initial_lr * 0.1},
+    ]
+    optimizer = torch.optim.AdamW(params, initial_lr)
+    #if pre_trained == True:
+        #print('Load pretrained weights...')
+        #checkpoint = torch.load('/mntnfs/med_data5/louwei/nips_comp/swin_stardist/swinunetr_3class/40.pth', map_location=torch.device(device))
+        #model.load_state_dict(checkpoint['model_state_dict'])
+    # start a typical PyTorch training
+    #checkpoint = torch.load("/data2/liuchenyu/log/convnextsmall/efficientunet_3class/510.pth", map_location=torch.device(device))
+    #model.load_state_dict(checkpoint['model_state_dict'])
+    print('distributed model')
+    model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters=True)  
+    print('successful model')
+    max_epochs = args.max_epochs
+    epoch_tolerance = args.epoch_tolerance
+    val_interval = args.val_interval
+    best_metric = -1
+    best_metric_epoch = -1
+    epoch_loss_values = list()
+    metric_values = list()
+    writer = SummaryWriter(model_path)
+    max_f1 = 0
+    for epoch in range(0, max_epochs):
+        model.train()
+        epoch_loss = 0
+        epoch_loss_prob = 0
+        epoch_loss_dist_2 = 0
+        epoch_loss_dist_1 = 0
+        for step, batch_data in enumerate(tqdm.tqdm(train_loader), 1):
+            inputs, labels = batch_data["img"],batch_data["label"]
+            print(step)
+            processes_labels = []
+            
+            for i in range(labels.shape[0]):
+                label = labels[i][0]
+                distances = star_dist(label,n_rays)
+                distances = np.transpose(distances,(2,0,1))
+                #print(distances.shape)
+                obj_probabilities = edt_prob(label.astype(int))
+                obj_probabilities = np.expand_dims(obj_probabilities,0)
+                #print(obj_probabilities.shape)
+                final_label = np.concatenate((distances,obj_probabilities),axis=0)
+                #print(final_label.shape)
+                processes_labels.append(final_label)
+            
+            labels = np.stack(processes_labels)
+
+            #print(inputs.shape,labels.shape)
+            inputs, labels = torch.tensor(inputs).to(device), torch.tensor(labels).to(device)
+            #print(inputs.shape,labels.shape)
+            optimizer.zero_grad()
+            output_dist,output_prob = model(inputs)
+            #print(outputs.shape)
+            dist_output = output_dist
+            prob_output = output_prob
+            dist_label = labels[:,:n_rays,:,:]
+            prob_label = torch.unsqueeze(labels[:,-1,:,:], 1)
+            #print(dist_output.shape,prob_output.shape,dist_label.shape)
+            #labels_onehot = monai.networks.one_hot(
+                #labels, args.num_class
+            #)  # (b,cls,256,256)
+            #print(prob_label.max(),prob_label.min())
+            loss_dist_1 = loss_dice(dist_output*prob_label,dist_label*prob_label)
+            #print(loss_dist_1)
+            loss_prob = loss_bce(prob_output,prob_label)
+            #print(prob_label.shape,dist_output.shape)
+            loss_dist_2 = loss_dist_mae(dist_output*prob_label,dist_label*prob_label)
+            #print(loss_dist_2)
+            loss = loss_prob + loss_dist_2*0.3 + loss_dist_1
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item()
+            epoch_loss_prob += loss_prob.item()
+            epoch_loss_dist_2 += loss_dist_2.item()
+            epoch_loss_dist_1 += loss_dist_1.item()
+            epoch_len = len(train_ds) // train_loader.batch_size
+            
+        epoch_loss /= step
+        epoch_loss_prob /= step
+        epoch_loss_dist_2 /= step
+        epoch_loss_dist_1 /= step
+        epoch_loss_values.append(epoch_loss)
+        print(f"epoch {epoch} average loss: {epoch_loss:.4f}")
+        writer.add_scalar("train_loss", epoch_loss, epoch)
+        print('dist dice: '+str(epoch_loss_dist_1)+' dist mae: '+str(epoch_loss_dist_2)+' prob bce: '+str(epoch_loss_prob))
+        checkpoint = {
+            "epoch": epoch,
+            "model_state_dict": model.module.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "loss": epoch_loss_values,
+        }
+        if epoch < 8:
+            continue
+        if epoch > 1 and epoch % val_interval == 0:
+            torch.save(checkpoint, join(model_path, str(epoch) + ".pth"))
+            model.eval()
+            with torch.no_grad():
+                val_images = None
+                val_labels = None
+                val_outputs = None
+                seg_metric = OrderedDict()
+                seg_metric['F1_Score'] = []
+                for val_data in tqdm.tqdm(val_loader):
+                    val_images, val_labels = val_data["img"].to(device), val_data[
+                        "label"
+                    ].to(device)
+                    roi_size = (512, 512)
+                    sw_batch_size = 4
+                    output_dist,output_prob = sliding_window_inference(
+                        val_images, roi_size, sw_batch_size, model
+                        )
+                    val_labels = val_labels[0][0].cpu().numpy()
+                    prob = output_prob[0][0].cpu().numpy()
+                    dist = output_dist[0].cpu().numpy()
+                    #print(val_labels.shape,prob.shape,dist.shape)
+                    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)
+            
+                    star_label = polygons_to_label(disti, points, prob=probi,shape=prob.shape)
+                    gt = remove_boundary_cells(val_labels.astype(np.int32)) 
+                    seg = remove_boundary_cells(star_label.astype(np.int32))           
+                    tp, fp, fn = eval_tp_fp_fn(gt, seg, threshold=0.5)
+                    if tp == 0:
+                        precision = 0
+                        recall = 0
+                        f1 = 0
+                    else:
+                        precision = tp / (tp + fp)
+                        recall = tp / (tp + fn)
+                        f1 = 2*(precision * recall)/ (precision + recall)
+                    f1 = np.round(f1, 4)
+                    seg_metric['F1_Score'].append(np.round(f1, 4))
+                avg_f1 = np.mean(seg_metric['F1_Score'])
+                writer.add_scalar("val_f1score", avg_f1, epoch)
+                if avg_f1 > max_f1:
+                    max_f1 = avg_f1
+                    print(str(epoch) + 'f1 score: ' + str(max_f1))
+                    torch.save(checkpoint, join(model_path, "best_model.pth"))
+    np.savez_compressed(
+        join(model_path, "train_log.npz"),
+        val_dice=metric_values,
+        epoch_loss=epoch_loss_values,
+    )
+
+
+if __name__ == "__main__":
+    main()

utils.py

@@ -0,0 +1,868 @@
+# 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 cv2
+import math
+import numpy as np
+import torch
+import torch.nn.functional as F
+import colorsys
+import itertools
+import matplotlib.pyplot as plt
+from matplotlib import cm
+
+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,
+)
+
+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
+from skimage.exposure import rescale_intensity
+from skimage.filters import sobel_h, sobel_v, gaussian
+from skimage.morphology import disk, binary_opening
+
+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 xentropy_loss(true, pred, reduction="mean"):
+    """Cross entropy loss. Assumes NHWC!
+
+    Args:
+        pred: prediction array
+        true: ground truth array
+    
+    Returns:
+        cross entropy loss
+
+    """
+    epsilon = 10e-8
+    # scale preds so that the class probs of each sample sum to 1
+    pred = pred / torch.sum(pred, -1, keepdim=True)
+    # manual computation of crossentropy
+    pred = torch.clamp(pred, epsilon, 1.0 - epsilon)
+    loss = -torch.sum((true * torch.log(pred)), -1, keepdim=True)
+    loss = loss.mean() if reduction == "mean" else loss.sum()
+    return loss
+
+
+####
+def dice_loss(true, pred, smooth=1e-3):
+    """`pred` and `true` must be of torch.float32. Assuming of shape NxHxWxC."""
+    inse = torch.sum(pred * true, (0, 1, 2))
+    l = torch.sum(pred, (0, 1, 2))
+    r = torch.sum(true, (0, 1, 2))
+    loss = 1.0 - (2.0 * inse + smooth) / (l + r + smooth)
+    loss = torch.sum(loss)
+    return loss
+
+
+####
+def mse_loss(true, pred):
+    """Calculate mean squared error loss.
+
+    Args:
+        true: ground truth of combined horizontal
+              and vertical maps
+        pred: prediction of combined horizontal
+              and vertical maps 
+    
+    Returns:
+        loss: mean squared error
+
+    """
+    loss = pred - true
+    loss = (loss * loss).mean()
+    return loss
+
+
+####
+def msge_loss(true, pred, focus):
+    """Calculate the mean squared error of the gradients of 
+    horizontal and vertical map predictions. Assumes 
+    channel 0 is Vertical and channel 1 is Horizontal.
+
+    Args:
+        true:  ground truth of combined horizontal
+               and vertical maps
+        pred:  prediction of combined horizontal
+               and vertical maps 
+        focus: area where to apply loss (we only calculate
+                the loss within the nuclei)
+    
+    Returns:
+        loss:  mean squared error of gradients
+
+    """
+
+    def get_sobel_kernel(size):
+        """Get sobel kernel with a given size."""
+        assert size % 2 == 1, "Must be odd, get size=%d" % size
+
+        h_range = torch.arange(
+            -size // 2 + 1,
+            size // 2 + 1,
+            dtype=torch.float32,
+            device="cuda",
+            requires_grad=False,
+        )
+        v_range = torch.arange(
+            -size // 2 + 1,
+            size // 2 + 1,
+            dtype=torch.float32,
+            device="cuda",
+            requires_grad=False,
+        )
+        h, v = torch.meshgrid(h_range, v_range)
+        kernel_h = h / (h * h + v * v + 1.0e-15)
+        kernel_v = v / (h * h + v * v + 1.0e-15)
+        return kernel_h, kernel_v
+
+    ####
+    def get_gradient_hv(hv):
+        """For calculating gradient."""
+        kernel_h, kernel_v = get_sobel_kernel(5)
+        kernel_h = kernel_h.view(1, 1, 5, 5)  # constant
+        kernel_v = kernel_v.view(1, 1, 5, 5)  # constant
+
+        h_ch = hv[..., 0].unsqueeze(1)  # Nx1xHxW
+        v_ch = hv[..., 1].unsqueeze(1)  # Nx1xHxW
+
+        # can only apply in NCHW mode
+        h_dh_ch = F.conv2d(h_ch, kernel_h, padding=2)
+        v_dv_ch = F.conv2d(v_ch, kernel_v, padding=2)
+        dhv = torch.cat([h_dh_ch, v_dv_ch], dim=1)
+        dhv = dhv.permute(0, 2, 3, 1).contiguous()  # to NHWC
+        return dhv
+
+    focus = (focus[..., None]).float()  # assume input NHW
+    focus = torch.cat([focus, focus], axis=-1)
+    true_grad = get_gradient_hv(true)
+    pred_grad = get_gradient_hv(pred)
+    loss = pred_grad - true_grad
+    loss = focus * (loss * loss)
+    # artificial reduce_mean with focused region
+    loss = loss.sum() / (focus.sum() + 1.0e-8)
+    return loss
+
+
+def __proc_np_hv(pred, np_thres, ksize, overall_thres, obj_size_thres):
+    """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 __proc_np_hv_2(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 = rescale_intensity(h_dir_raw, out_range=(0, 1)).astype('float32')
+    v_dir = rescale_intensity(v_dir_raw, out_range=(0, 1)).astype('float32')
+
+    sobelh = sobel_v(h_dir).astype('float64')
+    sobelv = sobel_h(v_dir).astype('float64')
+
+    sobelh = 1 - rescale_intensity(sobelh, out_range=(0, 1)).astype('float32')
+    sobelv = 1 - rescale_intensity(sobelv, out_range=(0, 1)).astype('float32')
+
+    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 = - gaussian(dist, sigma=0.8)
+
+    overall = np.array(overall >= overall_thres, dtype=np.int32)
+
+    marker = blb - overall
+    marker[marker < 0] = 0
+    marker = binary_fill_holes(marker).astype("uint8")
+    kernel = disk(2)
+    marker = binary_opening(marker, 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(
+    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)
+
+    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
+
+        # 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='cpu'))
+                count_map_list.append(torch.zeros([1, 1] + output_shape[2:], dtype=compute_dtype, device='cpu'))
+                _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
+                #print(output_image_list[ss][original_idx_zoom].device,importance_map_zoom.cpu().device,seg_prob.cpu().device)
+                output_image_list[ss][original_idx_zoom] += importance_map_zoom.cpu() * seg_prob[idx - slice_g].cpu()
+                count_map_list[ss][original_idx_zoom] += (
+                    importance_map_zoom.unsqueeze(0).unsqueeze(0).expand(count_map_list[ss][original_idx_zoom].shape).cpu()
+                )
+
+    # 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  # type: ignore
+    if isinstance(inputs, MetaTensor):
+        final_output = convert_to_dst_type(final_output, inputs)[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)

utils_modify.py

@@ -0,0 +1,369 @@
+# 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.big import _grid_divisible, BlockND, OBJECT_KEYS#, repaint_labels
+from stardist.matching import relabel_sequential
+from stardist import dist_to_coord, non_maximum_suppression, polygons_to_label
+from stardist import random_label_cmap,ray_angles
+from stardist 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,
+)
+
+tqdm, _ = optional_import("tqdm", name="tqdm")
+
+__all__ = ["sliding_window_inference"]
+
+
+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)
+        points, probi, disti = non_maximum_suppression(dist,prob,prob_thresh=0.5, nms_thresh=0.4)
+
+        coord = dist_to_coord(disti,points)
+            
+        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)
+
+    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
+
+        # 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)