OCR-(DB+CRNN)-代码分析-实践环境配置及运行

代码下载及环境配置

首先是代码下载,DBNet代码

使用 git clone 的方式下载上述代码即可。

使用conda来install一些必要的库

我对代码文件夹中的README.MD文件做了一个简化,只需要按照接下来的步骤按照虚拟环境即可,若中途出现超时报错,请多换几次源(对此深感无奈)。

先执行一下命令,构建一个初步的虚拟环境

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
conda create -n dbnet python=3.6
activate dbnet
conda install ipython pip
pip anyconfig==0.9.10
pip future==0.18.2
pip imgaug==0.4.0
pip matplotlib==3.1.2
pip numpy==1.17.4
pip opencv-python==4.1.2.30
pip Polygon3==3.0.8
pip pyclipper==1.1.0.post3
pip PyYAML==5.2
pip scikit-image==0.16.2
pip tensorboard==2.1.0
pip tqdm==4.40.1
pip install natsort
pip install addict
pip install cudatoolkit==10.2.89 #这个需要根据自己的电脑配置来确定

接下来的几个库,需要我们下载.whl来辅助下载,否则会非常慢,甚至找不到下载资源。

他们分别是这几个库

1
2
3
4
cudnn-7.6.5-cuda10.2_0 #这个需要根据自己的电脑配置来确定
Shapely-1.6.4.post2-cp36-cp36m-win_amd64
torch-1.7.1-cp36-cp36m-win_amd64
torchvision-0.8.2-cp36-cp36m-win_amd64

我在这里放一下百度网盘的连接

链接:https://pan.baidu.com/s/1r1mpHcDWKvNwUqR60h5-Vw
提取码:8ypv

训练步骤

首先需要修改一些配置文件

进入git clone好的代码文件中。

首先看config文件夹下的一些配置文件:这些配置文件中的内容包括了backbone使用何种网络模型来进行特征提取,

是否采用dcn空洞卷积来增大感受野,训练数据和测试数据的一个存放位置,是否启用预训练模式,epochs,各种参数等…

接着查看tools文件夹下的train.py

将原来的”config/…….”配置文件的相对路径改为具体采用的配置方案的配置文件的绝对路径

例如我采用的是icdar2015_resnet18_FPN_DBhead_polyLR.yaml

改好之后,按照图片中的格式,修改所选择的配置文件中一些关于数据的存放路径(注:train.txt和test.txt文件夹还未生成,先按代码文件夹下的README.MD,在datasets文件夹下,分别放好train,test文件夹下的gt和img)

数据集下载链接

接下来按照load.py代码,生成train.txttest.txt两个文件

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import os
def get_images(img_path):
'''
find image files in data path
:return: list of files found
'''
img_path = os.path.abspath(img_path)
files = []
exts = ['jpg', 'png', 'jpeg', 'JPG', 'PNG']
for parent, dirnames, filenames in os.walk(img_path):
for filename in filenames:
for ext in exts:
if filename.endswith(ext):
files.append(os.path.join(parent, filename))
break
print('Find {} images'.format(len(files)))
return sorted(files)

def get_txts(txt_path):
'''
find gt files in data path
:return: list of files found
'''
txt_path = os.path.abspath(txt_path)
files = []
exts = ['txt']
for parent, dirnames, filenames in os.walk(txt_path):
for filename in filenames:
for ext in exts:
if filename.endswith(ext):
files.append(os.path.join(parent, filename))
break
print('Find {} txts'.format(len(files)))
return sorted(files)

if __name__ == '__main__':
import json
#img_path = './data/ch4_training_images'
#img_path = './train/img'
img_path = './test/img'
files = get_images(img_path)
#txt_path = './data/ch4_training_localization_transcription_gt'
#txt_path = './train/gt'
txt_path = './test/gt'
txts = get_txts(txt_path)
n = len(files)
assert len(files) == len(txts)
with open('test.txt', 'w') as f:
for i in range(n):
line = files[i] + '\t' + txts[i] + '\n'
#line = files[i] + ' ' + txts[i] + '\n'
f.write(line)
print('dataset generated ^_^ ')

运行tools文件夹下的train.py,即可开始训练

