一种基于视觉词汇的文本分类方法

2017-09-18 本文已影响2197人老曹撸代码

一年多以前我脑子一热，想做一款移动应用：一款给学生朋友用的“错题集”应用，可以将错题拍照，记录图像的同时，还能自动分类。比如拍个题目，应用会把它自动分类为"物理/力学/曲线运动"。当然，这个项目其实不靠谱，市场上已经有太多“搜题”类应用了。但过程很有趣，导致我过了一年多，清理磁盘垃圾时，还舍不得删掉这个项目的“成果”，所以干脆回收利用一下，写篇文章圈圈粉（源码在这里, 欢迎 Star）。

这个项目，核心要解决的问题就是文本分类。所以最初想到的方案是先 OCR 图片转文本，然后分词，再计算 tf-idf，最后用 SVM 分类。但这个方案的问题是：开源 OCR 普遍需要自己训练，且需要做大量的优化、调校和训练，才能在中文识别上有不错的效果，加上图像上还会有公式、几何图形，这些特征也会决定分类，这又提高了对 OCR 的要求。所以我最终选择的方案是，不使用 OCR，而是直接从图像中寻找有区分性的、鲁棒的特征，作为视觉词汇。之后再通过传统文本分类的方法，训练分类器。

下面将展示整个训练过程，训练的样本来自《2016 B版 5年高考3年模拟:高考理数》(样本数据可在此下载)，并手工标注了14个分类，每个分类下约50个样本，每个样本为一个题目, 图像为手机拍摄。

本文中大部分算法库来自numpy、scipy、opencv、skimage、sklearn。

1. 预处理

为了获取到稳定的特征，我们需要对图像进行预处理，包括调整图像大小，将图像缩放到合适尺寸；旋转图像，或者说调整成水平；二值化，去除色彩信息，产生黑白图像。

1.1. 调整图像大小

调整的目的是为了让图像中文字的尺寸保持大致相同的像素尺寸。这里做了一个简单假设，即：图像基本是一段完整的文本，比如一个段落，或者一页文档，那么不同的图像中，每行文本的字数相差不会很大。这样我就可以从我所了解的、少得可怜的图像工具库里找到一个工具了：直线拟合。即通过拟合的直线(线段)长度与图像宽度的比例，调整图像的大小。下图为两张不同尺寸图像，经过多次拟合+调整大小后的结果，其中红色算法检查到的直线(线段)。

下面是使用 opencv 直线拟合的代码：

# Canny算法提取边缘特征, image是256灰度图像
image = cv2.Canny(image, 50, 200)
# 霍夫线变换提取直线
lines = cv2.HoughLinesP(image, 2, math.pi / 180.0, 40, numpy.array([]), 50, 10)[0]

1.2. 图像二值化

二值算法选用skimage.filters.threshold_adaptive（局部自适应阀值的二值化）, 试下来针对这种场景，这个算法效果最好，其他算法可以去scikit-image文档了解。下图为全局阀值和局部自适应阀值的效果对比：

1.3. 旋转图像

从第一步获取到的直线，可以计算出图像的倾斜角度，针对只是轻微倾斜的图像，可以反向旋转进行调整。由于可能存在干扰线条，所以这里取所有直线倾斜角度的中值比平均值更合适。下图展示了图像旋转跳转前后的效果：

2. 提取特征

这里的思路是，首先通过形态学处理，可以分割出文本行（的图像），再从文本行中分割出词汇（的图像），然后从"词汇"中提取特征。但这里的需要克服的困难是：

很多汉字分左右部，容易被错分，比如你好, 可能被分割成以4块图像：亻、尔、女、子。
独立的“字”并不适合于文本分类，还需能学习出词汇。

针对以上问题的解决方案是：

将小的图像块进行组合，组合后的新图像块和原来的小块图像一起作为原始图像的特征，如你好将得到10个特征：亻、你、你女，你好，尔、尔女、尔好、女、好、子。
得益于上面的方案，词汇信息也被保留了下来，所以第二个问题也就解决了，同时增加了算法的鲁棒性。

下面将介绍具体实现。

2.1. 提取文本行

由于预处理过程中已经将样本的图像尺寸基本调整一致，所以可以比较容易的利用形态学的处理方法，分割出文本行。过程如下：

# cv2.Canny 可提取边缘，并去除噪点
# image为调整过大小，但没有调整水平和二值化的图像
# 二值化后会影响 cv2.Canny 算法效果，所以这里用还没有二值化的图片
image = cv2.Canny(image, 100, 200)
# 二值化后调整水平
image = ndimage.rotate(image, slope)

