算法小用:我有一张图,想知道哪个PDF里出现过

777 阅读14分钟

算法不远,AI就在身边。下面是算法服务生活时间,我是主持人TF男孩。

信息时代,我们不能抵触信息,否则会落后于时代,我们需要做好信息筛选。

我加了好多群,群里好多PDF格式的分析报告,文件自动下载,堆积如山,因此只能选着看。

为此,我写了一个关键词筛选的功能。只需要输入“人工智能”,程序就可以从几百个文件中,找到哪个文件的第几页提到了“人工智能”。

这样,我就能大概知道哪个领域关注人工智能了。

一、文字搜索

听着很牛,看着不错,其实代码不到20行。

import fitz  # PyMuPDF
import os

def search_word_file(file_path, key_words):
    _, file_name = os.path.split(file_path)
    doc = fitz.open(file_path)
    page_size = len(doc)
    for page_num in range(page_size):
        page = doc[page_num]
        page_text = page.get_text()
        if key_words in page_text:
            print(f"文档 《{file_name}》 第 [{page_num}] 页包含关键字 '{key_words}'")

def search_word_dir(dataset_dir, key_words):
    for file_name in os.listdir(dataset_dir):
        file_path = os.path.join(dataset_dir,file_name)
        print(f"开始搜索文件 《{file_name}》...")
        search_word_file(file_path, key_words)

search_word_dir("资料库", "人工智能")

这里面用到了一个PDF处理库PyMuPDF。需要pip安装一下:

pip install PyMuPDF

利用fitz.open()打开PDF文件,可获得每页数据,page.get_text()可得此页文字。随后,搜索关键字是否在里面。最后,文件夹里的PDF都执行下这个操作,就完成了。

后来,又有一个新需求。我整理资料时,偶尔会看到一些新奇的图片。图片多是内部资料。这些图片,有人用过吗?谁用的?出现在哪个文档里?用在什么场景下?

图片,也可以像文字一样搜索吧?

二、图片读取

事实上是可以的。我做出来了。

而且搜索的结果也准确。

PyMuPDF不仅能读取PDF里的文本内容,也可以获取其中的图片元素。

page.get_text()是返回这一页的文本。那么,page.get_images(full=True)则是返回此页的所有图片。图片,可以有多张。

因此,获取PDF文件里的图片就变得简单了。

import fitz 

doc = fitz.open("test.pdf")
page_size = len(doc)
# 遍历每一页提取图片
for page_num in range(page_size):
    page = doc[page_num]
    images = page.get_images(full=True)

    for img_index, img in enumerate(images):
        xref = img[0]  # 图片的引用 ID
        base_image = doc.extract_image(xref)
        image_bytes = base_image["image"]
        image_ext = base_image["ext"]  # 图片格式(如 'png', 'jpeg')
        if image_ext not in ["jpg", "jpeg", "png"]: continue
        image_name = f"page{page_num+1}_img{img_index+1}.{image_ext}"
        with open(image_name, "wb") as image_file:
            image_file.write(image_bytes)

以上代码,用PyMuPDF获取PDF文件每一页里内嵌的图片,然后筛选jpg和png格式,保存为图片文件。命名格式为:page页码_img第几张图。

我们找一个PDF文档,试试看。

好像是拿到了整个文档中每一页的图片。下面该想办法做图像搜索了。

我手里有一张图,这算是个模板。旁边有一堆图,这堆图里有没有出现过这个模板呢?

三、特征匹配算法

这个需求,有一些低成本的解决方案。比如特征匹配,它就是将一个模板图片与一堆待搜索的图片进行比较,找出相似度最高的那一张图片。

OpenCV提供了多种特征匹配的算法,如SIFT、SURF、ORB等。他们各有优劣,今天我们不做测评,指定ORB算法。

3.1 ORB算法

ORB,全称是 Oriented FAST and Rotated BRIEF,是一种基于FAST关键点检测和BRIEF特征描述的算法。它属于中不溜的那种算法,效果还行,速度也不慢。虽然鲁棒性不如SIFT,但应对这个需求也是绰绰有余。

3.1.1 鲁棒性

鲁棒性是什么?老听你们说这个词儿。英文是Robustness,RoBu直接音译叫“鲁棒”,意思是稳健性、耐用性。其实就是抗造,就好像有山东人拿根大棒子过来,你也一样泰然自若。这就是有很强的鲁棒性。