训练结果model_best.pth位置如下

预测步骤

将红框中的路径,全部改为对应的model_best.pth,准备预测的图片,和预测结果的存放文件的绝对地址。

注意第二个红框,原代码中拼写有误。进行修改之后即可运行predict.py,然后即可在输出文件中查看训练结果。

整体网络模型构建代码

models/model.py

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class Model(nn.Module):
def __init__(self, model_config: dict):
"""
PANnet
:param model_config: 模型配置
"""
super().__init__()
model_config = Dict(model_config)
backbone_type = model_config.backbone.pop('type')
neck_type = model_config.neck.pop('type')
head_type = model_config.head.pop('type') #**将字典model_config中的backbone键的值全部解析出来
# backbone主干特征提取网络(和普通的resnet18网络模型一样,只是可以选择是否使用空洞卷积dcn来扩大感受野)
# 主干网络大多时候指的是提取特征的网络,其作用就是提取图片中的信息,供后面的网络使用
self.backbone = build_backbone(backbone_type, **model_config.backbone)

# neck使用的是FPN特征金字塔(FPN是一种利用常规CNN模型来高效提取图片中各维度特征的方法。)
# 其中neck的in_channels是由backbone的out_channels传进来的
# neck是放在backbone和head之间的,是为了更好的利用backbone提取的特征
self.neck = build_neck(neck_type, in_channels=self.backbone.out_channels, **model_config.neck)

# head是获取网络输出内容的网络,利用之前提取的特征,head利用这些特征,做出预测。
self.head = build_head(head_type, in_channels=self.neck.out_channels, **model_config.head)

# 给参数配置文件取一个名
self.name = f'{backbone_type}_{neck_type}_{head_type}'

# 前向传播
def forward(self, x):
_, _, H, W = x.size()
backbone_out = self.backbone(x)
neck_out = self.neck(backbone_out)
y = self.head(neck_out)
#为了保险起见,再对y进行一次插值操作,使其尺寸与原图相同。
y = F.interpolate(y, size=(H, W), mode='bilinear', align_corners=True)
return y # 整个模型最终得到的就是从head模型中输出的y
# 这个y最终会给到tools文件夹下的predict.py里的preds,
# 这个preds会喂给post_processing文件夹下的SegDetectorRepresenter()
# 并通过这个post_processing来获得后处理得到的文本框
# 如果是训练时,y会传给tools文件夹中的train.py文件

models/backbone/resnet.py Resnet18的网络模型构建

特征提取

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
# resnet18和resnet34用的是BasicBlock,
class BasicBlock(nn.Module):
expansion = 1
# 卷积步长stride = 1,扩张大小dilation = 1(也就是padding),
# in_planes和out_planes分别是输入和输出的通道数,
def __init__(self, inplanes, planes, stride=1, downsample=None, dcn=None):
super(BasicBlock, self).__init__()
self.with_dcn = dcn is not None # dcn参数指deformable convolution(空洞卷积)的参数设置
self.conv1 = conv3x3(inplanes, planes, stride) # 普通尺度不变的卷积核
self.bn1 = BatchNorm2d(planes) #BN
self.relu = nn.ReLU(inplace=True) #Relu
self.with_modulated_dcn = False
if not self.with_dcn: # 不使用dcn空洞卷积
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)
else: # 使用dcn空洞卷积
from torchvision.ops import DeformConv2d
deformable_groups = dcn.get('deformable_groups', 1)
offset_channels = 18
self.conv2_offset = nn.Conv2d(planes, deformable_groups * offset_channels, kernel_size=3, padding=1)
self.conv2 = DeformConv2d(planes, planes, kernel_size=3, padding=1, bias=False)
self.bn2 = BatchNorm2d(planes)
self.downsample = downsample
self.stride = stride

def forward(self, x):
residual = x

out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)

# out = self.conv2(out)
if not self.with_dcn: #如果不使用dcn空洞卷积
out = self.conv2(out)
else:
offset = self.conv2_offset(out)
out = self.conv2(out, offset)
out = self.bn2(out)
# 缩小图像(或称为下采样(subsampled)或降采样(downsampled))的主要目的是两个:
# 使得图像符合显示区域的大小;
# 生成对应图像的缩略图;
# 目的是为了让residual和out的尺度相同,只有相同才能进行out += residual操作
if self.downsample is not None:
residual = self.downsample(x)