# 进行四次膨胀和腐蚀操作
# 水平方向膨胀和腐蚀，联通字与字之间的空间
# 垂直方向做较小的膨胀和腐蚀，填补行内的空隙
image = cv2.dilate(image, cv2.getStructuringElement(cv2.MORPH_RECT, (40, 3)))
image = cv2.erode(image, cv2.getStructuringElement(cv2.MORPH_RECT, (40, 3)))
image = cv2.erode(image, cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5)))
image = cv2.dilate(image, cv2.getStructuringElement(cv2.MORPH_RECT, (6, 5)))

下图展示了每一步的变化：

接下来可以利用scipy库中的measurements.label方法，标记出不同的的区域，下图展示了标注后的效果，不同区域以不同的灰度表示。

2.2. 提取特征(视觉词汇)

裁剪出单行文本图像后，我们可以将图像中各列的像素的值各自累加，得到一个一纬数组，此数组中的每个局部最小值所在的位置，即为文字间的空隙。如下图所示，其中蓝色线为像素值的累加值，绿色线为其通过高斯滤波平滑后的效果，红色线为最终检测到的分割点。

详细过程见下面代码：

# 1. 将图像中每一列的所有像素的值累加
orisum = image.sum(axis=0) / 255.0

# 2. 累加后的数组通过高斯滤波器做平滑处理，减少干扰
filtered = filters.gaussian_filter(orisum, 8)

# 3. 找出拐点(上升转下降、下降转上升的点)
trend = False  # False 下降，True上升
preval = 0  # 上一个值
points = []  # 拐点
pos = 0
for i in filtered:
    if preval != i:
        if trend != (i > preval):
            trend = (i > preval)
            points += [[pos if pos == 0 else pos - 1, preval, orisum[pos]]]
    pos = pos + 1
    preval = i
 
# 4. 下降转上升的拐点即为分割点
  ... 代码略 ...

将单行的图像按上述方法获取的分割点进行裁剪，裁剪出单个字符，然后再把相邻的单个字符进行组合，得到最终的特征数据。组合相邻字符是为了使特征中保留词汇信息，同时增加鲁棒性。下图为最终获得的特征信息：

本文中使用的所有样本，最终能提取出约30万个特征。

2.3. 选择特征描述子

选择合适的特征描述子通常需要直觉+运气+不停的尝试（好吧我承认这里没有什么经验可分享），经过几次尝试，最终选中了HOG（方向梯度直方图）描述子。HOG 最让人熟悉的应用领域应该是行人检测了，它很适合描述钢性物体的边缘特征（方向），而印刷字体首先是刚性的，其次其关键信息都包含在边缘的方向上，所以理论上也适合用 HOG 描述。更多关于HOG的介绍请点击这里。下图为文字图像及其 HOG 描述子的可视化：

代码如下：

# 提取边缘
canny = cv2.Canny(numpy.uint8(img), 50, 200)
# 计算描特征描述子
desc, hog_image = hog(
    canny, 
    orientations=6, 
    pixels_per_cell=(4, 4), 
    cells_per_block=(2, 2), 
    visualise=True)

3. 训练词汇分类器

对词汇进行人工标注工作量太大，所以最好能做到自动分类。我的做法是先聚类，再基于聚类的结果训练分类器。但有个问题，主流的聚类算法中，除了 K-Means 外，其他都不适合处理大量样本（目前有30万+样本），但 K-Means 在这个场景上聚类效果不佳，高频但不相关的词汇容易被聚成一类，而 DBSCAN 效果很好，但样本数一多，所需时间几何级增长（在我的机器上，超过两万个样本就需要耗费数个小时）。下图来自sklearn 文档，对各聚类算法做了比较：

2017/09/21 修改：原此处选择的聚类方法（即先使用先用 K-Means 做较少的分类然后对每个分类单独使用 DBSCAN 聚类并单独训练 SVC 分类器），准确率保持在70%左右，很难提高，故改用了下面描述的新方法。

为解决这一问题，我的做法是：
1. 先对每类样本下的词汇用 DBSCAN 聚类（约1万个词汇样本），得到一级分类。
2. 聚类后，计算每个一级分类的中心，然后以所有中心为样本再用DBSCAN聚类，得到二级分类。完成后，原一级分类中心的新分类，即代表其原一级分类下所有元素的分类。

聚类的过程为，使用前面提取的 HOG 特征，先 PCA 降纬，再 DBSCAN 聚类。这里注意，计算二级分类时，PCA应使用全局样本计算。

