Unsupervised image segmentation by backpropagation算法

文章目录

1.算法主体

无监督图像分割 Unsupervised image segmentation

其中，Net() ，作者使用了一个全卷积网络，接受输入图片完成特征提取，这个网络由三层卷积网络组成，如下：

原作者竟把ReLU放在了BN前，把线性整流放在批归一化前面，会影响BN对方差的调整。

其中，PreSeg() ，（原文为 GetSuperPixels，使用了slic算法），即 使用经典的机器学习无监督语义分割算法对输入图片进行预先分类，比如 Python 的 skimage.segmentation中的多个算法，以及felzenszwalb算法。值得注意的是，作者的原始代码中，为slic算法选取了比较极端的参数，选取这个极端参数是有原因的：

在slic算法中，当分区数n_segments 越高，算法对输入图片的分块越多：

• 由于具有相同语义的像素通常存在于一张图片中的连续区域
• 从而推测：位置相近的像素，很大概率是属于相同语义
• 因此，我们在预分类中，为相邻的像素分配相同的语义标签

2.算法理解

先使用经典的机器学习算法，为输入图片进行『预分类』：调整算法参数，为语义信息明显相同的小区域分配相同的语义标签。由于具有相同语义的像素通常存在于一张图片中的连续区域，所以我们可以假设：颜色接近，纹理接近，且位置接近的像素，必定可以为其分配相同的语义标签。

然后使用深度学习结合自动编码器结构，对输入图片进行分类，分类的目标是：使输出的语义分割结果，尽可能地符合『预分类』的结果。训练到收敛。

最后，深度学习的语义分割结果，会在符合『预分类结果』基础上，将具备相同语义信息的小区块进行合并，得到大区块。

可以观察前面放出的gif图片，我的个人的理解是：深度学习（神经网络）在整个无监督语义分割任务中，承担的任务是：对经典机器学习无监督语义分割的细粒度预分类结果进行处理。并在迭代中，逐步对小区块进行合并，最后得到符合人类预期的语义分割结果。

需要改进的地方，如：『虎而不橘』的虎尾、虎眼，在迭代中被错误地分到了和『草』一样的标签，这不是我们希望看到的结果。

作者的原始代码中，使用随机梯度下降法（SDG）对网络进行训练，并选择了0.1的学习率（默认值是0.001），使得前期的迭代中，算法对像素的合并非常快。

3. 代码改进（仅针对运行效率，使运行时间缩短，不改变主体算法）

• 重写了算法中的三个for循环，（注意我是根据算法重写，并不是修改）
• 修改了经典机器学习无监督图片分类的算法：使用felz算法替代slic算法。
• 修改了卷积网络：使用四层卷积，仿照SENet ，使用3x3与1x1交替，膨胀64 与压缩32
• 将SDG修改为 Hinton 的
• RMSprop，迭代次数大大减少，最终效果精度下降(因此改进的方法依然使用SDG，但是处理大图的时候可自行选用 RMSprop)。

为何我推荐使用felz算法替代slic算法？

在预分类阶段，需要进行足够细粒度的分类，分出足够多的区域（保证该分类的地方被分类到，不该分类的地方，神经网络可以帮助它合并），才能使最终结果更准确。如果分出的类别过多，那么算法需要更多的迭代次数。

使用felz算法替代slic算法，是因为它可以在分出较少区域的情况下，命中更多的『正确边界』，并且felz分出的边界更加准确。无论是选用felz算法或是slic算法，在分出足够多的区域的情况下，对精度影响不大，但迭代次数却有很大不同。我们看图说话吧：

第一列使用了slic算法，分区数n_segments=1000，可以看到虽然分出了很多区域，但是老虎尾巴没有很好地与草丛分开。第二列使用的slic算法，分区数n_segments=100，没有命中我们想要得到的分类边界。

下面是适合了合适参数的预分类算法（对比felz与slic算法）

slic，边界条纹不够精细。而felz算法，它甚至连每一条虎纹都分出来了，这是我推荐这个算法的原因之一。

4. 优化结果（迭代128次，40秒→4秒）

由于修改后的代码，可以使用更少的迭代次数达到相同的效果，因此需要的时间比4秒更短。

测试用的图片

修改（魔改）后，不仅耗时缩短，图片分割质量也有所上升，下面是我随便从 法国自动化所的卫星图片数据集Inria Aerial Image[4] 里面的bellingham_x.tif 截取得到的图片 1000x1000，图中包括了 树林，草地，道路，建筑，以及右下角与能够cosplay草地的湖泊（偏绿色）。

5. 算法缺点

首先，这个算法还不够稳健（robust），算法受参数影响大（包括梯度下降法的参数，与机器学习的预分类算法的参数），并且算法多次随机重启的结果会有不同。为了展示这个缺点，我制作了『橘猫望橘图』：（ 问：这个方案能把老虎和橘子分开吗？答：有时候能，有时候不能，这个是该算法的缺点。）