out += residual
out = self.relu(out)

return out
#==================================================================================================
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)

x2 = self.layer1(x)
x3 = self.layer2(x2)
x4 = self.layer3(x3)
x5 = self.layer4(x4)
# 这里的x2...5对应的就是原图尺寸1/2,1/4,1/8,1/16,1/32的特征向量
# 会作为backbone_out传入neck中,在FPN特征金字塔做upsample_add和upsample_cat
return x2, x3, x4, x5

models/neck/FPN.py FPN特征金字塔

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def forward(self, x):
c2, c3, c4, c5 = x
# 此处的x含有backbon的各层信息
# 这里将x的1/32,1/16,1/8,1/4各层分别交给c5,c4,c3,c2
# 这里的1/32,1/16,1/8,1/4指的是原图尺寸的多少份之一
# Top-down
# 经过reduce_conv_c5(c5),p5尺度不变,channels变为256
p5 = self.reduce_conv_c5(c5)
# 经过_upsample_add(p5, self.reduce_conv_c4(c4)),
# p4是由“‘c5增加通道数为256之后得到的p5'和'增加通道数为256的c4’进行一个相加,并上采样得到的”
p4 = self._upsample_add(p5, self.reduce_conv_c4(c4))
p4 = self.smooth_p4(p4)
p3 = self._upsample_add(p4, self.reduce_conv_c3(c3))
p3 = self.smooth_p3(p3)
p2 = self._upsample_add(p3, self.reduce_conv_c2(c2))
p2 = self.smooth_p2(p2)
# 这里的p4,p3,p2是对应neck网络模型中的3个部分

#x此时获得了{p2,p3,p4,p5先upsample到和p2一样原图1/2的尺寸,然后再进行conv-cat融合}
x = self._upsample_cat(p2, p3, p4, p5)
# 融合几个张量后再卷积
x = self.conv(x)
return x

#上采样过程是用pytorch库中的import torch.nn.functional as F中的interpolate完成的,将上采样的结果再与相同尺寸的数据进行相加
def _upsample_add(self, x, y):
return F.interpolate(x, size=y.size()[2:]) + y
#将x的size插值为y的尺寸(使得尺寸相同)(upsample),再与y相加(add),合起来就是upsample_add

#通过上采样将p3,p4,p5的尺寸均插值为p2的1/4尺寸。将结果赋值给x,之后x再进行一个conv(x)
def _upsample_cat(self, p2, p3, p4, p5):
h, w = p2.size()[2:]
p3 = F.interpolate(p3, size=(h, w))
p4 = F.interpolate(p4, size=(h, w))
p5 = F.interpolate(p5, size=(h, w))
# 再进行p2,p3,p4,p5这几个张量tensor形式的图进行融合,
# 并且是按维数1拼接(横着拼)
return torch.cat([p2, p3, p4, p5], dim=1)
# cat是concatnate的意思:拼接,联系在一起。
# 按维数1拼接(横着拼),将多个tensor向量进行拼接联系起来

models/head/DBHead.py DBNet核心模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class DBHead(nn.Module):
def __init__(self, in_channels, out_channels, k = 50):
super().__init__()
self.k = k