在计算机视觉处理领域,一张图会有很多不确定性。比如光照,白天拍的照亮一些,晚上拍的暗一些。一个鲁棒的算法,应当能提取出同一样的特征。再比如形变,包括大小缩放,角度旋转,视角变化等,不能说一张苹果照片倒过来就不是苹果,那样不鲁棒。另外,还有一些其他干扰,比如图像的噪点,模糊,遮挡等等。

举个简单的例子吧,匹配里面有个初级的算法叫“模板匹配”,它一点也不鲁棒。

3.1.2 反面例子:模板匹配

拿一张图举例子。这是我十一假期在青岛拍的照片(远处的小岛就是父母爱情的取景地),这张图叫image.png

咱们将右边带驾驶室的那艘渔船,从图上裁剪出来,作为模板template.png

咱们让算法试试能不能在图片里找到它。

import cv2
# 加载图像和模板
image = cv2.imread('image.png')
template = cv2.imread('template.png')
# 转换为灰度图
gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
gray_template = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)
# 使用模板匹配
result = cv2.matchTemplate(gray_image, gray_template, cv2.TM_CCOEFF_NORMED)
# 获取最匹配位置
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
# 标记匹配区域
top_left = max_loc
h, w = gray_template.shape
bottom_right = (top_left[0] + w, top_left[1] + h)
# 绘制矩形框
cv2.rectangle(image, top_left, bottom_right, (0, 255, 0), 2)
cv2.imwrite('output.png', image)

运行前,先安装cv2。

pip install opencv-python

结果,在图片里找到了。

随后,我们把image.png缩小一半,再试试。结果就不行了。

咱们尺寸不变,把原图到过来,也搜不到。

这就是鲁棒性不行。具体点就是不具备旋转不变性,尺度不变性。

这些问题对于ORB算法而言,影响很小。

3.2 特征点和描述

传统算法,不走神经网络,一般都是根据规则计算和比对特征点。

ORB算法改良了FAST算法的角点检测,这是原本用于检测图像中显著的特征点。人类的身体构造就有特征点,比如头部从肩膀处突了出来,腿从腰部分叉了。图像也一样,通过某个像素与其周围16个像素的亮度值关系,会判断是否为角点。ORB算法从一张图片上会提取500个特征点,而且额外加入了方向的信息,让它做到方向不变性。

特征点提取出来,只是一个区域,下一步就是特征描述:这个区域有什么特点,它和别的点有什么差异。类比到人体,就是某个区域的重量、长度、肤色等等。对于描述,ORB用的是BRIEF算法,全称Binary Robust Independent Elementary Features,中文叫二元稳健独立基本特征。对于ORB提取的500个关键点,它分别再以32位的数值进行描述。

既然每张图都提取出500个特征点,每个特征点都有32位的描述。那就可以进行比对了。有了数据其实很好比对。类比到人的信息,我们已经把每个人的详细信息都提取出来了,要判断两个人像不像,不难。甚至孩子像妈妈,即便身高、体重不一样,通过面部特征也能判断出来他们好像同一个人。

3.3 代码解读

上面模板匹配的例子,用ORB算法改写一下,就是这样。

