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

最新推荐文章于 2024-07-13 08:48:44 发布

旅途中的宽~

最新推荐文章于 2024-07-13 08:48:44 发布

阅读量1.1k

点赞数 1

本文链接：https://blog.csdn.net/wzk4869/article/details/130542913

版权

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

摘要由CSDN通过智能技术生成

文章目录

一、前言
二、导包
三、加载预训练的 OFA 模型
四、模型EDA
五、Inference
六、安装并导入所有依赖项
七、设置配置
八、加载示例提交
九、Build index from images
十、CLIP interrogator tool
十一、Create a sample submission with a constant prompt prediction
- 11.1 Encode prompts
十二、Release GPU resource
十三、Ensemble(1)
十四、ViT-B-16
十五、Ensemble(2)

一、前言

此次代码集成了 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.