Boundry attention: 泛化能力很强的边缘检测模块

原文链接：Boundary attention:: Learning to Find Faint Boundaries at Any Resolution

提出像素级别的表示junction space,来表示像素的局部几何结构。并通过attention来学习该表示。
本文提出的模型泛化性好，即使只在合成的简单图案上训练，在复杂的真实图上做测试也能得到较好的边界。

细节部分：
使用了邻居交叉注意力 neighborhood cross-attention。

不同于viT把图片切成小patch，然后映射为token，而是每个像素都有一个token。
（文章说的dense，stride-1的token）

每个像素的token，设计为表达局部几何结构的特征空间，表征以自己为中心的正方形区域内的局部结构。

提出bottleneck：将图像的token映射到低维空间，这里设计这个空间为junction space。（受Field of Junctions的启发）

junction space：

用一些参数，能代表大范围的几何边缘结构。比如什么都没有，边，角，T结，Y结。（这些几何结构够了吗？能不能通过kmeans聚类得到最佳的参数空间大小？）

3.1 表达边界的元参数

g=(u, theta, omega) 定义对局部空间的分割。

1. 其中点位置u=（u，v），相对patch中的中心像素的；

2. 线的方向theta；线顺时针旋转产生另外两条线，夹角依次为omega1、omega2、omega3（因为最多产生三条线，所以最多表达三部分，比如T、Y结）。
   ⁃ 这些参数表达的几何结构种类是有限的，每一种叫做 wedge。

f-n代表输入图像
$s_{nj}(x;g)$为二值量。代表像素x在像素n的领域内，且属于第j个wedge，就=1.（j能离散？）（具体函数在附录，代表）
$d_n(x;g)>=0$ 代表像素x离g中的最近的边界的欧式距离。这里的边界就是jspace中的线。
图中为junction例子，最右边是distance space.

3. 每个像素表示的region的大小，设计成自适应的。用二值函数window function：$w_n(x;p)$表示。这个函数定义为一个凸函数：分三个级别衡量size（D=3，9，17）。如果距离中心像素的大小小于某个size，就加上对应的系数。

3.2 算子

网络通过不停迭代 { g t [ n ] , p t [ n ] } \{g^t[n], p^t[n]\} {gt[n],pt[n]} 。结合两个算子来计算。

gather operator：每个局部的wedge计算其包含的像素值。新的图的特征fkj为，window内包括的像素值的加权平均。
⁃ input：feature图、window function图p、wedge图g。
⁃ output：更新的feature

slice operator：对每个像素，计算包含他在内的所有region的某个值的mean，variance。这个值有两个：distance map和wedge feature。
在这里对每个patchk要进行加权计算，对他有贡献的是包含这个patch中心点n的所有patches，不只他自己。
⁃ input：distance map d，window function图p，wedge图g；feature map一样的
⁃ output：distance map d的均值、方差；feature map的均值、方差。
这两个operator得到distance map、feature的方差。

3.3 可视化输出

期望重叠区（每个pixel有一个region）的边界是一样的。所以让上面得到的两个输出方差很小。

对上面distance map和wedge feature经过attention输出的均值的：
⁃ 像素n的feature mean是对包含他的领域中每个wedge feature的权重求和。（attention是对邻居的不同的注意力）

图片解释：
⁃ 第二行的图，对应选取的区域，颜色都一样，所以他的第一个window function很大？对应window中的边缘（第二张图）是快到花瓣边缘了。

4 网络结构

整个网络要迭代g，p表示。迭代window function的系数和junction 空间的表示。为什么不能直接计算？因为直接计算太难了？是，因为根本无法计算，只能让网络迭代得到最好的g（关注该范围的参数是多少）最好的p（关注多大的范围），以及考虑邻居信息的feature

neighbor MLP-Mixer

网络输入像素，经过MLPmixer输出hidden state。

boundry attention