分类器使用SGDClassifier，原因是其支持分批计算，不至于导致内存不足。

本文中使用的样本，最终得到3000+词汇类型。下图为分类效果，其中每一行为一个分类：

4. 训练文本分类器

有了词汇分类器，我们终于可以识别出每个文本样本上所包含的词汇了（事实上前面步骤的中间过程也能得到每个样本的词汇信息），于是我们可以给每个样本计算一个词袋模型（即用每个词出现的次数表示一篇文本），再通过池袋模型计算TF-IDF模型（即用每个词的 TF*IDF 值表示一篇文本），并最终训练 SVM 分类器。下面展示了此过程的主要代码：

# bow 保存每篇文本的词袋模型
# labs 保存每篇文本类型
    
# 计算 tfidf 模型
tfidf = TfidfTransformer(smooth_idf=True, sublinear_tf=True, use_idf=True)
datas = numpy.array(tfidf.fit_transform(bow).toarray())
    
# 使用 PCA 降纬
pca = RandomizedPCA(n_components=20, whiten=True).fit(datas)
    
# 样本数据20%用于测试，80%用于训练
X_train, X_test, y_train, y_test = train_test_split(data, labs, test_size=0.2, random_state=42)

# 训练SVM  
print("Fitting the classifier to the training set")

param_grid = {'C': [1e3, 5e3, 1e4, 5e4, 1e5],
          'gamma': [0.0001, 0.0005, 0.001, 0.005, 0.01, 0.1], }
gscv = GridSearchCV(SVC(kernel='linear', probability=True), param_grid=param_grid)

svc = gscv.fit(X_train, y_train)
print("Best estimator found by grid search:")
print(svc.best_estimator_)
print "score :", svc.score(data, labs)
  
# 测试训练结果
print("Predicting on the test set")

y_pred = svc.predict(X_test)

print(classification_report(y_test, y_pred, target_names=target_names))
print(confusion_matrix(y_test, y_pred))

执行结果如下:

Fitting the classifier to the training set
done in 0.034s
score : 0.918639053254
Predicting on the test set
done in 0.004s
               precision    recall  f1-score   support

    物理-电学-静电场       1.00      0.67      0.80         3
   物理-力学-互相作用       0.56      1.00      0.71         5
  物理-机械振动和机械波       0.83      1.00      0.91         5
   物理-电学-电磁感应       0.71      1.00      0.83         5
   物理-电学-恒定电流       1.00      1.00      1.00         5
   物理-力学-曲线运动       0.88      0.78      0.82         9
 物理-机械能及其守恒定律       0.62      0.56      0.59         9
        物理-光学       1.00      0.50      0.67         2
物理-力学-万有引力与航天       1.00      0.75      0.86         4
 物理-力学-牛顿运动定律       0.62      0.71      0.67         7
   物理-电学-交变电流       1.00      1.00      1.00         1
     物理-电学-磁场       1.00      0.25      0.40         4
        物理-热学       1.00      1.00      1.00         2
物理-力学-质点的直线运动       0.86      0.86      0.86         7

  avg / total       0.81      0.78      0.77        68

[[2 0 0 0 0 0 1 0 0 0 0 0 0 0]
 [0 5 0 0 0 0 0 0 0 0 0 0 0 0]
 [0 0 5 0 0 0 0 0 0 0 0 0 0 0]
 [0 0 0 5 0 0 0 0 0 0 0 0 0 0]
 [0 0 0 0 5 0 0 0 0 0 0 0 0 0]
 [0 1 0 0 0 7 0 0 0 0 0 0 0 1]
 [0 2 0 0 0 1 5 0 0 1 0 0 0 0]
 [0 0 0 0 0 0 0 1 0 1 0 0 0 0]
 [0 0 0 0 0 0 1 0 3 0 0 0 0 0]
 [0 1 0 0 0 0 1 0 0 5 0 0 0 0]
 [0 0 0 0 0 0 0 0 0 0 1 0 0 0]
 [0 0 0 2 0 0 0 0 0 1 0 1 0 0]
 [0 0 0 0 0 0 0 0 0 0 0 0 2 0]
 [0 0 1 0 0 0 0 0 0 0 0 0 0 6]]

测试集上正确率 81%，召回率 78%。个别分类正确率较低，可能是因为样本数太少，另外训练过程大多使用默认参数，若进行细致调校，应该还有提高空间。

5. 结束

此项目完整代码及样本数据可点击这里下载。任何想在实际项目中使用此方法的朋友请注意，以上方法目前只在一个样本库中测试过，在其他样本库中表现如何还不知道，但愿没把你带坑里。