import cv2
# 加载图像和模板
image = cv2.imread('image.png')
template = cv2.imread('template.png')
# 转为灰度图
gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
gray_template = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)
# 初始化ORB检测器
orb = cv2.ORB_create()
# 检测关键点和描述符
kp_image, des_image = orb.detectAndCompute(gray_image, None)
kp_template, des_template = orb.detectAndCompute(gray_template, None)
# 使用Brute Force Matcher进行匹配
bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
matches = bf.match(des_image, des_template)
# 按匹配度排序
matches = sorted(matches, key = lambda x:x.distance)
# 在图像上绘制匹配结果
matched_image = cv2.drawMatches(image, kp_image, template, kp_template, matches[:30], None, flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
cv2.imwrite('output.png', matched_image)

先看效果,从一张图上找到了另一张图的关键点。

而且,根据关键点,可以计算出位于哪个区域。

首先,orb.detectAndCompute是提取关键点和描述符,它只接收灰度图片,彩色的它无法处理。第二个参数是None,表示没有指定预置的特征点,让系统自己寻找。

返回值有两个,kp_image是关键点,des_image是描述符。

关键点的结构,你可以这样打印,调用和结果如下:

# 遍历关键点并读取信息
for i, kp in enumerate(kp_image):
    print(f"KeyPoint {i + 1}:")
    print(f"  坐标 (x, y): {kp.pt}")
    print(f"  尺寸 (size): {kp.size}")
    print(f"  方向 (angle): {kp.angle}")

KeyPoint 1:
  坐标 (x, y): (246.0, 227.0)
  尺寸 (size): 31.0
  方向 (angle): 307.1259460449219
KeyPoint 2:
  坐标 (x, y): (687.0, 153.0)
  尺寸 (size): 31.0
  方向 (angle): 55.012447357177734
  ……
KeyPoint 500:
  坐标 (x, y): (569.7258911132812, 200.65818786621094)
  尺寸 (size): 111.0786361694336
  方向 (angle): 132.6415557861328

des_image是这500个关键点的描述:

# 500条描述数据:
array([[107, 181, 177, ..., 206, 219, 223],
       [ 24,  94,  23, ...,   0, 165,  57],
       ...,
       [ 43, 206, 108, ..., 142, 248,  76]], dtype=uint8)

# 取其中一条数据,是32个数字描述:
array([107, 181, 177, 230, 237,  29, 220, 255, 203, 107, 189, 104, 254,
       253, 220, 166, 159,  91, 255,  13, 215, 107, 250, 211, 191,  95,
        63, 254, 206, 206, 219, 223], dtype=uint8)

特征点找到了,描述符找到了,下一步就是匹配。

先用cv2.BFMatcher构建一个匹配器bf。然后用bf.match对两个文件的描述进行比对。

cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)

在构建匹配器时,第一个参数是cv2.NORM_HAMMING,表示使用汉明距离来计算差异。计算差异有很多种方法,ORB算法用的是二进制描述,所以用汉明距离。SIFT算法用的是浮动描述,它用cv2.NORM_L2L2距离(欧几里得距离)。

crossCheck=True表示进行交叉验证匹配。白话就是我身上的鼻子和你最像,你身上也是鼻子和我最像。不然容易误匹配,比如你渔船上的钉子和我皮包上的纽扣最像,这不是一个物体。

匹配完之后,按照距离排序。distance表示两者的距离,越小越好,为0表示完全一样。排完序,取前30个匹配结果。用cv2.drawMatches在图像上绘制匹配结果。

算法往往都是匹配数据,不会出结论。结论是我们自己出的。

下面这段代码就是一种出结论方式。

# 设置距离阈值,保留高质量匹配
good_matches = [m for m in matches if m.distance < 30]
# 判断匹配是否足够好
if len(good_matches) > 60:  # 根据实际需求设置数量阈值
    print("找到较多匹配项,可能是相关图像")
else:
    print("匹配点过少,可能不是同一对象")

首先从匹配点matches里面选取距离差异小于30的点,作为最佳匹配点。然后,如果匹配点数大于60,就认为找到较多匹配项,是相关图像。否则,两个图像不相关。

为什么要做这个操作?因为就算两个不相关的图片,也可能有相似点。

但是,有相似点不一定能过验证这一关。验证这一关的把控,就是我们调试出来的。

下面的两艘船,它们相似特征点挺多的。只所以选中右边的那个,是因为选取了差异更小的、匹配更多,而且加了交叉验证的一方。

如果是完全相同的两个图,他们的差异是0。

我们再验证几个图像,看看效果。

让他缩放。

让他旋转。

好了,算法方面解决了,下面融入到项目中。

四、项目实战

首先,针对PDF文件的其中一页,提取出所有内嵌的图片。

def get_page_imgs(doc, page_num):
    image_arr = []
    page = doc[page_num]
    images = page.get_images(full=True)
    for img in images:
        xref = img[0]  # 图片的引用 ID
        base_image = doc.extract_image(xref)
        ……
        image_array = np.frombuffer(image_bytes, dtype=np.uint8)
        image = cv2.imdecode(image_array, cv2.IMREAD_GRAYSCALE)
        image_arr.append(image)
        ……
    return image_arr

此处图片文件进行灰度化后,留在内存,放到image_arr数组里。

然后,写一个针对两张灰度图片做对比的函数。

def is_same_img(gray_image, gray_template, MIN_DIS=30, MIN_MATCHE=30):
    orb = cv2.ORB_create()
    kp_image, des_image = orb.detectAndCompute(gray_image, None)
    kp_template, des_template = orb.detectAndCompute(gray_template, None)
    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
    matches = bf.match(des_image, des_template)
    matches = sorted(matches, key = lambda x:x.distance)
    good_matches = [m for m in matches if m.distance < MIN_DIS]
    return len(good_matches) > MIN_MATCHE

