【Kaggle】Stable Diffusion - Image to Prompts竞赛代码初步理解

最新推荐文章于 2025-08-28 14:23:09 发布

原创最新推荐文章于 2025-08-28 14:23:09 发布 · 1.5k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#stable diffusion #python #kaggle

kaggle 专栏收录该内容

13 篇文章

订阅专栏

本文深入解析了Kaggle上的Stable Diffusion项目，涵盖了从预处理OFA模型、模型EDA、图像到提示文本转换的全过程。文章通过实例展示了如何加载预训练模型，对图像进行归一化处理，使用SentenceTransformer和CLIP模型进行文本嵌入，以及如何生成和评估图像的提示描述。此外，还探讨了如何在Kaggle环境中安装和使用特定版本的transformers库。

一、前言

此次代码集成了 CLIP Interrogator、OFA 模型和 ViT 模型。

首先安装指定版本的 transformers 库：

transformers-4.18.0.dev0-py3-none-any.whl 是一个 transformers 库的文件，它的命名方式表示这是一个开发版本（dev）的预构建轮子（wheel）文件。

轮子文件是 Python 包的一种打包格式，可以通过 pip 安装。

如果您想要安装这个特定版本的 transformers 库，可以使用以下命令：

pip install transformers-4.18.0.dev0-py3-none-any.whl

请确保您位于包含该文件的目录，并且在运行该命令之前已经安装了适当的依赖项。

请注意，这是一个开发版本的预构建文件，可能包含尚未正式发布的功能或存在 bug。如果您只是想使用稳定版本的 transformers 库，建议使用正式发布的版本，例如：

pip install transformers

这将安装最新的稳定版本，而不是开发版本。

在比赛中，我们可以引入的方法为：

!pip install -q /kaggle/input/stable-diffusion-data/transformers-4.18.0.dev0-py3-none-any.whl

在 Kaggle 环境中使用 Jupyter Notebook 或 JupyterLab 的一个指令。

! 符号是 Jupyter Notebook 或 JupyterLab 中的一个魔术命令前缀，用于执行系统级命令。
pip install 是用于安装 Python 包的 pip 命令。
-q 或 –quiet 参数是 pip 命令的选项之一，用于使安装过程静默，即不显示安装的详细信息。
/kaggle/input/stable-diffusion-data/transformers-4.18.0.dev0-py3-none-any.whl 是一个文件路径，表示要安装的 transformers 库的轮子文件（.whl 文件）。根据路径中的前缀 /kaggle/input，可以推断这是在 Kaggle 环境中安装位于 /kaggle/input 目录下的本地文件。

综上所述，该命令的含义是在 Kaggle 环境中安装指定路径下的 transformers-4.18.0.dev0-py3-none-any.whl 轮子文件，并且在安装过程中不显示详细信息。这将通过使用 pip 命令将该轮子文件作为本地文件进行安装。

二、导包

import os
import sys
import glob
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from PIL import Image
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from transformers import OFATokenizer, OFAModel
from transformers.models.ofa.generate import sequence_generator
import gc

这是一个 Python 脚本中的导入语句，它用于将一些常用的 Python 和深度学习库导入到脚本中。以下是每个导入库的简要介绍：

os：提供了一个与操作系统交互的简单方法，例如读取和写入文件。
sys：提供了一个与 Python 解释器交互的方法，例如修改 sys.path 导入路径、获取命令行参数等。
glob：提供了一个通用文件和目录的匹配模式方法。
pathlib.Path：提供了一个简单的面向对象的路径类，用于操作文件和目录。
numpy：提供了一个用于进行数学和科学计算的 Python 库。
pandas：提供了一个用于数据分析的 Python 库，可以用于处理和操作大型数据集。
matplotlib.pyplot：提供了一个用于绘制数据可视化图形的 Python 库。
PIL.Image：提供了一个 Python 图像处理库，可以用于处理和操作图像。
torch：提供了一个 PyTorch 深度学习框架库，用于创建和训练神经网络模型。
torch.utils.data.Dataset：提供了一个 PyTorch 数据集抽象类，用于创建自定义数据集类。
torchvision.transforms：提供了一些 PyTorch 中用于数据增强和预处理的转换类。
transformers.OFATokenizer：提供了一个 OneFlow Aware Tokenizer（OFA）类，用于在 PyTorch 中进行模型搜索和架构优化。
transformers.OFAModel：提供了一个 OneFlow Aware Model（OFA）类，用于在 PyTorch 中进行模型搜索和架构优化。
transformers.models.ofa.generate.sequence_generator：提供了一个序列生成器，用于生成一系列优化的模型架构。
gc：是 Python 中的垃圾回收库，可以用于管理内存。

CKPT_DIR = "/kaggle/input/stable-diffusion-data/OFA-large-caption/"
IMAGE_DIR = "/kaggle/input/stable-diffusion-image-to-prompts/images"

BATCH_SIZE = 24

这些是一些变量的赋值语句：

CKPT_DIR：设置为 “/kaggle/input/stable-diffusion-data/OFA-large-caption/”，表示一个目录路径，指向存储 OneFlow Aware (OFA) 模型的检查点文件的位置。
IMAGE_DIR：设置为 “/kaggle/input/stable-diffusion-image-to-prompts/images”，表示一个目录路径，指向存储图像文件的位置。
BATCH_SIZE：设置为 24，表示批量处理数据时的批量大小，即一次传递给模型的样本数量。

这些变量的值可以根据需要进行调整，用于指定数据的路径、模型的保存位置以及批量处理数据时的批量大小。

三、加载预训练的 OFA 模型

mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]
resolution = 480
patch_resize_transform = transforms.Compose([
        lambda image: image.convert("RGB"),
        transforms.Resize((resolution, resolution), interpolation=Image.BICUBIC),
        transforms.ToTensor(), 
        transforms.Normalize(mean=mean, std=std)
    ])

tokenizer = OFATokenizer.from_pretrained(CKPT_DIR)
model = OFAModel.from_pretrained(CKPT_DIR, use_cache=False).cuda()
txt = " what does the image describe?"
inputs = tokenizer([txt], return_tensors="pt").input_ids

让我们逐行解读这段代码：

mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]：这行代码定义了 mean 和 std 两个变量，它们是归一化图像时使用的均值和标准差。这里设置的均值和标准差都是 [0.5, 0.5, 0.5]，表示将图像的每个通道的像素值缩放到范围 [-1, 1]。
resolution = 480：这行代码定义了 resolution 变量，它表示将图像调整为的分辨率大小。在这里，图像将被调整为 480 x 480 像素。
patch_resize_transform = transforms.Compose([…])：这行代码定义了一个转换序列 patch_resize_transform，它将应用于输入图像。转换序列中的每个操作按顺序应用于图像。
image.convert(“RGB”) 将图像转换为 RGB 模式，确保图像有三个通道。
transforms.Resize((resolution, resolution), interpolation=Image.BICUBIC) 调整图像大小为给定的 resolution，使用双三次插值方法进行调整。
transforms.ToTensor() 将图像转换为张量形式，将像素值缩放到范围 [0, 1]。
transforms.Normalize(mean=mean, std=std) 对图像进行标准化，将像素值归一化为均值为 mean，标准差为 std 的分布。
tokenizer = OFATokenizer.from_pretrained(CKPT_DIR)：这行代码创建一个 OFATokenizer 对象，从预训练模型的检查点文件中加载 tokenizer。
model = OFAModel.from_pretrained(CKPT_DIR, use_cache=False).cuda()：这行代码创建一个 OFAModel 对象，并加载预训练模型的权重。use_cache = False 表示不使用缓存。
txt = " what does the image describe?"：这行代码定义了一个字符串变量 txt，它包含了图像描述的文本。
inputs = tokenizer([txt], return_tensors=“pt”).input_ids：这行代码使用之前创建的 tokenizer 对象将文本 txt 编码为模型输入的张量形式。tokenizer([txt]) 将文本转换为 tokens，并返回一个字典对象。.input_ids 提取了输入 tokens 的张量表示。最终，inputs 变量将包含文本编码后的输入张量。

总的来说，这段代码的目的是为了准备图像和文本数据，以便将它们输入到 OFA 模型中进行处理和生成。

四、模型EDA

sample_images = glob.glob("/kaggle/input/stable-diffusion-image-to-prompts/images/*")[:7]
fig, ax = plt.subplots(7,1, figsize=(4,35))

for i,impath in enumerate(sample_images):
    image = Image.open(impath)
    image_t = patch_resize_transform(image).cuda().unsqueeze(0)
    out = model.generate(inputs.cuda(), patch_images=image_t.cuda(), num_beams=5, no_repeat_ngram_size=2)
    out_captions = tokenizer.batch_decode(out, skip_special_tokens=True)
    ax[i].imshow(image)
    ax[i].text(1.1, .5, out_captions[0], horizontalalignment='left', verticalalignment='center', transform=ax[i].transAxes)

让我们逐行解释这段代码：