# binarize把从backbone和neck进行特征融合后的原图尺寸1/4的结果进行卷积操作
# 之后得到的probability map(概率图)
self.binarize = nn.Sequential(
nn.Conv2d(in_channels, in_channels // 4, 3, padding=1),
nn.BatchNorm2d(in_channels // 4),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(in_channels // 4, in_channels // 4, 2, 2),
nn.BatchNorm2d(in_channels // 4),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(in_channels // 4, 1, 2, 2),
nn.Sigmoid())
self.binarize.apply(self.weights_init)

# thresh把从backbone和neck进行特征融合后的原图尺寸1/4的结果
# 用来计算threshold map
self.thresh = self._init_thresh(in_channels)
self.thresh.apply(self.weights_init)

def forward(self, x): # x是从neck中的到的1/4原图尺寸的一个张量结果,也就是neck_out
# shrink_maps也就是probability map
# 将从backbone和neck结合后得到的x再分别传入binarize函数,和thresh函数
# 得到probability map和threshold_maps
shrink_maps = self.binarize(x) #probability map其实就是shrink_map
threshold_maps = self.thresh(x)
# 如果是在训练过程中使用模型,则self.training为True
# 则将进行一个可微分的二值化操作,通过probability map和threshold_maps,使用DB函数进行计算得到结果
if self.training:
# 则将进行一个可微分的二值化操作,通过probability map和threshold_maps,
# 使用step_function函数(DB函数进)行计算得到结果
# 将该结果赋值给binary_maps,并在融合时,除了shrink_maps,threshold_maps
# 还需要融合binary_maps
binary_maps = self.step_function(shrink_maps, threshold_maps)
# 再进行两个张量tensor形式的图进行融合,并且是按维数1拼接(横着拼)
y = torch.cat((shrink_maps, threshold_maps, binary_maps), dim=1)
else: # 预测过程
# 再进行两个张量tensor形式的图进行融合,并且是按维数1拼接(横着拼)
y = torch.cat((shrink_maps, threshold_maps), dim=1)
return y # y是最终的一个预测结果
1
2
3
def step_function(self, x, y):
return torch.reciprocal(1 + torch.exp(-self.k * (x - y)))
# 这个step_function就是Differentiable binarization的公式

流程:主model.py最终从backbone–neck–DB得到的y,如果是训练,则传给train.py;如果是预测,则传给predict.py。

训练部分代码

tools/train.py

一些训练的设置,重点是其在评估方法时,采用的是quad_metric.py中的方法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# 后处理,获取文本框,这里也是一个重要的步骤
post_p = get_post_processing(config['post_processing'])

# 在utils文件夹下的ocr_metric中,传入metric配置信息,实例化一个metric
# 从配置信息可以看处,config['metric']中配置的信息是
# utils文件夹下,ocr_metric文件夹下,icdar2015文件夹下的quad_metric.py文件
# 所使用的方法也就是quad_metric
# 用法就是:每多少个epoch,调用其measure函数,就进行一下评估,
# 评估的东西可以从trainer.train()中看出,是boxes和scores两个东西,
# metric_cls.validate_measure(batch, (boxes, scores))
# boxes, scores = self.post_process(batch, preds,is_output_polygon=self.metric_cls.is_output_polygon)
# 这两个是通过后处理post_process得到的"构成原文本框的的所有点boxes"和"一个文本框的概率均值scores"

# 当然不止,quad_metric总共有三个评估措施
# validate_measure(),evaluate_measure(),gather_measure()
# 返回一个result
metric = get_metric(config['metric'])

trainer/trainer.py 重点:其中评估方法quad_metric.py文件

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# preds:从网络模型的head中输出的一个二维的包含probability map和threshold map两个张量的张量
preds = self.model(batch['img'])

def _eval(self, epoch):
#==========================================================================================
# 将每一个批量的boxes(预测框的点)和scores传入metric_cls中的评估方法中
# 这个metric_cls是在utils文件夹下,ocr_metric文件夹下,
# icdar2015文件夹下的quad_metric.py文件中的方法
# 疯狂套娃,各种递归传回,太累了
# 这里raw_metric最终得到的是经过最终在utils/ocr_metric/icdar2015/detection/iou.py文件传回的
# 含有recall, precision,iou等一系列的值
# 这一系列的值用一个字典的形式传回给raw_metric

# raw_metric是从iou.py文件中返回的一个batch的评估信息字典
raw_metric = self.metric_cls.validate_measure(batch, (boxes, scores))
raw_metrics.append(raw_metric) # 有很多个batch
# 用quad_metric.py文件中的gather_measure函数
# 求整个数据集的recall,precision等评估结果
# metrics 才是整个数据集的评估信息字典
# 这里的metrics是调用了quad_metric.py文件中的gather_measure方法
metrics = self.metric_cls.gather_measure(raw_metrics)

训练的batch以及总体的评估

utils/ocr_metric/icdar2015/quad_metric.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
class QuadMetric():
def __init__(self, is_output_polygon=False):
self.is_output_polygon = is_output_polygon
#------------------------------------------------------------------------------------
# DetectionIoUEvaluator是评估器(核心,重点),最终的评估是在这里面完成
# 然后将结果返回给evaluator
# 这个DetectionIoUEvaluator方法是utils文件夹下,ocr_metric文件夹下,icdar2015文件夹下,
# 的detection文件夹下的iou.py文件中的一个检测评估类
self.evaluator = DetectionIoUEvaluator(is_output_polygon=is_output_polygon)
#------------------------------------------------------------------------------------

# 主方法,下面三个方法其实都是调用measure的
def measure(self, batch, output, box_thresh=0.6):
# 从predict.py中传入的batch是一个字典,[(h,w)]的键是shape
# h和w代表了原始图片的高度和宽度
# batch = {'shape': [(h, w)]}
# 也就是第二个batch

# 从trainer.py中传入的batch可以理解为笼统的batch含义,也就是若干图片

results = []
# 表示有若干个图片,每个图片都有若干个文字框的gt,点坐标
# 正确的标注就是是groundtruth,用图像标注工具标注的
gt_polyons_batch = batch['text_polys']
ignore_tags_batch = batch['ignore_tags']


# for循环遍历每一张图片,得到这张图片的所有多边形框
# 包括gt,预测框,分数等内容
for polygons, pred_polygons, pred_scores, ignore_tags in zip(gt_polyons_batch, pred_polygons_batch, pred_scores_batch, ignore_tags_batch):
# 将gt整理成一个列表,每个元素是个字典,代表对应文字框的点坐标和是否忽略
gt = [dict(points=np.int64(polygons[i]), ignore=ignore_tags[i]) for i in range(len(polygons))]
if self.is_output_polygon: # 如果要输出的是多边形
pred = [dict(points=pred_polygons[i]) for i in range(len(pred_polygons))]
else: # 如果输出的是矩形,则还要做一个过滤(if pred_scores[i] >= box_thresh(0.6))
# 意思是说:如果矩形框(文本框)内的文字区域如果是一个极度细长且极度弯曲的文字区域,那么
# 即使文字区域概率全为1,非文字区域概率全为0,和gt(手工标记的)结果是一样的
# 但由于pred_scores[i]是计算得到这个文本框内的概率均值,它是由
# 文字区域的概率,和非文字区域的概率共同决定的一个均值
# 因为这个矩形文本框中,绝大多数都是非文字区域,
# 所以计算出来的文本框概率均值可能就会小于box_thresh
# 如果小于,也就说明这个文字区域是极度弯曲细长的,这样的文本框送给crnn做识别
# 或者其实给人肉眼看,都没有什么实际意义,所以就将其忽略。

# 如果大于box_thresh(0.6),则将其放入pred列表中

# 将上述的经过极端情况过滤后的pred列表,以及gt列表,一起传入evaluator方法中
# 来计算一个batch的
# 经过for循环,将每个batch中的评估信息,
# 以字典的形式传入results列表中
results.append(self.evaluator.evaluate_image(gt, pred))
# 这里重点:evaluator这个是什么,它是从./detection/iou中
# import 的一个 DetectionIoUEvaluator类
# 最终返回results列表给trainer.py
return results

# 该方法返回的是一个batch的评估结果,以字典的形式
def validate_measure(self, batch, output, box_thresh=0.6):
return self.measure(batch, output, box_thresh)

def evaluate_measure(self, batch, output):
return self.measure(batch, output), np.linspace(0, batch['image'].shape[0]).tolist()

# 这个方法则是返回所有batch组合(也就是一个数据集)的评估结果
# raw_metrices是从trainer.py文件中传过来的,
# 它里面存的是每一个batch所得到的一个评估结果
# 也就是for循环调用validate_measure之后得到的一个评估合集
# 现在将这个合集传入gather_measure方法中,
# 通过这个方法,将这个合集中每一个batch的评估结果进行一个整合
# 这个方法返回给trainer.py的是整个数据集的一个评估结果
def gather_measure(self, raw_metrics):
# 展平,raw_metrices中的每个元素都是一个batch的评估结果
raw_metrics = [image_metrics
for batch_metrics in raw_metrics
for image_metrics in batch_metrics]

# 将raw_metrices传入iou.py文件中的combine_results方法
result = self.evaluator.combine_results(raw_metrics)
#=================================================================================

# 将整个数据集的评估结果返回个trainer.py文件中的metrics
return {
'precision': precision,
'recall': recall,
'fmeasure': fmeasure
}
utils/ocr_metric/icdar2015/detection/iou.py
1
2
3
4
    
# gt是从quad_metric.py中一张图片的多个正确文本框标记
# pred也是从quad_metric.py中传入的,是一个存有一张图片中多个probability map的列表
def evaluate_image(self, gt, pred): # 核心评估函数,评估的是一张图中的若干文字框
1
2
3
4
5
6
7
8
9
10
def get_union(pD, pG): # 求并集
return Polygon(pD).union(Polygon(pG)).area

# iou 计算的是 “预测的边框” 和 “真实的边框” 的交集和并集的比值
# 即iou = get_intersection(pD, pG) / get_union(pD,pG)
def get_intersection_over_union(pD, pG): # 求iou
return get_intersection(pD, pG) / get_union(pD, pG)

def get_intersection(pD, pG): # 求交集
return Polygon(pD).intersection(Polygon(pG)).area

Precision精确率, Recall召回率,是二分类问题常用的评价指标。

1
2
3
4
5
6
7
8
# Num det = n
# Num pred = m
# care det = n'
# care pred = m'
# matched = k #表示匹配的个数
# recall = k / n'
# precision = k / m'
# 计算recall和precision

预测部分代码

tools/predict.py

1
2
3
4
5
6
7
8
9
10
# -----------------------------------------------------------------------
#配置重点:get_post_processing 是获得预测的文本框,
# 其中根据已配置的文件config中确定的post_processing,
# 我们选择的post_processing是post_processing文件夹下的
# seg_detector_representer.py文件
self.post_process = get_post_processing(config['post_processing'])
self.post_process.box_thresh = post_p_thre
# 目前只是确定了(配置好)后处理获取文本框所需要用到的方法是
# SegDetectorRepresenter,在下面的预测中才实际使用它
# -----------------------------------------------------------------------
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
#--------------------------------------------------------------------------
# preds:从网络模型的head中输出的一个二维的包含probability map和threshold map两个张量的张量
# preds是models文件下的model.py中的forward函数中return的y
# preds的第0个维度为probability map,第二个维度为threshold map
preds = self.model(tensor)
if str(self.device).__contains__('cuda'):
# 等待CUDA上所有流中的所有核心上的操作都完成,使得测试出来的时间更加准确
torch.cuda.synchronize(self.device)
#-------------------------------------------------------------------------------
#-------------------------------------------------------------------------------
#获取文本框,实际重点:
# 后处理后得到最终的文本框box_list
# 将batch和preds传到post_processing文件夹下的seg_detector_representer.py文件夹下的
# def __call__():
box_list, score_list = self.post_process(batch, preds, is_output_polygon=is_output_polygon)
# 因为预测的图片只有一张,所以只需要获取第一个即可
box_list, score_list = box_list[0], score_list[0]
#--------------------------------------------------------------------------------

重点:进行后处理操作

训练或者预测中====

用来获取训练train.py时的train数据集中的文本框,

预测predict.py待预测图片的文本框的操作

post_processing/seg_detector_representer.py

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# thresh=0.3将用于一个传统的二值化操作,box_thresh筛掉一些无用的文字区域,
# max_candidates是指一张图片的文字区域最多为1000
# unclip_ratio(膨胀比率):因为probability map是shrink后的图片,
# 需要通过unclip_ratio经过一个公式D‘计算后将probability map经过dilated(膨胀)
# 来得到一个和文字区域相同的map,unrlip_ratio就是该公式D’的r',设置为1.5
def __init__(self, thresh=0.3, box_thresh=0.7, max_candidates=1000, unclip_ratio=1.5):

#====================================================================================================
# 这里的的batch和pred是从tools文件夹下的predict.py
# 或者trainer文件夹下的trainer.py中传过来的batch和preds

# 最终要将boxes, scores返回给predict.py的两个变量box_list, score_list
# 或者返回给trainer.py中的boxes和scores
# 即 return boxes, scores
def __call__(self, batch, pred, is_output_polygon=False):

#=====================================================================================================
# 将box存入boxes中,并传入
boxes.append(box)
# 将函数box_score_fast()计算出来的一张图片的一个文本框的概率均值存入scores
scores.append(score)
# 将两者分别传回给tools文件夹中的predict.py中的box_list, score_list
# 并在predict.py中
# 如果是训练阶段,则似乎将boxes,scores传回给trainer.py中的boxes,scores
return boxes, scores

label标签制作

根据人工标记的gt框(一系列坐标点),进行一些膨胀(dilate)和缩小(shrink)的操作

膨胀和缩小的距离:

并且做一些gt框内的计算来得到probability_mapthreshold_map

data_loader/modules/make_shrink_map.py

一个shrink操作:

红框为gt框,蓝框为shrink后的结果

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# 所采用的shrink方法         多边形的一系列点, 需要shrink的比率,按照ppt上,所采用的比率为0.4
def shrink_polygon_pyclipper(polygon, shrink_ratio):
from shapely.geometry import Polygon
import pyclipper # 该库是创建一个多边形来对原图进行修改,既可以缩小,也可以膨胀
polygon_shape = Polygon(polygon) # 创建多边形
# distance的公式是ppt20页的,用来计算具体的一个缩小范围
distance = polygon_shape.area * (1 - np.power(shrink_ratio, 2)) / polygon_shape.length

subject = [tuple(l) for l in polygon]
padding = pyclipper.PyclipperOffset()
padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)

# 因为distance是一个整数,我们需要的是shrink而不是膨胀,所以输入-distance
shrinked = padding.Execute(-distance)
if shrinked == []:
shrinked = np.array(shrinked)
else:
shrinked = np.array(shrinked[0]).reshape(-1, 2)
return shrinked # 返回ppt20页中的构成红框内的“蓝框”的所有点的坐标,蓝框是将红框shrink后得到的

data_loader/modules/make_border_map.py

shrink操作 + dilate操作:

红框为gt框,蓝框为shrink框,绿框为dilate框

1
2
3
4
5
6
7
8
9
# 和shrink_map中的公式一样
distance = polygon_shape.area * (1 - np.power(self.shrink_ratio, 2)) / polygon_shape.length
subject = [tuple(l) for l in polygon]
padding = pyclipper.PyclipperOffset()
padding.AddPath(subject, pyclipper.JT_ROUND,
pyclipper.ET_CLOSEDPOLYGON)

# 传入正的distance,即膨胀操作,膨胀后得到的是ppt40页中的绿色框
padded_polygon = np.array(padding.Execute(distance)[0])

处理后的Distance map

1
2
3
4
5
6
7
8
9
10
11
12
13
# 这样一来,在绿框中,绿框到蓝框边的所有区域的点在distance_map中取值为0
# 蓝框内的所有点在distance_map中取值为1
distance_map[i] = np.clip(absolute_distance / distance, 0, 1)

# 1 - distance_map 其实算是一个将distance_map内的每个点的值的取反的操作。
# 这里的canvas可以看作gt图,下面的操作,其实就是根据上面求出点的distance_map图中的值
# 对canvas进行一个填充,填充的结果图在ppt46页左下角的gt图,
# gt图中很大区域是灰色的原因是因为前面初始化gt图时,将里面的点的值设为了0-1之间,所以才为灰色
canvas[ymin_valid:ymax_valid + 1, xmin_valid:xmax_valid + 1] = np.fmax(
1 - distance_map[
ymin_valid - ymin:ymax_valid - ymax + height,
xmin_valid - xmin:xmax_valid - xmax + width],
canvas[ymin_valid:ymax_valid + 1, xmin_valid:xmax_valid + 1])

两个map的计算

shrink_map 和 dilate_map的一个计算


本博客所有文章除特别声明外,转载请注明出处!