本文为稀土技能社区首发签约文章，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 依据图片意义查找图片

相关文章

ROS2 tf2广播器(broadcaster)、侦听器(listener)

美团2024届暑期实习第一轮后端笔试详解

Topic太多！RocketMQ炸了！

【干货】使用Canal 解决数据同步场景中的疑难杂症！！！

作者信息

推广