本文为稀土技能社区首发签约文章，30天内制止转载，30天后未获授权制止转载，侵权必究！

一、前言

上一篇介绍以图搜图的完结：/post/725585… ，咱们利用了卷积神经网络提取特征，然后比照特征类似度，并运用向量数据库加快查找。本文咱们将介绍依据文本查找图片的完结。

首要需求知道依据文本查找图片详细是什么问题，这儿能够有两个层面。第一个则是图片中包含的文本内容，这个能够用OCR辨认提取出来。第二个则是深层次的对图片描绘的文本，比方红色的狗、跑步的猪、骑猪的人。这些都是对图片内容的描绘，相比之下第二种要复杂得多。

二、OCR+文字搜图

OCR是指光学字符辨认，也便是咱们常说的文字辨认。OCR的完结方法是多样的，这儿运用Tesseract或许各种神经网络。OCR不是文本要点，因而这儿只简略介绍其运用。概况可见：/post/696437…

OCR+文字搜图的原理十分简略，便是先辨认文字，然后依据文字含糊查询找到相关图片即可。为了便利查询，这儿需求运用数据库。

2.1 文字辨认

运用pytesseract模块能够很便利完结OCR，详细代码如下：

import os, cv2
import numpy as np
import pytesseract
from tqdm import tqdm
from PIL import Image
from sqlalchemy import create_engine, String, select
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, Session
base_path = "G:\datasets\emoji"
files = [os.path.join(base_path, file) for file in os.listdir(base_path) if file.endswith(".jpg")]
for file in files:
try:
    image = Image.open(file)
    string = pytesseract.image_to_string(image, lang='chi_sim')
    print(file, ":", string.strip())
except Exception as e:
    pass

其间string便是辨认到的文本内容。pytesseract中也提供了批量辨认的接口，由于这儿存在一些过错图片，因而这儿不适用批量接口。

2.2 存储数据库

为了便利查询，能够把图片路径和图片中包含的文本内容存储到数据库中。这儿运用sqlalchemy+sqlite，代码如下：

from sqlalchemy import create_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy import String
class Base(DeclarativeBase):
    pass
class ImageInformation(Base):
    __tablename__ = "image_information"
    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    filepath: Mapped[str] = mapped_column(String(255))
    content: Mapped[str] = mapped_column(String(255))
    def __repr__(self) -> str:
        return f"User(id={self.id!r}, filepath={self.filepath!r}, content={self.content!r})"
engine = create_engine("sqlite:///image_search.db", echo=False)
Base.metadata.create_all(engine)

其间ImageInformation类对应咱们需求创立的数据库表。创立好后，辨认图片的文字，然后存储到图片数据库中：

base_path = "G:\datasets\emoji"
files = [os.path.join(base_path, file) for file in os.listdir(base_path) if file.endswith(".jpg")]
bar = tqdm(total=len(files))
for file in files:
    try:
        # 辨认文字
        image = Image.open(file)
        string = pytesseract.image_to_string(image, lang='chi_sim').strip()
        file = file[:255] if len(file) > 255 else file
        string = string[:255] if len(string) > 255 else string
        # 存储数据库
        with Session(engine) as session:
            info = ImageInformation(filepath=file, content=string)
            session.add_all([info])
            session.commit()
    except Exception as e:
        pass
    bar.update(1)

这个进程会比较久。

2.3 依据文字查找图片

完结上面的存储操作后，就能够开始依据文字查找图片了。这儿只需求运用简略的数据库查询操作即可完结，代码如下，咱们先把你好作为输入文本：

keyword = '你好'
w, h = 224, 224
with Session(engine) as session:
    stmt = select(ImageInformation).where(ImageInformation.content.contains(keyword)).limit(8)
    images = [cv2.resize(cv2.imread(ii.filepath), (w, h)) for ii in session.scalars(stmt)]
    if len(images) > 0:
        result = np.hstack(images)
        cv2.imwrite("result.jpg", result)
    else:
        print("没有找到成果")

下面是查询到的成果图片：

假如关键词改为喜欢，得到成果如下：

通过测试，发现在一些短文本查找中，这种方法比较奏效，但是在长文本则常常查找不到成果。一种改善方法是不存储文本自身，而是运用Bert等模型把文本转换成Embedding，然后存储Embedding。这样咱们就不能再运用sqlite了，而需求运用向量数据库。

三、依据Transformer的改善

在前面的比方中，查找成果十分依赖字符串匹配。比方查找鸡，只要图片中有鸡字才会被查找到，而与坤相关的图片则查找不到。为此咱们用Transformer对上面进行改善，主要思路便是先辨认文字，然后把文字交给文本编码器，转换成Embedding，然后在查找时查找输入文本和Embedding的类似度，这样就能够缓解上述问题。

3.1 创立数据库

这儿咱们还是运用向量数据库，向量数据库有许多挑选，这儿运用Milvus数据库，详细运用能够参考：milvus.io/docs/instal…

首要创立数据库和调集：