图中的橘子和橘猫的颜色是一样的。下面三行是我随便调整了参数后得到的不同结果。
结果图中，橘猫颜色的橘色要比橘子浅，是因为橘猫在计算平均像素的时候，把黑色的条纹也纳入计算了，并非 橘猫橘 与 橘子橘 不同。我特地用了PS去证明两种橘色是相同的，结果图中，橘猫的平均颜色比橘子浅，是因为橘猫的平均颜色中包含了黑色的虎纹。深度学习能把橘子与橘猫区分开，很大的原因是卷积网络能比较好的感知纹理的不同，而不仅仅是依靠颜色去分类。

其次，算法还不够成熟，随着迭代，算法会逐步被合并各个分区。但是算法里却没有设定限制，去抑制神经网络对小区域的合并。

6. code

​​https://github.com/Yonv1943/Unsupervised-Segmentation/tree/master​​

import os
import time

import cv2
import numpy as np
from skimage import segmentation

import torch
import torch.nn as nn


class Args(object):
    input_image_path = 'image/00000022.tif'  # image/coral.jpg image/tiger.jpg
    train_epoch = 2 ** 6
    mod_dim1 = 64  #
    mod_dim2 = 32
    gpu_id = 1

    min_label_num = 4  # if the label number small than it, break loop
    max_label_num = 256  # if the label number small than it, start to show result image.


class MyNet(nn.Module):
    def __init__(self, inp_dim, mod_dim1, mod_dim2):
        super(MyNet, self).__init__()

        self.seq = nn.Sequential(
            nn.Conv2d(inp_dim, mod_dim1, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(mod_dim1),
            nn.ReLU(inplace=True),

            nn.Conv2d(mod_dim1, mod_dim2, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(mod_dim2),
            nn.ReLU(inplace=True),

            nn.Conv2d(mod_dim2, mod_dim1, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(mod_dim1),
            nn.ReLU(inplace=True),

            nn.Conv2d(mod_dim1, mod_dim2, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(mod_dim2),
        )

    def forward(self, x):
        return self.seq(x)


def run():
    start_time0 = time.time()

    args = Args()
    torch.cuda.manual_seed_all(1943)
    np.random.seed(1943)
    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id)  # choose GPU:0
    image = cv2.imread(args.input_image_path)

    '''segmentation ML'''
    seg_map = segmentation.felzenszwalb(image, scale=64, sigma=0.5, min_size=64)
    # seg_map = segmentation.slic(image, n_segments=10000, compactness=100)
    seg_map = seg_map.flatten()
    seg_lab = [np.where(seg_map == u_label)[0]
               for u_label in np.unique(seg_map)]

    '''train init'''
    device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')

    tensor = image.transpose((2, 0, 1))
    tensor = tensor.astype(np.float32) / 255.0
    tensor = tensor[np.newaxis, :, :, :]
    tensor = torch.from_numpy(tensor).to(device)

    model = MyNet(inp_dim=3, mod_dim1=args.mod_dim1, mod_dim2=args.mod_dim2).to(device)
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=5e-2, momentum=0.9)
    # optimizer = torch.optim.RMSprop(model.parameters(), lr=1e-1, momentum=0.0)

    image_flatten = image.reshape((-1, 3))
    color_avg = np.random.randint(255, size=(args.max_label_num, 3))
    show = image

    '''train loop'''
    start_time1 = time.time()
    model.train()
    for batch_idx in range(args.train_epoch):
        '''forward'''
        optimizer.zero_grad()
        output = model(tensor)[0]
        output = output.permute(1, 2, 0).view(-1, args.mod_dim2)
        target = torch.argmax(output, 1)
        im_target = target.data.cpu().numpy()

        '''refine'''
        for inds in seg_lab:
            u_labels, hist = np.unique(im_target[inds], return_counts=True)
            im_target[inds] = u_labels[np.argmax(hist)]

        '''backward'''
        target = torch.from_numpy(im_target)
        target = target.to(device)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

        '''show image'''
        un_label, lab_inverse = np.unique(im_target, return_inverse=True, )
        if un_label.shape[0] < args.max_label_num:  # update show
            img_flatten = image_flatten.copy()
            if len(color_avg) != un_label.shape[0]:
                color_avg = [np.mean(img_flatten[im_target == label], axis=0, dtype=np.int) for label in un_label]
            for lab_id, color in enumerate(color_avg):
                img_flatten[lab_inverse == lab_id] = color
            show = img_flatten.reshape(image.shape)
        cv2.imshow("seg_pt", show)
        cv2.waitKey(1)

        print('Loss:', batch_idx, loss.item())
        if len(un_label) < args.min_label_num:
            break

    '''save'''
    time0 = time.time() - start_time0
    time1 = time.time() - start_time1
    print('PyTorchInit: %.2f\nTimeUsed: %.2f' % (time0, time1))
    cv2.imwrite("seg_%s_%ds.jpg" % (args.input_image_path[6:-4], time1), show)


if __name__ == '__main__':
    run()