然后经过attention模块，8次迭代g和p（wedge feature和window size参数）学习hidden state

每个boundry attention：

使用跳跃连接；
把hidden state copy两份，一份+window token（${\pi }^t$），一份+global feature $f^t$。分别经过MLP。
window token作为Q，global feature作为KV。
window token（size， 1）：前面提到过每个pixel n都作为query。
global feature分割为11大小的patch，两者做cross-attention：每个像素n（query）的考虑了邻居的特征的新特征是什么？

还给global feature的patch加了positional encoding。

attention输出的hidden state，其中后8维是window token。他们两个经过线性decoder，映射回junction space的参数表示g。得到的g进行gather和slice opartor。

4.1 训练过程：

一共有一个MLPmix，两个boundry attention。总共进行8次迭代。
第一阶段迭代4次：线MLP-mixer和第一个attention一起训练；在第3，4迭代使用loss
第二阶段在迭代4次：然后再加第二个attention进行端到端的训练。在第7，8迭代使用loss
最后一次迭代的loss比前几次迭代的都权重大。
（一开始输入-MLPmixer，attention-再返回给attention-在返回给attention（此时计算loss1）-在返回给attention（此时计算loss2），加上第二个attetnion，返回头部端到端输出-在返回头部端到端输出-在返回头部端到端输出（此时计算loss3）-在返回头部端到端输出（此时计算loss4），最后loss应该是loss1-4的和）
（可能因为全部端到端训练，太难了？）

loss12是监督loss：包括feature和distance map的GT（这个他怎么得到的？）

每个loss包括像素级和patch级别的feature loss和distance map loss。让输入接近GT。还有他们对应的方差loss，让其尽量为0.

实验细节：
合成有15-20种几何图形的图片，变成灰度图，参见裁剪25*125大小，作为输入。
数据集10^5，90%作为训练。
⁃ 给图片加noise，让网络实现得到和无噪声一样的feature map，distance map。

结果：
1. 验证对噪声的抗干扰性能
和EDTRE、去噪的EDTRE对比
和类似原理但不是深度学习的filed of junction做对比。

2.	验证去噪检测时仍然存在的亚像素精度

在小image上检测，上采样放大边缘也很好。

3.	可视化hidden state，发现确实学到了junction space。而且直接对值进行差值，能得到连续的junction（可视化就是直接把hiddenstate值作为参数进行画？

4.	在真实图像上

没和其他网络比较对真实图像的检测，而是比较了加噪声后的表现。

5.	推理时间等等

补充材料：
1. junction space
junction空间的几何表示，可以有M种，本文用的是M=3

参数空间有一些等价表示。
这里没看懂

Mjunction是更大的M‘的子集（训练时还可以提高M试一试）

在Mjucntion参数空间中，能直接计算出central 方向和边界boundary方向

和filed of junction那篇在junction space上的比较：参数空间多了oemga，这样就不用限制fai的范围。

这个参数空间的表示允许线性插值

2.	训练数据

合成数据集Kaleidoshapes

3.	模型细节

输入3通道图像，每个像素映射为64维。
经过一个zero-padding处理（补充边界）。
然后MLPmix：两个neighbor mixing blocks：每个包括空间patch mixer和channelmixer。前者是用两个3*3空间卷积+GELU实现。
channelmixer是每个像素输入的MLP。
最后剪掉包含了原图之外信息的patches。

训练细节：
第一阶段
训练只用kal数据集的简化版（100*100大小，加了高斯噪声，）用等式9，10的loss。让网络学到有意义的hidden state gamma，避免jucntion的collapse。让网络关注边界一致性loss（方差loss-0），学习无边界patch。

数据存在不平衡：只有一部分有边界信息。所以加一个额外的空间importance mask给有边界的区域更多关注。（用过gaussian产生这个mask？）把mask以常数C加在loss上。

第二阶段：
加另一个block，用全部iamge和loss。
也用了空间importance mask