from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility
# 创立数据库
connections.connect(host='127.0.0.1', port='19530')
def create_milvus_collection(collection_name, dim):
    if utility.has_collection(collection_name):
        utility.drop_collection(collection_name)
    fields = [
        FieldSchema(name='id', dtype=DataType.INT64, descrition='ids', max_length=500, is_primary=True,
                    auto_id=True),
        FieldSchema(name='filepath', dtype=DataType.VARCHAR, description='filepath', max_length=512),
        FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, descrition='embedding vectors', dim=dim),
    ]
    schema = CollectionSchema(fields=fields, description='reverse image search')
    collection = Collection(name=collection_name, schema=schema)
    # create IVF_FLAT index for collection.
    index_params = {
        'metric_type': 'L2',
        'index_type': "IVF_FLAT",
        'params': {"nlist": 2048}
    }
    collection.create_index(field_name="embedding", index_params=index_params)
    return collection
collection = create_milvus_collection('image_information', 768)

调集中主要有filepath和embedding两个字段。其间embedding有768个维度，这是由Transformer决议的，我这儿挑选的Transformer输出768个维度，因而这儿填768。

3.2 text2vec

创立完结后，便是读取图片、辨认文字、文字编码、存入数据库。其间文字编码能够运用Transformers模块或许text2vec完结，这儿运用text2vec，其操作如下：

from text2vec import SentenceModel
model = SentenceModel('shibing624/text2vec-base-chinese')
embeddings = model.encode(['不要温柔地走进那个良夜'])
print(embeddings.shape)

在创立SentenceModel时传入对应的模型，然后调用model.encode方法即可。输出如下成果：

(1, 768)

其他操作则不详细解释，详细代码如下：

from text2vec import SentenceModel
model = SentenceModel('shibing624/text2vec-base-chinese', device="cuda")
base_path = "G:\datasets\emoji"
files = [os.path.join(base_path, file) for file in os.listdir(base_path) if file.endswith(".jpg")]
bar = tqdm(total=len(files))
for idx, file in enumerate(files):
    try:
        image = Image.open(file)
        string = pytesseract.image_to_string(image, lang='chi_sim').strip()
        embedding = model.encode([string])[0]
        collection.insert([
            [file],
            [embedding]
        ])
    except Exception as e:
        pass
    bar.update(1)

3.3 依据文字查找图片

在插入数据后，直接运用数据库的查询操作即可完结查找操作，详细代码如下：

import cv2
import numpy as np
from text2vec import SentenceModel
from pymilvus import connections, Collection
# 加载模型
model = SentenceModel('shibing624/text2vec-base-chinese', device="cuda")
# 衔接数据库，加载调集
connections.connect(host='127.0.0.1', port='19530')
collection = Collection(name='image_information')
search_params = {"metric_type": "L2", "params": {"nprobe": 10}, "offset": 5}
collection.load()
# 用来查询的文本
keyword = "今日不开心"
embedding = model.encode([keyword])
print(embedding.shape)
# 在数据库中查找
results = collection.search(
    data=[embedding[0]],
    anns_field='embedding',
    param=search_params,
    output_fields=['filepath'],
    limit=10,
    consistency_level="Strong"
)
collection.release()
# 展现查询成果
w, h = 224, 224
images = []
for result in results[0]:
    entity = result.entity
    filepath = entity.get('filepath')
    image = cv2.resize(cv2.imread(filepath), (w, h))
    images.append(np.array(image))
result = np.hstack(images)
cv2.imwrite("result.jpg", result)

向量数据库在查询时，能够依据向量的类似度回来查询成果。在前面咱们存储了句向量，所以咱们能够把查询文本转换成句向量，然后利用向量数据库的查询功用，查找类似成果。在上面代码中，咱们查询“今日不开心”，这次不再是字符串层面的查询，而是语句意义层面的查询，因而能够查询的不包含这些字符的图片，下面是查询成果：

把关键词修正为“我想吃饭”后得到下面的成果：

整体作用还是十分不错的。

不过前面的成果是建立在能在图片中辨认到文本的情况下，假如是咱们随手拍的照片，那么就不能运用上面的方法来完结文字查找图片。

四、依据图片意义的文字搜图

在多模态领域有许多组合模型，而咱们需求的便是Image-to-Text类模型。假如要手艺给图片增加画面描绘会十分麻烦，因而咱们挑选运用Image-to-Text模型完结自动辨认。

4.1 完结原理

依据图片意义的文字搜图的完结与前面依据OCR的类似，只不过需求把OCR修正为Image Captioning网络。在前面咱们的流程是：

读取图片
OCR辨认
把辨认成果转换成向量
存入数据库

现在只需求把第二步修正为运用Image Captioning生成图片描绘即可。后面部分则是完全一致的。

4.2 Image Captioning

像这类输入图片，输出画面描绘的使命叫做Image Captioning，用于这一使命的模型十分多。包含CNN+LSTM，Vit等都能够完结Image Captioning。两者都是一个Encoder-Decoder架构，运用CNN、Vit作为图片Encoder，将图片转换成特征图或许特征向量。然后把Encoder的输出作为Decoder的输入，并输入，然后顺次生成图片描绘。

