使用🤗数据集进行图像搜索
使用🤗数据集进行图像搜索
🤗 datasets 是一个可以轻松访问和共享数据集的库。它还可以高效处理数据，包括处理不适合内存的数据。
当 datasets 最初推出时，它主要与文本数据相关。然而，最近，datasets 还增加了对音频和图像的支持。特别地，现在有一个用于图像的 datasets 特征类型。之前的一篇博文展示了如何使用 🤗 transformers 和 datasets 来训练图像分类模型。在这篇博文中，我们将看到如何将 datasets 和其他几个库结合起来创建一个图像搜索应用。
首先，我们将安装 datasets。由于我们将要处理图像，我们还需要安装 pillow。我们还需要 sentence_transformers 和 faiss。我们将在下面更详细地介绍它们。我们还要安装 rich – 我们在这里只会简单使用它，但它是一个非常方便的包，我真的建议进一步探索它！
!pip install datasets pillow rich faiss-gpu sentence_transformers 首先，让我们来看一下图像特征。我们可以使用很棒的 rich 库来查看 Python 对象（函数、类等）
from rich import inspect
import datasets
inspect(datasets.Image, help=True)
╭───────────────────────── <class 'datasets.features.image.Image'> ─────────────────────────╮
│ class Image(decode: bool = True, id: Union[str, NoneType] = None) -> None: │
│ │
│ Image feature 用于从图像文件中读取图像数据。 │
│ │
│ 输入：Image feature 接受以下输入： │
│ - 一个 :obj:`str`：图像文件的绝对路径（即允许随机访问）。 │
│ - 一个带有以下键的 :obj:`dict`： │
│ │
│ - path：包含到归档文件的图像文件的相对路径的字符串。 │
│ - bytes：图像文件的字节。 │
│ │
│ 这对于具有顺序访问的归档文件非常有用。 │
│ │
│ - 一个 :obj:`np.ndarray`：表示图像的 NumPy 数组。 │
│ - 一个 :obj:`PIL.Image.Image`：PIL 图像对象。 │
│ │
│ 参数： │
│ decode (:obj:`bool`，默认为 ``True``)：是否解码图像数据。如果为 `False`，则返回基础字典，格式为 {"path": image_path, "bytes": │
│ image_bytes}。 │
│ │
│ decode = True │
│ dtype = 'PIL.Image.Image' │
│ id = None │
│ pa_type = StructType(struct<bytes: binary, path: string>) │
╰───────────────────────────────────────────────────────────────────────────────────────────╯我们可以看到有几种不同的方式可以传入我们的图像。我们一会儿会回到这个问题。
datasets 库的一个非常好的功能（除了处理数据、内存映射等功能之外）是你可以获得一些额外的好东西。其中之一是向数据集添加 faiss 索引的能力。 faiss 是一个“用于高效相似度搜索和密集向量聚类的库”。
datasets 的文档展示了使用 faiss 索引进行文本检索的示例。在本文中，我们将看看我们是否可以对图像执行同样的操作。
数据集：“数字化图书 – 被识别为装饰品的图像。c. 1510 – c. 1900”
这是一个数据集，其中的图像来自英国图书馆的一组数字化图书。这些图像来自不同时间段和领域的图书。这些图像是通过使用每本书的 OCR 输出中的信息提取的。因此，我们知道这些图像来自哪本书，但不一定知道图像中显示了什么。
为了帮助克服这个问题，一些尝试包括将图像上传到flickr。这样可以让人们给图像打标签或将它们放入不同的类别中。
还有一些项目使用机器学习来标记数据集。这项工作使得可以通过标签进行搜索，但我们可能希望有更”丰富”的搜索能力。对于这个特定的实验，我们将使用包含”装饰”的集合的子集。这个数据集要小一些，所以用来进行实验更好。我们可以从英国图书馆的数据仓库获取完整的数据：https://doi.org/10.21250/db17 。由于完整的数据集仍然相当大，您可能想从一个较小的样本开始。
创建我们的数据集
我们的数据集由一个包含子目录的文件夹组成，其中包含图像。这是一种共享图像数据集的相当标准的格式。多亏了最近合并的拉取请求，我们可以使用 datasets ImageFolder 加载器直接加载这个数据集 🤯
from datasets import load_dataset
dataset = load_dataset("imagefolder", data_files="https://zenodo.org/record/6224034/files/embellishments_sample.zip?download=1")让我们看看我们得到了什么。
dataset
DatasetDict({
train: Dataset({
features: ['image', 'label'],
num_rows: 10000
})
})我们可以获得一个 DatasetDict，其中包含一个带有图像和标签特征的数据集。由于这里没有训练/验证拆分，让我们获取数据集的训练部分。我们还可以查看数据集中的一个示例，看看它是什么样子的。
dataset = dataset["train"]
dataset[0]
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=358x461 at 0x7F9488DBB090>,
'label': 208}让我们从标签列开始。它包含了我们的图像的父文件夹。在这种情况下，标签列表示从中提取图像的书籍的出版年份。我们可以使用 dataset.features 查看这个映射：
dataset.features['label']在这个特定的数据集中，图像文件名还包含一些关于提取图像的书籍的元数据。我们可以以几种方式获得这些信息。
当我们查看数据集中的一个示例时，可以看到 image 特征是一个 PIL.JpegImagePlugin.JpegImageFile 。由于 PIL.Images 有一个 filename 属性，我们可以通过访问这个属性来获取文件名。
dataset[0]['image'].filename
/root/.cache/huggingface/datasets/downloads/extracted/f324a87ed7bf3a6b83b8a353096fbd9500d6e7956e55c3d96d2b23cc03146582/embellishments_sample/1920/000499442_0_000579_1_[The Ring and the Book etc ]_1920.jpg由于以后我们可能需要轻松访问这些信息，让我们创建一个新列来提取文件名。为此，我们将使用 map 方法。
dataset = dataset.map(lambda example: {"fname": example['image'].filename.split("/")[-1]})我们可以查看一个示例，看看现在是什么样子。
dataset[0]
{'fname': '000499442_0_000579_1_[The Ring and the Book etc ]_1920.jpg',
'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=358x461 at 0x7F94862A9650>,
'label': 208}我们现在有了元数据。让我们快点看看一些图片吧！如果我们访问一个示例并索引到 image 列，我们将看到我们的图像 😃
dataset[10]['image']
注意：在此博文的早期版本中，下载和加载图像的步骤要复杂得多。新的 ImageFolder 加载器使这个过程变得更加容易 😀 特别是，由于 datasets 已经为我们处理了这个问题，我们无需担心如何加载图像。
将所有东西推到中心！
🤗生态系统中的一个超级棒的功能是Hugging Face Hub。我们可以使用Hub来访问模型和数据集。它经常用于与他人分享工作，但也可以作为工作进展的有用工具。最近，datasets添加了一个push_to_hub方法，允许您以最小的麻烦将数据集推送到Hub。这可以通过允许您传递已完成所有转换等的数据集来真正有帮助。
现在，我们将数据集推送到Hub，并在初始时保持私有。
根据您运行代码的位置，您可能需要进行身份验证。您可以使用huggingface-cli login命令进行身份验证，或者如果您在笔记本中运行，则可以使用notebook_login
from huggingface_hub import notebook_login
notebook_login()
dataset.push_to_hub('davanstrien/embellishments-sample', private=True)注意：在以前的博客文章版本中，我们需要执行一些额外的步骤，以确保在使用
push_to_hub时嵌入图像。由于这个拉取请求，我们不再需要担心这些额外的步骤。我们只需要确保embed_external_files=True（默认行为）。
切换机器
到目前为止，我们已经创建了一个数据集并将其移到了Hub。这意味着我们可以在其他地方继续工作/使用数据集。
在这个特定的示例中，拥有GPU的访问权限很重要。使用Hub作为传递数据的方式，我们可以从笔记本电脑开始，在Google Colab上继续工作。
如果我们切换到新的机器，可能需要重新登录。一旦我们这样做了，就可以加载我们的数据集
from datasets import load_dataset
dataset = load_dataset("davanstrien/embellishments-sample", use_auth_token=True)创建嵌入🕸
现在我们有一个包含许多图像的数据集。为了开始创建我们的图像搜索应用程序，我们需要对这些图像进行嵌入。有多种方法可以尝试做到这一点，但其中一种可能的方法是使用sentence_transformers库中的CLIP模型。OpenAI的CLIP模型学习了图像和文本的联合表示，这对我们想要做的事情非常有用，因为我们希望输入文本并返回图像。
我们可以使用SentenceTransformer类下载模型。
from sentence_transformers import SentenceTransformer
model = SentenceTransformer('clip-ViT-B-32')该模型将接受图像或一些文本作为输入，并返回一个嵌入。我们可以使用datasets的map方法使用此模型对所有图像进行编码。当我们调用map时，我们返回一个字典，其中包含模型返回的嵌入的键embeddings。当我们调用模型时，我们还传递device='cuda'，这确保我们在GPU上进行编码。
ds_with_embeddings = dataset.map(
lambda example: {'embeddings':model.encode(example['image'], device='cuda')}, batched=True, batch_size=32)我们可以使用push_to_hub将我们的工作“保存”到Hub。
ds_with_embeddings.push_to_hub('davanstrien/embellishments-sample', private=True)如果我们要切换到另一台机器，我们可以通过从Hub加载它来重新获取我们的工作😃
from datasets import load_dataset
ds_with_embeddings = load_dataset("davanstrien/embellishments-sample", use_auth_token=True)现在我们有了一个包含我们图像嵌入的新列。我们可以手动搜索这些嵌入并将它们与一些输入嵌入进行比较，但数据集有一个add_faiss_index方法。这使用faiss库为搜索嵌入创建一个高效的索引。关于此库的更多背景信息，您可以观看这个YouTube视频
ds_with_embeddings['train'].add_faiss_index(column='embeddings')
数据集({
特征: ['fname', 'year', 'path', 'image', 'embeddings'],
行数: 10000
})图像搜索
注意，这些示例是从数据集的完整版本生成的，因此您可能会得到稍微不同的结果。
现在我们拥有了创建简单图像搜索所需的一切。我们可以使用之前用于编码图像的模型来编码一些输入文本。这将作为我们尝试寻找相似示例的提示。让我们从“蒸汽机”开始。
prompt = model.encode("蒸汽机")我们可以使用数据集库中的另一种方法get_nearest_examples来获取与我们输入提示嵌入接近的图像。我们可以传入我们希望返回的结果数量。
scores, retrieved_examples = ds_with_embeddings['train'].get_nearest_examples('embeddings', prompt, k=9)我们可以提取返回结果中的第一个示例：
retrieved_examples['image'][0]
这并不完全是一台蒸汽机，但也不是完全奇怪的结果。我们可以绘制其他结果来查看返回的内容。
import matplotlib.pyplot as plt
plt.figure(figsize=(20, 20))
columns = 3
for i in range(9):
image = retrieved_examples['image'][i]
plt.subplot(9 / columns + 1, columns, i + 1)
plt.imshow(image)
其中一些结果看起来与我们的输入提示相当接近。我们可以将其包装在一个函数中，以便更轻松地尝试不同的提示。
def get_image_from_text(text_prompt, number_to_retrieve=9):
prompt = model.encode(text_prompt)
scores, retrieved_examples = ds_with_embeddings['train'].get_nearest_examples('embeddings', prompt, k=number_to_retrieve)
plt.figure(figsize=(20, 20))
columns = 3
for i in range(9):
image = retrieved_examples['image'][i]
plt.title(text_prompt)
plt.subplot(9 / columns + 1, columns, i + 1)
plt.imshow(image)
get_image_from_text("太阳背后的山的插图")
尝试一系列提示 ✨
现在我们有了一个获取几个结果的函数，我们可以尝试一系列不同的提示：
-
对于其中一些提示，我会选择广义的“类别”，例如：“一种乐器”或“一个动物”，其他提示是具体的，例如：“一把吉他”。
-
出于兴趣，我还尝试了一个布尔运算符：“一只猫或一只狗的插图”。
-
最后，我尝试了一些更抽象的内容：“一个空洞的深渊”
prompts = ["一种乐器", "一把吉他", "一个动物", "一只猫或一只狗的插图", "一个空洞的深渊"]
for prompt in prompts:
get_image_from_text(prompt)
我们可以看到这些结果并不总是完全正确的，但它们通常是合理的。它似乎已经可以用于在该数据集中搜索图像的语义内容。但是，我们可能会暂时不分享这个结果…
创建一个Hugging Face Space? 🤷🏼
对于这种项目，一个明显的下一步是创建一个Hugging Face Space演示。这是我为其他模型做过的事情。
从我们到达这里的点来看，设置Gradio应用程序是一个相当简单的过程。这是该应用程序的屏幕截图：
然而，我对立即公开这一点有些犹豫。通过查看CLIP模型的模型卡，我们可以了解其主要预期用途：
主要预期用途
我们主要设想该模型将被研究人员用于更好地理解计算机视觉模型的鲁棒性、泛化能力以及其他能力、偏见和约束。来源
这与我们在这里所感兴趣的内容非常接近。特别是我们可能对该模型如何处理我们数据集中的图像类型（主要是来自19世纪的书籍插图）感兴趣。我们数据集中的图像与训练数据（可能）相当不同。由于一些图像中还包含文本，这可能有助于CLIP，因为它显示了一些OCR能力。
然而，查看模型卡中的超出范围的用例：
超出范围的用例
目前，任何部署的模型用例（无论是商业用途还是非商业用途）都超出了范围。非部署的用例，如在受限环境中进行图像搜索，也不建议，除非对模型进行了特定的领域测试，并使用了特定的、固定的类别分类法。这是因为我们的安全评估表明，尤其是考虑到CLIP在不同类别分类法下的性能变化的可变性，需要进行任务特定的测试。这使得目前在任何用例中对模型进行未经测试和不受限制的部署可能会有潜在的危害。来源
暗示“部署”不是一个好主意。尽管我得到的结果很有趣，但我还没有对该模型进行足够多的探索（也没有进行更系统化的评估其性能和偏见），无法对“部署”它感到自信。另一个额外的考虑因素是目标数据集本身。这些图像来自涵盖各种主题和时期的书籍。有很多书籍代表了殖民态度，因此其中一些图像可能以消极的方式代表某些人群。这可能与允许将任意文本输入编码为提示的工具相结合，可能会产生不好的组合。
可能有办法解决这个问题，但这需要更多的思考。
结论
尽管我们没有漂亮的演示来展示，但我们已经看到了如何使用datasets来：
- 将图像加载到新的
Image特征类型中 - 使用
push_to_hub保存我们的工作，并将数据在不同机器/会话之间移动 - 为图像创建一个
faiss索引，我们可以使用它从文本（或图像）输入中检索图像。