sample_images = glob.glob(“/kaggle/input/stable-diffusion-image-to-prompts/images/*”)[:7]：这行代码使用 glob.glob 函数获取 /kaggle/input/stable-diffusion-image-to-prompts/images/ 目录下的图像文件路径，并选择前 7 个图像文件。这些文件路径被存储在 sample_images 列表中。
fig, ax = plt.subplots(7, 1, figsize=(4, 35))：这行代码创建了一个包含 7 行、1 列的子图布局，每个子图的大小为 (4, 35)。返回的 fig 对象表示整个图像，ax 对象是包含 7 个子图的数组。
for i, impath in enumerate(sample_images):：这是一个循环，遍历 sample_images 列表中的图像文件路径。enumerate 函数用于同时迭代列表中的元素和它们的索引。在每次迭代中，i 是索引，impath 是当前的图像文件路径。
image = Image.open(impath)：这行代码使用 PIL 库的 Image.open 函数打开图像文件，并将图像对象存储在 image 变量中。
image_t = patch_resize_transform(image).cuda().unsqueeze(0)：这行代码对打开的图像 image 应用之前定义的 patch_resize_transform 转换序列，将图像调整大小并进行标准化处理。然后，使用 .cuda() 将图像张量移动到 GPU 上，并使用 unsqueeze(0) 在批次维度上添加一个维度。
out = model.generate(inputs.cuda(), patch_images=image_t.cuda(), num_beams=5, no_repeat_ngram_size=2)：这行代码使用预训练的 OFA 模型 model 生成文本。它接收一个输入文本的张量 inputs，以及调整大小和标准化后的图像张量 image_t。num_beams=5 表示使用束搜索方法生成多个可能的文本输出，no_repeat_ngram_size=2 表示生成的文本中不会有连续重复的 2-gram。
out_captions = tokenizer.batch_decode(out, skip_special_tokens=True)：这行代码使用 tokenizer 将模型生成的输出 out 解码为文本。skip_special_tokens = True 表示跳过特殊标记，如起始和结束标记。
ax[i].imshow(image)：这行代码在第 i 个子图中显示图像。
ax[i].text(1.1, .5, out_captions[0], horizontalalignment=‘left’, verticalalignment=‘center’, transform=ax[i].transAxes)：这行代码在第 i 个子图中添加文本标注。1.1, .5 是文本的位置坐标，horizontalalignment = ‘left’ 表示文本水平对齐方式为左对齐，verticalalignment = ‘center’ 表示文本垂直对齐方式为居中对齐。ax[i].transAxes 表示使用子图坐标系进行转换。

这段代码的目的是展示图像并在每张图像上显示由 OFA 模型生成的文本描述。它首先遍历了前 7 个图像文件的路径，然后对每张图像进行处理：调整大小、标准化，并传递给 OFA 模型生成文本描述。生成的文本描述被解码后，用作文本标注，并与相应的图像一起显示在子图中。

最终的结果是，在一个具有 7 行、1 列的图像布局中，显示了每张图像以及由 OFA 模型生成的相应文本描述。

五、Inference

sys.path.append('../input/sentence-transformers-222/sentence-transformers')
from sentence_transformers import SentenceTransformer, models

comp_path = Path('../input/stable-diffusion-image-to-prompts/')
st_model = SentenceTransformer('/kaggle/input/sentence-transformers-222/all-MiniLM-L6-v2')

让我们逐行解读上面的代码：

sys.path.append(‘…/input/sentence-transformers-222/sentence-transformers’)：这行代码将 ‘…/input/sentence-transformers-222/sentence-transformers’ 目录添加到 Python 的模块搜索路径中，以便能够导入其中的模块。
from sentence_transformers import SentenceTransformer, models：这行代码从 sentence_transformers 模块中导入 SentenceTransformer 和 models 类。这些类提供了用于处理和生成句子向量的功能。
comp_path = Path(‘…/input/stable-diffusion-image-to-prompts/’)：这行代码创建了一个 Path 对象 comp_path，指向 ‘…/input/stable-diffusion-image-to-prompts/’ 目录。Path 对象提供了一些方便的方法来处理文件路径。
st_model = SentenceTransformer(‘/kaggle/input/sentence-transformers-222/all-MiniLM-L6-v2’)：这行代码创建了一个 SentenceTransformer 对象 st_model，并加载了预训练的句子转换模型。模型文件位于 ‘/kaggle/input/sentence-transformers-222/all-MiniLM-L6-v2’ 路径下。该模型使用 MiniLM-L6 架构，用于将句子转换为向量表示。

这段代码的目的是导入句子转换器模块，并创建一个 SentenceTransformer 对象，以便在后续的文本处理中使用。同时，还定义了一个路径对象 comp_path，指向图像和文本数据的目录。

class ImageGen(Dataset):
    def __init__(self, root, batch_size=32):
        self.root = root
        self.im_paths = os.listdir(self.root)
        self.batch_size = batch_size
        self.sz = len(self.im_paths)
        self.genlen = self.sz//self.batch_size + int(self.sz%self.batch_size > 0)
        
    def __getitem__(self, index):
        if index >= self.genlen:
            raise IndexError("Out of bounds")
        
        l, r = index*self.batch_size, min(self.sz, (index+1)*self.batch_size)
        
        f_paths = [os.path.join(self.root, self.im_paths[i]) for i in range(l,r)]
        f_ids = [self.im_paths[i][:-4] for i in range(l,r)]
        
        ims = [Image.open(f_path) for f_path in f_paths]
        ims = [patch_resize_transform(im).cuda().unsqueeze(0) for im in ims]
        ims = torch.cat(ims)
        
        return ims, f_ids
    
    def __len__(self):
        return self.genlen

让我们逐行解读上面的代码：

class ImageGen(Dataset):：这行代码定义了一个名为 ImageGen 的类，它继承自 Dataset 类，表示一个自定义的数据集。
def init(self, root, batch_size=32):：这是 ImageGen 类的构造函数。它接受一个 root 参数，表示图像文件的根目录路径，以及一个可选的 batch_size 参数，表示批量大小，默认为 32。
self.root = root：这行代码将构造函数中传入的 root 参数赋值给 self.root，表示图像文件的根目录路径。
self.im_paths = os.listdir(self.root)：这行代码使用 os.listdir 函数列出指定目录 self.root 中的所有文件和文件夹，并将结果存储在 self.im_paths 列表中。这里假设 self.root 目录下仅包含图像文件。
self.batch_size = batch_size：这行代码将构造函数中传入的 batch_size 参数赋值给 self.batch_size，表示批量大小。
self.sz = len(self.im_paths)：这行代码计算图像文件列表 self.im_paths 的长度，并将结果存储在 self.sz 中，表示图像文件的总数。
self.genlen = self.sz // self.batch_size + int(self.sz%self.batch_size > 0)：这行代码计算生成器的长度 self.genlen。它通过将总图像数 self.sz 除以批量大小 self.batch_size，再加上一个余数为 0 或 1 的值，来确定生成器的迭代次数。
def getitem(self, index):：这是 ImageGen 类的 getitem 方法，用于获取给定索引 index 处的图像数据。
if index >= self.genlen:：这行代码检查索引 index 是否超出生成器的长度，如果超出，则抛出 IndexError 异常，表示索引越界。
l, r = index * self.batch_size, min(self.sz, (index + 1) * self.batch_size)：这行代码计算图像文件索引范围。它根据当前索引 index 和批量大小 self.batch_size 计算起始索引 l 和结束索引 r。l 表示当前批次图像文件的起始索引，r 表示当前批次图像文件的结束索引（不包括 r）。
f_paths = [os.path.join(self.root, self.im_paths[i]) for i in range(l, r)]：这行代码根据起始索引 l 和结束索引 r，获取对应的图像文件路径，并将它们存储在 f_paths 列表中。os.path.join 函数用于将目录路径和文件名组合成完整的文件路径。
f_ids = [self.im_paths[i][:-4] for i in range(l,r)]：这行代码使用列表推导式生成一个列表 f_ids，其中包含当前批次图像文件的文件名（不包括扩展名）。它通过遍历起始索引 l 到结束索引 r 之间的索引，并使用切片操作 [:-4] 来去除文件名的扩展名部分（假设文件名的扩展名为 4 个字符）。
ims = [Image.open(f_path) for f_path in f_paths]：这行代码使用列表推导式打开当前批次图像文件，并将打开的图像对象存储在列表 ims 中。它遍历 f_paths 列表中的文件路径，并使用 Image.open 函数打开图像文件。
ims = [patch_resize_transform(im).cuda().unsqueeze(0) for im in ims]，ims = torch.cat(ims)：这几行代码对图像进行预处理和转换。它遍历列表 ims 中的每个图像对象 im，并依次应用 patch_resize_transform 变换。变换包括将图像转换为 RGB 格式、调整图像大小、转换为张量，并进行均值标准化。转换后的图像张量被添加到新的列表中。最后，使用 torch.cat 函数将所有图像张量在维度 0 上进行拼接，形成一个批次的图像数据。
return ims, f_ids：这行代码返回当前批次的图像数据 ims 和对应的文件名列表 f_ids。
def len(self): return self.genlen：这是 ImageGen 类的 len 方法，用于返回生成器的长度，即生成器可以迭代的次数。

ImageGen 类的作用是作为自定义的数据集类，提供了对图像数据的加载和处理功能。通过索引获取图像数据时，它根据索引确定当前批次的图像文件范围，并将对应的图像文件路径转换为图像张量。同时，它还提供了生成器的长度，以便在训练过程中迭代地获取图像数据。

from gensim.parsing.preprocessing import remove_stopwords

sub_ids = []
sub_embeds = []

imgen = ImageGen(IMAGE_DIR, BATCH_SIZE)

for b in imgen:
    for j in range(len(b[1])):
        sub_ids.extend([f"{b[1][j]}_{i}" for i in range(384)])
    
    img_batch = b[0]
    out = model.generate(inputs.repeat(len(img_batch), 1).cuda(), patch_images=img_batch, num_beams=5, no_repeat_ngram_size=2)
    out_captions = tokenizer.batch_decode(out, skip_special_tokens=True)
    out_captions = [remove_stopwords(text) for text in out_captions]
    
    embeddings = st_model.encode(out_captions).flatten()
    sub_embeds.extend(embeddings)

让我们逐行解读上面的代码：

from gensim.parsing.preprocessing import remove_stopwords：这行代码导入 gensim 库中的 remove_stopwords 函数，用于去除句子中的停用词。
sub_ids = []：这行代码创建一个空列表 sub_ids，用于存储生成的子图像标识。
sub_embeds = []：这行代码创建一个空列表 sub_embeds，用于存储生成的子图像的嵌入向量。
imgen = ImageGen(IMAGE_DIR, BATCH_SIZE)：这行代码创建了一个 ImageGen 对象 imgen，并传入图像文件的根目录路径 IMAGE_DIR 和批量大小 BATCH_SIZE。
for b in imgen:：这是一个迭代循环，它遍历 imgen 对象，获取每个批次的图像数据。
for j in range(len(b[1])):：这是一个嵌套的循环，它遍历当前批次的文件名列表 b[1] 的索引。
sub_ids.extend([f"{b[1][j]}_{i}" for i in range(384)])：这行代码使用列表推导式生成子图像的标识，并将它们添加到 sub_ids 列表中。生成的标识由当前文件名 b[1][j] 和索引 i 组成，形式为 “文件名_索引”。
img_batch = b[0]：这行代码将当前批次的图像数据存储在变量 img_batch 中。
out = model.generate(inputs.repeat(len(img_batch), 1).cuda(), patch_images=img_batch, num_beams=5, no_repeat_ngram_size=2)：这行代码使用预训练的模型 model 生成图像的描述文本。它将文本输入张量 inputs 进行重复，使其与图像批次大小匹配，并调用模型的 generate 方法生成描述文本。生成的文本通过 num_beams 和 no_repeat_ngram_size 参数进行限制。
out_captions = tokenizer.batch_decode(out, skip_special_tokens=True)：这行代码使用 tokenizer 对生成的文本进行解码，将其转换为可读的文本格式，并将结果存储在 out_captions 列表中。同时，通过 skip_special_tokens = True 参数，跳过特殊标记。
out_captions = [remove_stopwords(text) for text in out_captions]：这行代码使用 remove_stopwords 函数去除生成文本中的停用词，并将结果存储在 out_captions 列表中。
embeddings = st_model.encode(out_captions).flatten()：这行代码使用句子转换模型 st_model 对生成的文本进行编码，将其转换为嵌入向量，并通过 .flatten() 方法将嵌入向量扁平化。
sub_embeds.extend(embeddings)：这行代码将扁平化的嵌入向量 embeddings 添加到 sub_embeds 列表中。
整个代码块结束，继续下一个批次的循环迭代。

这段代码的作用是使用模型生成图像的描述文本，并将生成的子图像标识和对应的嵌入向量存储在 sub_ids 和 sub_embeds 列表中。它通过迭代处理图像批次，对每个图像生成多个描述文本，并将每个描述文本的停用词去除后进行编码，得到对应的嵌入向量。最终，sub_ids 存储了每个子图像的标识，sub_embeds 存储了每个子图像的嵌入向量。

embeddings1 = np.array(sub_embeds)

sub_embeds 应该是一个列表，其中每个元素都是一个嵌入向量。因此，可以使用 np.array 函数将 sub_embeds 转换为 NumPy 数组。

del model, tokenizer, st_model
torch.cuda.empty_cache()
gc.collect()

这段代码用于清理和释放内存资源。

del model, tokenizer, st_model：这行代码删除了之前创建的模型对象 model、分词器对象 tokenizer 和句子转换器对象 st_model，从而释放它们所占用的内存。
torch.cuda.empty_cache()：这行代码用于清空 CUDA 缓存，释放由 PyTorch 使用的 GPU 内存。
gc.collect()：这行代码使用 Python 的垃圾回收机制进行内存回收，以便尽可能地释放未使用的内存。

这些操作可以帮助释放和回收已经使用的内存资源，以便在后续的代码中重新分配和使用更多的内存。这在处理大型数据集或使用大型模型时特别有用，可以减少内存占用和提高系统性能。

六、安装并导入所有依赖项

wheels_path = "/kaggle/input/clip-interrogator-wheels-x"
clip_interrogator_whl_path = f"{wheels_path}/clip_interrogator-0.4.3-py3-none-any.whl"

根据给定的代码，wheels_path 是包含 CLIP Interrogator 的 wheels 文件的路径。假设 wheels_path 指向目录 /kaggle/input/clip-interrogator-wheels-x。

接下来，clip_interrogator_whl_path 是一个字符串变量，它通过将 wheels_path 和 CLIP Interrogator wheels 文件的文件名拼接而成。假设 CLIP Interrogator wheels 文件名为 clip_interrogator-0.4.3-py3-none-any.whl，那么 clip_interrogator_whl_path 的值将是 /kaggle/input/clip-interrogator-wheels-x/clip_interrogator-0.4.3-py3-none-any.whl。

!pip install --no-index --find-links $wheels_path $clip_interrogator_whl_path -q

这行代码使用 pip 命令安装 CLIP Interrogator，具体参数解释如下：

–no-index: 指示 pip 不从索引（如 PyPI）获取软件包，而是使用本地的 wheels 文件进行安装。
–find-links $wheels_path: 指定 pip 在给定的路径 $wheels_path 中查找软件包。这个路径是包含 CLIP Interrogator wheels 文件的目录。
$clip_interrogator_whl_path: 指定要安装的 CLIP Interrogator wheels 文件的路径。
-q: 使用静默模式进行安装，减少输出信息。

综合起来，这行代码的作用是使用指定的 wheels 文件安装 CLIP Interrogator 软件包。通过提供本地路径来安装软件包可以避免从网络下载软件包，同时可以确保安装指定版本的软件包。

!pip list | grep transformers

这行代码使用 pip list 命令列出当前环境中安装的所有 Python 包，并通过管道符 | 将输出传递给 grep 命令进行过滤。

grep transformers 表示在列表中搜索包含 “transformers” 字符串的行。grep 命令在这里用于过滤出与 “transformers” 相关的包信息。

因此，这行代码的作用是列出当前环境中安装的所有 Python 包，并仅显示包名中包含 “transformers” 的行。这样可以快速查看与 Transformers 相关的包是否已经安装。

import inspect
import importlib

from blip.models import blip
from clip_interrogator import clip_interrogator

这段代码导入了两个模块：blip.models 和 clip_interrogator。

import inspect: 导入 inspect 模块，它提供了用于获取有关活动对象的信息的函数。可以使用它来检查模块、类、函数等的属性和源代码。
from blip.models import blip: 从 blip.models 模块导入 blip。这表示在后续代码中可以直接使用 blip 来访问 blip.models 模块中定义的对象、类或函数。
from clip_interrogator import clip_interrogator: 从 clip_interrogator 模块导入 clip_interrogator。这表示在后续代码中可以直接使用 clip_interrogator 来访问 clip_interrogator 模块中定义的对象、类或函数。

注意，为了成功导入这些模块，确保它们已经安装并位于 Python 模块搜索路径中。

blip_path = inspect.getfile(blip)

fin = open(blip_path, "rt")
data = fin.read()
data = data.replace(
    "BertTokenizer.from_pretrained('bert-base-uncased')", 
    "BertTokenizer.from_pretrained('/kaggle/input/clip-interrogator-models-x/bert-base-uncased')"
)
fin.close()

fin = open(blip_path, "wt")
fin.write(data)
fin.close()

# reload module
importlib.reload(blip)

这段代码的目的是修改 blip 模块中的代码，并重新加载该模块。

blip_path = inspect.getfile(blip): 使用 inspect.getfile() 函数获取 blip 模块的文件路径，并将其赋值给 blip_path 变量。这将返回 blip 模块所在的文件路径。
fin = open(blip_path, “rt”): 使用 open() 函数打开 blip 模块文件，并以只读文本模式（“rt”）打开文件，将文件对象赋值给 fin 变量。
data = fin.read(): 使用文件对象的 read() 方法读取文件的内容，并将其赋值给 data 变量。
data = data.replace(…)：在读取的文件内容中使用 replace() 方法，将字符串 “BertTokenizer.from_pretrained(‘bert-base-uncased’)” 替换为 “BertTokenizer.from_pretrained(‘/kaggle/input/clip-interrogator-models-x/bert-base-uncased’)”。这样做的目的是修改代码中的特定部分。
fin.close(): 关闭文件对象，释放资源。
fin = open(blip_path, “wt”): 使用 open() 函数以写入文本模式（“wt”）重新打开 blip 模块文件，并将文件对象赋值给 fin 变量。
fin.write(data): 使用文件对象的 write() 方法将修改后的内容写入文件。
fin.close(): 关闭文件对象，释放资源。
importlib.reload(blip): 使用 importlib.reload() 函数重新加载 blip 模块，以使修改后的代码生效。

综合起来，这段代码的作用是通过修改 blip 模块的特定部分代码，并重新加载模块，以实现对 blip 模块的修改生效。

clip_interrogator_path = inspect.getfile(clip_interrogator.Interrogator)

fin = open(clip_interrogator_path, "rt")
data = fin.read()
data = data.replace(
    'open_clip.get_tokenizer(clip_model_name)', 
    'open_clip.get_tokenizer(config.clip_model_name.split("/", 2)[0])'
)
fin.close()

fin = open(clip_interrogator_path, "wt")
fin.write(data)
fin.close()

# reload module
importlib.reload(clip_interrogator)

这段代码的目的是修改 clip_interrogator.Interrogator 类所在模块的代码，并重新加载该模块。

clip_interrogator_path = inspect.getfile(clip_interrogator.Interrogator): 使用 inspect.getfile() 函数获取 clip_interrogator.Interrogator 类所在模块的文件路径，并将其赋值给 clip_interrogator_path 变量。这将返回 clip_interrogator.Interrogator 类所在模块的文件路径。
fin = open(clip_interrogator_path, “rt”): 使用 open() 函数打开 clip_interrogator.Interrogator 类所在模块文件，并以只读文本模式（“rt”）打开文件，将文件对象赋值给 fin 变量。
data = fin.read(): 使用文件对象的 read() 方法读取文件的内容，并将其赋值给 data 变量。
data = data.replace(…)：在读取的文件内容中使用 replace() 方法，将字符串 ‘open_clip.get_tokenizer(clip_model_name)’ 替换为 ‘open_clip.get_tokenizer(config.clip_model_name.split(“/”, 2)[0])’。这样做的目的是修改代码中的特定部分。
fin.close(): 关闭文件对象，释放资源。
fin = open(clip_interrogator_path, “wt”): 使用 open() 函数以写入文本模式（“wt”）重新打开 clip_interrogator.Interrogator 类所在模块文件，并将文件对象赋值给 fin 变量。
fin.write(data): 使用文件对象的 write() 方法将修改后的内容写入文件。
fin.close(): 关闭文件对象，释放资源。
importlib.reload(clip_interrogator): 使用 importlib.reload() 函数重新加载 clip_interrogator 模块，以使修改后的代码生效。

综合起来，这段代码的作用是通过修改 clip_interrogator.Interrogator 类所在模块的特定部分代码，并重新加载模块，以实现对该模块的修改生效。

sys.path.append('../input/sentence-transformers-222/sentence-transformers')
from sentence_transformers import SentenceTransformer, models
comp_path = Path('/kaggle/input/stable-diffusion-image-to-prompts/')

这段代码的目的是导入相关模块并设置变量。

sys.path.append(‘…/input/sentence-transformers-222/sentence-transformers’): 将路径 …/input/sentence-transformers-222/sentence-transformers 添加到系统路径中，以便导入该路径下的模块。
from sentence_transformers import SentenceTransformer, models: 导入sentence_transformers 模块中的 SentenceTransformer 和 models。
comp_path = Path(‘/kaggle/input/stable-diffusion-image-to-prompts/’): 创建一个 Path 对象，表示路径 /kaggle/input/stable-diffusion-image-to-prompts/，并将其赋值给变量 comp_path。

综合起来，这段代码的作用是导入 sentence_transformers 模块的相关部分，并设置变量 comp_path 为指定的路径。

七、设置配置

class CFG:
    device = "cuda"
    seed = 42
    embedding_length = 384
    sentence_model_path = "/kaggle/input/sentence-transformers-222/all-MiniLM-L6-v2"
    blip_model_path = "/kaggle/input/clip-interrogator-models-x/model_large_caption.pth"
    ci_clip_model_name = "ViT-H-14/laion2b_s32b_b79k"
    clip_model_name = "ViT-H-14"
    clip_model_path = "/kaggle/input/clip-interrogator-models-x/CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin"
    cache_path = "/kaggle/input/clip-interrogator-models-x"

这段代码定义了一个名为 CFG 的类，并设置了一些类变量。

device = “cuda”: 设置类变量 device 的值为字符串 “cuda”，表示使用 CUDA 设备。
seed = 42: 设置类变量 seed 的值为整数 42，表示随机种子。
embedding_length = 384: 设置类变量 embedding_length 的值为整数 384，表示嵌入长度。
sentence_model_path = “/kaggle/input/sentence-transformers-222/all-MiniLM-L6-v2”: 设置类变量 sentence_model_path 的值为字符串 “/kaggle/input/sentence-transformers-222/all-MiniLM-L6-v2”，表示句子模型的路径。
blip_model_path = “/kaggle/input/clip-interrogator-models-x/model_large_caption.pth”: 设置类变量 blip_model_path 的值为字符串 “/kaggle/input/clip-interrogator-models-x/model_large_caption.pth”，表示 BLIP 模型的路径。
ci_clip_model_name = “ViT-H-14/laion2b_s32b_b79k”: 设置类变量 ci_clip_model_name 的值为字符串 “ViT-H-14/laion2b_s32b_b79k”，表示 CI_CLIP 模型的名称。
clip_model_name = “ViT-H-14”: 设置类变量 clip_model_name 的值为字符串 “ViT-H-14”，表示 CLIP 模型的名称。
clip_model_path = “/kaggle/input/clip-interrogator-models-x/CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin”: 设置类变量 clip_model_path 的值为字符串 “/kaggle/input/clip-interrogator-models-x/CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin”，表示 CLIP 模型的路径。
cache_path = “/kaggle/input/clip-interrogator-models-x”: 设置类变量 cache_path 的值为字符串 “/kaggle/input/clip-interrogator-models-x”，表示缓存路径。

综合起来，这段代码定义了一个配置类 CFG，其中包含了一些用于配置的类变量，用于指定设备、随机种子、模型路径等信息。

八、加载示例提交

df_submission = pd.read_csv(comp_path / 'sample_submission.csv', index_col='imgId_eId')
df_submission.head()

九、Build index from images

images = os.listdir(comp_path / 'images')
imgIds = [i.split('.')[0] for i in images]

eIds = list(range(CFG.embedding_length))

imgId_eId = [
    '_'.join(map(str, i)) for i in zip(
        np.repeat(imgIds, CFG.embedding_length),
        np.tile(range(CFG.embedding_length), len(imgIds))
    )
]

assert sorted(imgId_eId) == sorted(df_submission.index)

这段代码的目的是进行图像和嵌入向量的索引匹配。

images = os.listdir(comp_path / ‘images’): 获取指定路径下 ‘images’ 文件夹中的所有文件名，并将其存储在列表 images 中。
imgIds = [i.split(‘.’)[0] for i in images]: 对 images 列表中的每个文件名执行操作，通过将文件名按 . 分割并取第一个部分，得到图像的 ID，将这些 ID 存储在列表 imgIds 中。
eIds = list(range(CFG.embedding_length)): 创建一个列表 eIds，其中包含了从 0 到 CFG.embedding_length-1 的整数。
imgId_eId = [‘_’.join(map(str, i)) for i in zip(np.repeat(imgIds, CFG.embedding_length), np.tile(range(CFG.embedding_length), len(imgIds)))]: 将 imgIds 和 range(CFG.embedding_length) 进行组合，使用下划线 ‘_’ 连接它们的元素，并将结果存储在列表 imgId_eId 中。这段代码实际上生成了一个图像 ID 和嵌入向量 ID 的组合，用于匹配图像和对应的嵌入向量。
assert sorted(imgId_eId) == sorted(df_submission.index): 检查生成的图像和嵌入向量索引是否与 df_submission 的索引一致，如果一致则通过断言，否则会引发异常。

综合起来，这段代码的作用是获取图像文件的 ID，并生成图像和嵌入向量之间的索引匹配关系，然后通过断言检查这些索引是否与给定的数据帧 df_submission 的索引一致。

st_model = SentenceTransformer(CFG.sentence_model_path)

这行代码使用 CFG.sentence_model_path 指定的路径加载 SentenceTransformer 模型，并将其实例化为 st_model。所使用的 SentenceTransformer 模型是预训练的模型，它用于将文本句子转换为对应的嵌入向量表示。通过加载预训练的模型，我们可以在后续的操作中使用该模型来计算文本的嵌入向量。

十、CLIP interrogator tool

10.1 Define CLIP interrogator config

model_config = clip_interrogator.Config(clip_model_name=CFG.ci_clip_model_name)
model_config.cache_path = CFG.cache_path

这段代码创建了一个 clip_interrogator.Config 的实例对象 model_config，并设置了其属性值。clip_interrogator.Config 是 Clip Interrogator 模型的配置类，用于存储和管理模型的各种配置参数。

在这里，通过 clip_model_name=CFG.ci_clip_model_name，将 CFG.ci_clip_model_name 的值传递给 model_config 的 clip_model_name 属性，用于指定 Clip Interrogator 模型的名称。

接下来，将 CFG.cache_path 的值赋给 model_config 的 cache_path 属性，用于指定模型的缓存路径。

通过对 model_config 的属性赋值，可以定制和配置 Clip Interrogator 模型的各种参数和选项，以满足特定的需求和设置。

10.2 Define BLIP model

configs_path = os.path.join(os.path.dirname(os.path.dirname(blip_path)), 'configs')
med_config = os.path.join(configs_path, 'med_config.json')
blip_model = blip.blip_decoder(
    pretrained=CFG.blip_model_path,
    image_size=model_config.blip_image_eval_size, 
    vit=model_config.blip_model_type, 
    med_config=med_config
)
blip_model.eval()
blip_model = blip_model.to(model_config.device)
model_config.blip_model = blip_model

这段代码主要用于加载和配置 BLIP 模型。

首先，通过 os.path.dirname(os.path.dirname(blip_path)) 获取 blip_path 的父目录的父目录，然后使用 configs_path 存储该目录下的 configs 子目录的路径。

接下来，将 configs_path 和 ‘med_config.json’ 拼接起来，得到 med_config 的路径，该文件用于配置 BLIP 模型的中介配置。

然后，使用 blip.blip_decoder() 函数加载 BLIP 模型，并通过传递参数进行配置：

pretrained = CFG.blip_model_path 指定预训练的 BLIP 模型的路径。

image_size = model_config.blip_image_eval_size 设置 BLIP 模型的图像尺寸，该值由 model_config 的 blip_image_eval_size 属性决定。

vit = model_config.blip_model_type 设置 BLIP 模型的类型，该值由 model_config 的 blip_model_type 属性决定。

med_config = med_config 指定 BLIP 模型的中介配置文件的路径。

接着，通过调用 blip_model.eval() 将 BLIP 模型设置为评估模式，并将其移动到 model_config 的指定设备（例如 “cuda”）上。

最后，将配置好的 BLIP 模型赋值给 model_config 的 blip_model 属性，以便在后续的操作中使用该模型。

10.3 Define CLIP model

clip_model = open_clip.create_model(CFG.clip_model_name, precision='fp16' if model_config.device == 'cuda' else 'fp32')
open_clip.load_checkpoint(clip_model, CFG.clip_model_path)
clip_model.to(model_config.device).eval()
model_config.clip_model = clip_model

这段代码主要用于加载和配置 OpenAI CLIP 模型。

首先，通过调用 open_clip.create_model() 函数创建 CLIP 模型，并传递参数进行配置：

CFG.clip_model_name 指定 CLIP 模型的名称。
precision = ‘fp16’ if model_config.device == ‘cuda’ else ‘fp32’ 根据 model_config 的设备类型（如果是 “cuda”），选择使用浮点数精度为 “fp16”，否则选择 “fp32”。

接着，使用 open_clip.load_checkpoint() 函数加载预训练的 CLIP 模型的权重，并传递参数：

clip_model 是之前创建的 CLIP 模型实例。
CFG.clip_model_path 指定预训练的 CLIP 模型的权重文件路径。

然后，通过调用 clip_model.to(model_config.device).eval() 将 CLIP 模型移动到 model_config 指定的设备上（例如 “cuda”），并将其设置为评估模式。

最后，将配置好的 CLIP 模型赋值给 model_config 的 clip_model 属性，以便在后续的操作中使用该模型。

clip_preprocess = open_clip.image_transform(
    clip_model.visual.image_size,
    is_train = False,
    mean = getattr(clip_model.visual, 'image_mean', None),
    std = getattr(clip_model.visual, 'image_std', None),
)
model_config.clip_preprocess = clip_preprocess

这段代码用于定义 CLIP 模型的图像预处理函数。

首先，通过调用 open_clip.image_transform() 函数创建图像预处理函数 clip_preprocess，并传递参数：

clip_model.visual.image_size 指定 CLIP 模型期望的图像尺寸。
is_train = False 表示图像预处理函数将在推断阶段使用。
mean = getattr(clip_model.visual, ‘image_mean’, None) 获取 CLIP 模型的图像均值。如果模型定义了 image_mean 属性，则使用该值作为图像均值；否则，传递 None。
std = getattr(clip_model.visual, ‘image_std’, None) 获取 CLIP 模型的图像标准差。如果模型定义了 image_std 属性，则使用该值作为图像标准差；否则，传递 None。

最后，将定义好的图像预处理函数赋值给 model_config 的 clip_preprocess 属性，以便在后续的操作中使用该预处理函数对图像进行处理。

10.4 Create CLIP interrogator object

ci = clip_interrogator.Interrogator(model_config)

这段代码创建了一个 clip_interrogator.Interrogator 对象 ci，并传递了之前定义的 model_config 对象作为参数。

clip_interrogator.Interrogator 是用于与 CLIP 模型进行交互和生成文本描述的工具类。通过将 model_config 作为参数传递给 Interrogator，可以配置该工具类的行为，包括使用的 CLIP 模型、BLiP 模型、缓存路径等。

通过创建 ci 对象，可以使用其中的方法来生成图像的描述、计算图像与文本之间的相似性等操作。

10.5 Define interrogate function

10.5.1 Get labels embeddings

原始 CLIP Interrogator 使用 image_features 和 text_embeds 矩阵乘法来细化相应图像和文本标签之间的相似性。但我发现使用余弦相似度要快得多，而且得到的分数几乎相同。所以考虑到这一点。

cos = torch.nn.CosineSimilarity(dim=1)

mediums_features_array = torch.stack([torch.from_numpy(t) for t in ci.mediums.embeds]).to(ci.device)
movements_features_array = torch.stack([torch.from_numpy(t) for t in ci.movements.embeds]).to(ci.device)
flavors_features_array = torch.stack([torch.from_numpy(t) for t in ci.flavors.embeds]).to(ci.device)

这段代码定义了一个余弦相似性计算器 cos，使用 torch.nn.CosineSimilarity 类进行初始化。这个计算器可以用于计算两个张量之间的余弦相似度。

接下来，通过将 ci.mediums.embeds、ci.movements.embeds 和 ci.flavors.embeds 转换为张量，并将它们堆叠在一起，创建了三个特征数组 mediums_features_array、movements_features_array 和 flavors_features_array。

具体地：

torch.from_numpy(t) 将 ci.mediums.embeds、ci.movements.embeds 和 ci.flavors.embeds 中的每个 NumPy 数组 t 转换为 PyTorch 张量。
torch.stack() 将转换后的张量堆叠在一起，创建一个包含所有特征的张量数组。
.to(ci.device) 将张量移动到指定的设备（在 model_config 中配置的设备，通常是 CUDA 设备）。

最终，mediums_features_array、movements_features_array 和 flavors_features_array 是三个包含特征向量的 PyTorch 张量，可以在后续的计算中使用。

10.5.2 Create main interrogation function

def interrogate(image: Image) -> str:
    caption = ci.generate_caption(image)
    image_features = ci.image_to_features(image)
    
    medium = [ci.mediums.labels[i] for i in cos(image_features, mediums_features_array).topk(1).indices][0]
    movement = [ci.movements.labels[i] for i in cos(image_features, movements_features_array).topk(1).indices][0]
    flaves = ", ".join([ci.flavors.labels[i] for i in cos(image_features, flavors_features_array).topk(3).indices])

    if caption.startswith(medium):
        prompt = f"{caption}, {movement}, {flaves}"
    else:
        prompt = f"{caption}, {medium}, {movement}, {flaves}"

    return clip_interrogator._truncate_to_fit(prompt, ci.tokenize)

这段代码定义了一个名为 interrogate 的函数，它接受一张图像作为输入，并返回一个字符串。

函数中的步骤如下：

使用 ci.generate_caption(image) 生成图像的标题（caption）。
使用 ci.image_to_features(image) 将图像转换为特征向量。
通过计算图像特征与 mediums_features_array、movements_features_array 和 flavors_features_array 之间的余弦相似度，找到与图像特征最相似的介质（medium）、运动（movement）和口味（flaves）标签。具体步骤如下：

对于介质（medium），通过计算图像特征与 mediums_features_array 的余弦相似度，并选择相似度最高的标签；

对于运动（movement），通过计算图像特征与 movements_features_array 的余弦相似度，并选择相似度最高的标签；

对于口味（flaves），通过计算图像特征与 flavors_features_array 的余弦相似度，并选择相似度最高的三个标签，并用逗号连接它们。

根据生成的标题（caption）的开头是否与介质（medium）相同，构建一个提示（prompt）字符串。如果标题的开头与介质相同，则提示字符串为 “{caption}, {movement}, {flaves}”；否则，提示字符串为 “{caption}, {medium}, {movement}, {flaves}”。
最后，将提示字符串截断为合适的长度，并返回结果。

总体来说，该函数利用图像特征与预定义的介质、运动和特征之间的余弦相似度，结合生成的图像标题，构建一个提示字符串，用于描述图像的内容。

10.6 Extract promt from images

prompts = []

images_path = "../input/stable-diffusion-image-to-prompts/images/"
for image_name in images:
    img = Image.open(images_path + image_name).convert("RGB")

    generated = interrogate(img)
    
    prompts.append(generated)

10.6.1 Check the result

def add_text_limiters(text: str) -> str:
    return " ".join([
        word + "\n" if i % 15 == 0 else word 
        for i, word in enumerate(text.split(" "), start=1)
    ])

def plot_image(image: np.ndarray, original_prompt: str, generated_prompt: str) -> None:
    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    plt.annotate(
        "Original prompt:\n" + add_text_limiters(original_prompt) + "\n\nGenerated prompt:\n" + add_text_limiters(generated_prompt), 
        xy=(1.05, 0.5), xycoords='axes fraction', ha='left', va='center', 
        fontsize=16, rotation=0, color="#104a6e"
    )

这段代码定义了两个函数：add_text_limiters 和 plot_image。

add_text_limiters(text: str) -> str 函数接受一个字符串作为输入，并在每15个单词后添加换行符，以便在绘制图像时更好地展示文本。具体步骤如下：

将输入的字符串按空格分割成单词列表。

使用列表推导式遍历单词列表，对于每个单词，如果其在单词列表中的索引模15等于0，则在单词后添加换行符，否则保持不变。

最后，使用空格将处理后的单词列表重新连接成字符串，并返回结果。

plot_image(image: np.ndarray, original_prompt: str, generated_prompt: str) -> None 函数用于绘制图像，并在图像旁边添加原始提示（original_prompt）和生成提示（generated_prompt）。具体步骤如下：

创建一个大小为（10, 10）的图像。

使用 plt.imshow 绘制输入的图像。

使用 plt.annotate 在图像旁边添加文本注释，注释内容包括原始提示和生成提示。其中，add_text_limiters 函数用于对提示进行处理，以确保在每15个单词后添加换行符。

设置注释的位置、字体大小、颜色等参数。

最后，无返回值，图像和注释会显示在当前的图像窗口中。

十一、Create a sample submission with a constant prompt prediction

11.1 Encode prompts

embeddings2 = st_model.encode(prompts).flatten()

十二、Release GPU resource

del ci
del blip_model, clip_model
del st_model
torch.cuda.empty_cache()
gc.collect()

十三、Ensemble(1)

embeddings12 = ratio_OFA * embeddings1 + ratio_CLIP_Interrogator * embeddings2

del embeddings1
del embeddings2
gc.collect()

十四、ViT-B-16

import numpy as np
import pandas as pd
from pathlib import Path
from PIL import Image
from tqdm.notebook import tqdm
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import timm
from sklearn.preprocessing import normalize

class CFG:
    model_path = '/kaggle/input/stable-diffusion-vit-baseline-train/vit_base_patch16_224.pth'
    model_name = 'vit_base_patch16_224'
    input_size = 224
    batch_size = 64

class DiffusionTestDataset(Dataset):
    def __init__(self, images, transform):
        self.images = images
        self.transform = transform
    
    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image = Image.open(self.images[idx])
        image = self.transform(image)
        return image

def predict(
    images,
    model_path,
    model_name,
    input_size,
    batch_size
):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    transform = transforms.Compose([
        transforms.Resize(input_size),
        transforms.RandomHorizontalFlip(p=0.5),
#         transforms.RandomRotation(degrees=10),

        # transforms.RandomVerticalFlip(p=0.5),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])
    dataset = DiffusionTestDataset(images, transform)
    dataloader = DataLoader(
        dataset=dataset,
        shuffle=False,
        batch_size=batch_size,
        pin_memory=True,
        num_workers=2,
        drop_last=False
    )

    model = timm.create_model(
        model_name,
        pretrained=False,
        num_classes=384
    )
    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    
    tta_preds = None
    for _ in range(2):
        preds = []
        for X in tqdm(dataloader, leave=False):
            X = X.to(device)

            with torch.no_grad():
                X_out = model(X).cpu().numpy()
                # L2 normalize -- Start
                X_out = X_out / ( np.abs(X_out).max(axis=-1, keepdims=True) + 0.0000001)  # To avoid to overflow at normalize()
                X_out = normalize( X_out )
                # L2 normalize -- End
                preds.append(X_out)
    if tta_preds is None:
            tta_preds = np.vstack(preds).flatten()
        else:
            tta_preds += np.vstack(preds).flatten()
    
    return tta_preds / 2

该代码定义了一个predict函数，用于对输入的图像进行预测。

函数参数包括：

images：输入的图像数据。
model_path：模型文件的路径。
model_name：模型的名称。
input_size：输入图像的尺寸。
batch_size：批处理的大小。

函数的主要步骤如下：

检查可用的设备，并将设备设置为CUDA（GPU）或 CPU。
定义数据转换的操作，包括图像的大小调整、随机水平翻转和张量转换等。
创建 DiffusionTestDataset 数据集对象，将图像数据和转换操作传递给数据集。
创建 DataLoader 数据加载器，用于批量加载数据，并设置相关参数，如批处理大小、是否打乱数据等。
创建指定模型名称和输出类别数的模型。
加载预训练模型的权重。
将模型移动到指定设备上，并设置为评估模式。
进行测试时间数据增强（TTA）：

对于每个图像批次，将图像数据传递给模型进行预测。

对预测结果进行 L2 归一化处理。

将预测结果添加到列表中。

计算 TTA 的平均预测结果。

最后，返回经过 TTA 处理后的预测结果。

images = list(Path('/kaggle/input/stable-diffusion-image-to-prompts/images').glob('*.png'))
imgIds = [i.stem for i in images]
EMBEDDING_LENGTH = 384
imgId_eId = [
    '_'.join(map(str, i)) for i in zip(
        np.repeat(imgIds, EMBEDDING_LENGTH),
        np.tile(range(EMBEDDING_LENGTH), len(imgIds)))]

embeddings3 = predict(images, CFG.model_path, CFG.model_name, CFG.input_size, CFG.batch_size)

该代码从指定路径（/kaggle/input/stable-diffusion-image-to-prompts/images）中获取所有以 .png 扩展名结尾的图像文件，并将它们存储在 images 列表中。然后，通过遍历 images 列表，获取每个图像文件的文件名（不包含扩展名），并将它们存储在 imgIds 列表中。

接下来，定义了一个常量 EMBEDDING_LENGTH，表示嵌入长度（embedding length），设定为 384。

然后，通过将 imgIds 列表中的图像文件名和 range(EMBEDDING_LENGTH) 组合，生成一个新的列表 imgId_eId，其中每个元素是由图像文件名和对应的嵌入索引组成，用下划线连接起来。

最后，调用 predict 函数，传入图像列表 images、模型路径 CFG.model_path、模型名称 CFG.model_name、输入尺寸 CFG.input_size 和批处理大小 CFG.batch_size，以获取预测的嵌入向量。将结果存储在 embeddings3 变量中。

十五、Ensemble(2)

embeddings = embeddings12 + ratio_ViT_B_16 * embeddings3

submission = pd.DataFrame(
    index=imgId_eId,
    data=embeddings,
    columns=['val']
).rename_axis('imgId_eId')

【Kaggle】Stable Diffusion - Image to Prompts竞赛代码初步理解

文章目录

一、前言

二、导包

三、加载预训练的 OFA 模型

四、模型EDA

五、Inference

六、安装并导入所有依赖项

七、设置配置

八、加载示例提交

九、Build index from images

十、CLIP interrogator tool

10.1 Define CLIP interrogator config

10.2 Define BLIP model

10.3 Define CLIP model

10.4 Create CLIP interrogator object

10.5 Define interrogate function

10.5.1 Get labels embeddings

10.5.2 Create main interrogation function

10.6 Extract promt from images

10.6.1 Check the result

十一、Create a sample submission with a constant prompt prediction

11.1 Encode prompts

十二、Release GPU resource

十三、Ensemble(1)

十四、ViT-B-16

十五、Ensemble(2)