以Vit为例，其结构如图：

Vit其实便是一个为图片设计的Transformer架构，在某些细节上为图片做了一些修正。

4.3 依据图片意义查找图片

首要咱们能够运用和前面相同的方法创立数据库，这儿不再重复，咱们复用前面的数据库image_information。然后需求修正插入数据的代码，首要来创立一个函数加载模型，并创立一个函数用于将图片转换成文本向量，代码如下：

import os
import torch
from tqdm import tqdm
from PIL import Image
from text2vec import SentenceModel
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
from pymilvus import Collection
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def load_model():
    """
    加载需求运用到的模型
    """
    sentence_model = SentenceModel('shibing624/text2vec-base-chinese', device="cuda")
    model = VisionEncoderDecoderModel.from_pretrained("bipin/image-caption-generator")
    image_processor = ViTImageProcessor.from_pretrained("bipin/image-caption-generator")
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    model.to(device)
    return sentence_model, model, image_processor, tokenizer
def get_embedding(filepath):
    """
    输入图片路径，将图片转成描绘向量
    """
    pixel_values = image_processor(images=[Image.open(filepath)], return_tensors="pt").pixel_values.to(device)
    output_ids = model.generate(pixel_values, num_beams=4, max_length=128)
    pred = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return sentence_model.encode(pred)

后续只需求调用get_embedding函数就能够完结图片到向量的转换。接下来便是修正插入数据的代码，详细如下：

connections.connect(host='127.0.0.1', port='19530')
collection = Collection("image_information")
collection.load()
sentence_model, model, image_processor, tokenizer = load_model()
base_path = "G:\datasets\people"
files = [os.path.join(base_path, file) for file in os.listdir(base_path)]
bar = tqdm(total=len(files))
for idx, file in enumerate(files):
    try:
        embedding = get_embedding(file)
        collection.insert([
            [file],
            [embedding]
        ])
    except Exception as e:
        pass
    bar.update(1)

最终则是搜图操作了，这个和前面是完全一样的：

search_params = {"metric_type": "L2", "params": {"nprobe": 10}, "offset": 5}
# 用来查询的文本
keyword = "girl"
embedding = sentence_model.encode([keyword])
# 在数据库中查找
results = collection.search(
    data=[embedding[0]],
    anns_field='embedding',
    param=search_params,
    output_fields=['filepath'],
    limit=10,
    consistency_level="Strong"
)
collection.release()
# 展现查询成果
w, h = 224, 224
images = []
for result in results[0]:
    entity = result.entity
    filepath = entity.get('filepath')
    image = cv2.resize(cv2.imread(filepath), (w, h))
    images.append(np.array(image))
result = np.hstack(images)
cv2.imwrite("result.jpg", result)

由于这儿挑选的Image Captioning模型输出为英文，因而这儿把英文作为关键字。这儿关键字为”girl”，下面是查找成果：

由于数据库中还存储了之前的表情包，因而表情包中关于与”girl”有关的表情包也查找出来了，比方”娘们”、”女性”等。

假如把关键字改为”smile girl”，查找成果如下：

假如图片数量足够，则能够得到一个比较好的查找成果。

上面的成果还能够有一些改善，在Image Captioning过程，咱们只生成了一个描绘。在许多情况下，这个描绘不一定准确，比方下面的图片：

能够描绘为“拿着话筒的姑娘”、“一个姑娘在浅笑”或许“一个拿着话筒的姑娘在浅笑”。因而咱们能够生成多个描绘，存入数据库，这样在查找时成果能够更准确。能够通过修正temperature参数来生成不同的描绘：

output_ids = model.generate(pixel_values, num_beams=4, max_length=128, temperature=0.8)

当temperature小于1时，生成成果带有随机性。temperature越小，成果越随机。

如何实现文字搜图

一、前言

二、OCR+文字搜图

2.1 文字辨认

2.2 存储数据库

2.3 依据文字查找图片

三、依据Transformer的改善

3.1 创立数据库

3.2 text2vec

3.3 依据文字查找图片

四、依据图片意义的文字搜图

4.1 完结原理

4.2 Image Captioning

4.3 依据图片意义查找图片

作者信息

如何实现文字搜图

一、前言

二、OCR+文字搜图

2.1 文字辨认

2.2 存储数据库

2.3 依据文字查找图片

三、依据Transformer的改善

3.1 创立数据库

3.2 text2vec

3.3 依据文字查找图片

四、依据图片意义的文字搜图

4.1 完结原理

4.2 Image Captioning

4.3 依据图片意义查找图片

相关文章

Flutter异常监控 – 肆 | Rollbar源码赏析

Wireshark 基础 | 显示过滤篇

[Stable Diffusion] 三句话，让Ai帮你画18万张图

2022年我的面试万字总结（JS篇下）

作者信息