Pytorch extract skeleton (differentiable)

Preface

There are many ready-made packages for extracting skeletons , The simplest and most direct is ：

from skimage.morphology import skeletonize, skeletonize_3d

But today we will introduce another method of extracting skeleton ！ It can also be understood as refinement
Use pytorch The goal is , This process is differentiable , In other words , You can reverse the gradient , For network prediction mask, You can extract the skeleton through this function , Then constrain the topology of the object on the skeleton

This method comes from the literature CVPR 2021：clDice - a Novel Topology-Preserving Loss Function for Tubular Structure Segmentation
Source code ：https://github.com/jocpae/clDice
（ The source code is very, very simple , This is one of the few so concise CVPR The source ）

Source code

import torch
import torch.nn as nn
import torch.nn.functional as F

def soft_erode(img):
    if len(img.shape)==4:
        p1 = -F.max_pool2d(-img, (3,1), (1,1), (1,0))
        p2 = -F.max_pool2d(-img, (1,3), (1,1), (0,1))
        return torch.min(p1,p2)
    elif len(img.shape)==5:
        p1 = -F.max_pool3d(-img,(3,1,1),(1,1,1),(1,0,0))
        p2 = -F.max_pool3d(-img,(1,3,1),(1,1,1),(0,1,0))
        p3 = -F.max_pool3d(-img,(1,1,3),(1,1,1),(0,0,1))
        return torch.min(torch.min(p1, p2), p3)

def soft_dilate(img):
    if len(img.shape)==4:
        return F.max_pool2d(img, (3,3), (1,1), (1,1))
    elif len(img.shape)==5:
        return F.max_pool3d(img,(3,3,3),(1,1,1),(1,1,1))

def soft_open(img):
    return soft_dilate(soft_erode(img))

def soft_skel(img, iter_):
    img1  =  soft_open(img)
    skel  =  F.relu(img-img1)
    for j in range(iter_):
        img  =  soft_erode(img)
        img1  =  soft_open(img)
        delta  =  F.relu(img-img1)
        skel  =  skel +  F.relu(delta-skel*delta)
    return skel

This is all the code for extracting the skeleton in this article , It's really simple

test

Take a CVPPP Of label For example ：

To test the code that extracts the skeleton , I take the boundary of the above image as the foreground , To refine the boundary .
Use the following function to extract the boundary （ Expand three times at the same time ）：

import numpy as np
from PIL import Image
from scipy import ndimage
from skimage.segmentation import find_boundaries

def extract_boundary(label, dilation=False):
    boundary = np.zeros_like(label, dtype=np.uint8)
    ids, counts = np.unique(label, return_counts=True)
    for i, id in enumerate(ids):
        if id == 0:
            boundary[label == 0] = 0
        else:
            tmp = np.zeros_like(label)
            tmp[label == id] = 1
            tmp_bound = find_boundaries(tmp != 0, mode='outer')
            if dilation:
                tmp_bound = ndimage.binary_dilation(tmp_bound, iterations=3, border_value=1)
            boundary[tmp_bound == 1] = 1
    return boundary

Be ready , Upper main function ：

if __name__ == "__main__":
    label = np.asarray(Image.open('plant154_label.png'))

    boundary = extract_boundary(label, dilation=True)
    boundary_uint8 = boundary.astype(np.uint8) * 255
    Image.fromarray(boundary_uint8).save('boundary.png')

    boundary = boundary.astype(np.float32)
    boundary = boundary[np.newaxis, np.newaxis, ...]
    boundary = torch.from_numpy(boundary)

    skel = soft_skel(boundary, 10)
    skel = np.squeeze(skel.numpy())
    skel = (skel * 255).astype(np.uint8)
    Image.fromarray(skel).save('skel.png')

Running results ：
The boundary visualization results before extracting the skeleton ：

Visualization results of extracted skeleton ：

use skimage.morphology Medium skeletonize The result of function extraction ：

summary

1. The proposed method is different from the traditional skeletonize（skimage） There is still a certain gap , especially , From the visualization results above , We can also see that , There are many faults in the estimation extracted by this method , This is unbearable for the skeleton
2. This method has a super parameter , That is, specify the number of iterations , This parameter also greatly affects the result , The above example is used 10, I tried 5 The situation of , The generated skeleton is even worse
3. But this method still has one biggest advantage , It is differentiable , Derivable , It can be used to output the skeleton in the network and then as a loss