EAST框架解析

EAST简介

EAST是一个文本检测模型，由于他的实现方式非常简洁且高效，所以吸引了很多人用他的方法做文本检测，具体实现方法已经有很多人讲过，下面将记录一下其中的具体代码处理过程。

EAST论文地址：https://arxiv.org/pdf/1704.03155.pdf
EAST代码复现：https://github.com/argman/EAST

一、图像预处理

 # 读取图片
im_fn = image_list[i]
im = cv2.imread(im_fn)
h, w, _ = im.shape

# 读取标签txt
label_name = os.path.splitext(os.path.basename(im_fn))[0]
label_dir = os.path.join(data_dir, "labels")
txt_fn = os.path.join(label_dir,label_name+".txt") 

text_polys, text_tags = load_annoataion(txt_fn) # 获取点的位置和内容信息
#保证框在图里
text_polys, text_tags = check_and_validate_polys(text_polys, text_tags, (h, w)) 

然后图像根据 background_ratio 分流，切一部分带文字的当前景，切不带文字的当背景

if np.random.rand() < background_ratio: #这些图当背景
	im, text_polys, text_tags =	crop_area(im, text_polys, text_tags, crop_background=True)
	......#一系列背景处理
else: #当前景
#切出的图如果带框，就当前景，数据增强的一种
im, text_polys, text_tags = crop_area(im, text_polys, text_tags, crop_background=True)

score_map, geo_map, training_mask = generate_rbox((new_h, new_w), text_polys, text_tags)
......

geo_map是最复杂的，具体格式如下图，对于0,1，2这每一个像素，都有5个值，ABCDE，分别代表像素点到框的上、下、左、右、框倾斜角度这五个值。包括d的4层和e的1层

具体矩阵如下图所示，

二、模型训练

1）前向计算

EAST的模型，代码中用的是ResNet50 + 多分支预测 的架构，论文中如下图所示

代码中具体像素如下：
然后用这个预测结果与前面标注的score_map, geo_map,结果进行对比，算loss，然后反向传播优化，最终得到一个好的模型。

Loss 计算

Loss的计算，根据论文分为2部分，包括score_map的loss和像素框geo_map的loss

其中：score_loss在代码中使用了dice_lose

 classification_loss = dice_coefficient(y_true_cls, y_pred_cls, training_mask)
# 其中，dice_lose 是
 def dice_coefficient(y_true_cls, y_pred_cls,training_mask):
    import numpy as np
    np.random.normal()

    '''
    dice loss
    coeff![在这里插入图片描述](https://img-blog.csdnimg.cn/20200426233015802.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2xvb3ZlbGo=,size_16,color_FFFFFF,t_70)
icient:协同因素
    :param y_true_cls:  是一个1/0图
    :param y_pred_cls:  是一个概率图，每个像素是一个0~1之间的概率值
    :param training_mask:
    :return:
    '''
    eps = 1e-5 #0.00001
    # reduce_sum是压缩求和，得到一个标量
    intersection = tf.reduce_sum(y_true_cls * y_pred_cls * training_mask) # training_mask都是1，只有文字框小的，模糊的会被屏蔽掉
    union = tf.reduce_sum(y_true_cls * training_mask) + \
            tf.reduce_sum(y_pred_cls * training_mask) + \
            eps
    loss = 1. - (2 * intersection / union) # loss在[-1,1]之间
    tf.summary.scalar('classification_dice_loss', loss)
    return loss

这是Lg的loss，就是如下公式

用图像表示如下：

三、前向推理

推理图像，首先resize到固定的32倍数的大小，如512/720等像素。经过模型后获得 score_map 和geo_map，通过这些信息，先推理出所有的矩形框，

正向预测图像：

（1）根据点到四边的距离，计算一个矩形
（2）根据角度旋转矩形
（3）根据中心点位置，移动到相应的位置

Locality-Aware NMS

局部感知NMS，这个主要就是将相近的先通过权重融合到一起，而不是直接去除，这样省了很多计算量，其中最关键的是：

def weighted_merge(g, p):
   # 取g,p两个几何体的加权（权重根据对应的检测得分计算得到）
   g[:8] = (g[8] * g[:8] + p[8] * p[:8])/(g[8] + p[8])
   
   #合并后的几何体的得分为两个几何体得分的总和
   g[8] = (g[8] + p[8])
   return g

总结：

相比于其他框架，这个简单粗暴，单效果还挺好，除了长文本和弯曲文本，对于普通的短文本非常合适，快而且准确，是非常好的文本检测模型。其中，看源码也有很大的启发，里面很多numpy函数、tf1.x时代的多gpu训练函数等，体现了argman扎实的编程基础，每次读源码都有收获。为这位大神点赞

参考文档：
1、https://github.com/argman/EAST #原始注释
2、https://github.com/zyh8306/EAST #这个代码注释详细、幽默，强烈推荐
3、EAST源码解析
4、【OCR实践系列】EAST理解及实现
5、Dice_Loss简介
6、EAST讲解 求面积的动图很好
7、目标检测之非极大值抑制(NMS)各种变体
8、知乎上EAST讲解
9、目标检测之非极大值抑制(NMS)各种变体