代码很简单。就是先用上面讲的detectAndCompute检测关键点,外加bf.match匹配。收录了差异距离小于MIN_DIS的点作为最佳匹配点。同时,如果最佳匹配点数量大于MIN_MATCHE,就认为是相关图像。

然后,针对一个PDF文件,让每一页都走下get_page_imgs获得此页的所有图片,并调用is_same_img同目标模板进行对照。

def search_img_file(file_path, gray_template):
    _, file_name = os.path.split(file_path)
    doc = fitz.open(file_path)
    page_size = len(doc)
    for page_num in range(page_size): # 查询一页
        image_arr = get_page_imgs(doc, page_num)
        for img in image_arr:
            if is_same_img(img, gray_template):
                print(f"文档 《{file_name}》 第 [{page_num+1}] 页找到这个图像")
                continue

目标模板gray_template是提前加载的,作为参数传入。这样就实现了要寻找的图,同一个PDF里所有的图片对照一遍。如果找到匹配的,就输出文档名和页码。

最最后,让一个文件夹里每个PDF都按照这个方式走一遍,就实现了批量文件的图片搜索。

def search_img_dir(dataset_dir, gray_path):
    gray_template = cv2.imread(gray_path, cv2.IMREAD_GRAYSCALE)
    for file_name in os.listdir(dataset_dir):
        if not file_name.endswith(".pdf"): continue
        file_path = os.path.join(dataset_dir,file_name)
        print(f"开始搜索文件 《{file_name}》...")
        search_img_file(file_path, gray_template)

search_img_dir("资料库", "temp.png")

这样就实现我们要的效果。

ORB算法的效果,还不错,就算是手机拍摄屏幕,也一样能识别。

五、总结

此类特征匹配算法非常之多,比如SIFT、SURF、BRIEF、ORB、AKAZE、BRISK、FREAK、Harris。他们各有优劣。

特性SIFTSURFBRIEFORBAKAZEBRISKFREAKHarris
鲁棒性高(尺度、旋转、光照、噪声)较高(尺度、旋转)低(不适应尺度、旋转、光照变化)较好(旋转不变,尺度适应较差)高(尺度、旋转)高(尺度、旋转)较低(旋转不变性强)低(对光照变化敏感)
计算效率低(计算复杂)较高高(计算和存储都快)高(实时应用)中(比SIFT快)高(实时应用)高(实时应用)高(简单快速)
描述符128维浮动描述符64维描述符二进制描述符二进制描述符二进制描述符二进制描述符二进制描述符无(只检测角点)
适应性高(尺度、旋转、仿射)较强(旋转、尺度)低(不变性差)较好(旋转不变)较好(尺度、旋转)较强(尺度、旋转)较强(旋转不变)低(不适应尺度、旋转)
应用场景高精度应用,如拼接、三维重建高效应用,如拼接、目标识别实时应用,如目标跟踪实时应用,如SLAM、AR实时应用,如SLAM、拼接实时应用,如目标检测实时应用,如SLAM角点检测,运动估计

其实在OpenCV中,算法替换很简单。比如将我们的ORB算法替换为SIFT算法,这样修改就行。

sift = cv2.SIFT_create()
kp_image, des_image = sift.detectAndCompute(gray_image, None)
kp_template, des_template = sift.detectAndCompute(gray_template, None)
bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)

创建检测器不一样,构造匹配器不一样。其他都差不多。你可以试一试效果。

具备高精度的算法,计算量大,比较慢。计算快的算法,精度又低。又快又准的,鲁棒性有短板。这都需要依据实际场景来选择。甚至最开始的模板匹配那么差的鲁棒性,依然也有适用空间。高端的技术投入也多,高端的生活成本也高。

我从来不过度推崇高端技术、新技术,因为老板们往往都不缺技术,缺钱。以图搜图效果最好是引入深度学习,但是那样不但对硬件有要求,更要命的还要有数据。

因此啊,两块钱一打的一次性手套,能解决某个特定场景的问题,是可行的,而且效果也不错。

学AI,解决自己生活中的问题。凑合用呗,你还想卖钱咋地。卖钱的路可长了。

好了,这期的AI生活有妙招就到这里。咱们下棋见……下期见。

我是爱生活、爱技术的TF男孩。