[机器学习] 图像字幕模型：输入图像，输出文本

介绍： 本次实验主要依据此参考文献进行，模型输入为270*200大小的公式图片，输出为图片中公式对应的latex码。

数据处理

1、提取交换图tex码

训练数据来源于stacks-project（一本书），下载其tex代码源码。通过python提取其中的\xymatrix{ }块，提取部分的关键代码取下：

import re
# 定义正则表达式模式
pattern = r'\$\$\n*\\xymatrix{.*?}\n*\$\$'
# 遍历目录中的文件
for filename in os.listdir(directory):
    with open(filepath, 'r') as file:
        # 读取文件内容
        contents = file.read()
        # 使用正则表达式搜索xymatrix块
        matches = re.findall(pattern, contents, re.DOTALL)

得到共3884个文件，论文中为4345个，原因是使用代码进行提取时的匹配方式导致部分latex代码块提取不到。考虑到时间因素，且3884个样本也足够继续模型的训练和验证，故不再继续提取剩余代码块。

2、交换图编码为pdf文件

下载tex编译器，配置环境变量。安装所需宏包。

将提取出的.tex文件批量进行编码，需要注意的是编码后会得到pdf文件，需要进一步转为png图片。使用以下cmd命令可以对单个.tex文件进行编译：

pdflatex -interaction nonstopmode -output-directory E:/TeXworkspace/diagramPdf E:/TeXworkspace/diagramCode/0_adequate.tex

对目录中的每个文件，可以使用subprocess库构建命令以将.tex批量编码为.pdf文件。

3、转为图片并裁剪大小

使用pdf2image库可以将pdf文件转换为图片，但是需要先安装poppler并配置环境变量。

from pdf2image import convert_from_path
images = convert_from_path(pdf_file)

由于生成的交换图在图片（1654*2339）中所占大小很小，不能直接对转换的图片进行缩放。考虑使用canny边缘检测算法定位需要裁剪的地方，使用canny对其进行边缘检测，对最大框附近进行裁剪，最后统一大小为270*200。

# 将图像转换为灰度图
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# 使用Canny边缘检测算法检测图像边缘
edges = cv2.Canny(gray, threshold1=50, threshold2=100)
# 查找轮廓
contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# 找到最大的轮廓
max_contour = max(contours, key=cv2.contourArea)
# 获取轮廓的边界框
x, y, width, height = cv2.boundingRect(max_contour)
# 裁剪图像
cropped_image = image[y-250:y + height + 225, x-300:x + width+350]
cropped_image = cv2.resize(cropped_image, (270, 200))

改进思路： 使用cv2.resize()方法后图像的清晰度变得很低，如果能够使用其他更好的方法修改裁剪后的图片大小，使得到的图像更清晰，应该能够进一步提升最终模型的准确度。

4、词元划分与编码

由于模型不能直接识别字符，故需对tex码中的单词进行提取并编码。提取的方式有两种，一种是将单个字符作为提取对象，一种是进行词元（tokenize）提取。这里采取词元提取的方法对tex码进行处理。需要先将tex码中的如,[{^_\']}()|~等符号与紧挨着的字符用空格分开，然后对所有提取出的tex码中出现的词元进行统计。对每个词元用数字进行编码（即分配索引）。此外，还需要加入特殊词元如、、、分别代表未知、开始、占位、结束。

由于篇幅问题，这里主要介绍思路，不贴代码了。可以参考这个链接，写得还是比较详细的：参考链接

大致思路如下：

1）按照以下步骤进行操作构建图像字幕的词表：

1. 遍历存放图像字幕的目录，读取每个字幕文件。
2. 对于每个字幕文件，读取其中的文本内容。
3. 对文本内容进行分词，将文本分割成单词或子词。

4. 统计每个单词或子词的出现频率，并记录到词频字典中。
5. 根据词频字典，选择出现频率大于阈值的单词或子词作为词汇表（词表）。
6. 使用字典数据结构为词汇表中的每个单词或子词分配一个唯一的索引，构建单词到索引的映射关系。
7. 添加特殊的标记符号，如起始标记（）和结束标记（），并将其加入词汇表中，并为它们分配相应的索引。
8. 添加其他特殊标记符号，如未知单词标记（）或填充标记（）。
9. 将构建好的词汇表保存到文件中，以备后续使用。

2）对数据进行词元化，使得文本序列中每个词元(tokenize)要么是一个词，要么是一个标点符号。

词元划分示例：
[‘\xymatrix’, ‘{’, ‘enter’, ‘A’, ‘^’, ‘{’, ‘\oplus’, ‘n’, ‘}’, ‘\ar’, ‘[’, ‘r’, ‘]’, ‘\ar’, ‘[’, ‘d’, ‘]’, ‘', ‘{’, ‘m’, '’, ‘1’, ‘,’, ‘\ldots’, ‘,’, ‘m’, ‘_’, ‘n’, ‘}’, ‘&’, ‘N’, “'”, ‘\ar’, ‘[’, ‘d’, ‘]’, ‘\\’, ‘enter’, ‘M’, ‘\ar’, ‘[’, ‘r’, ‘]’, ‘&’, ‘N’, ‘enter’, ‘}’]

3）读取所有的tex文件，对其中的每个词元进行计数，构建vocab类，这个类为词元添加索引并支持查询词表长度、通过索引查询词元、词元查询索引。

词元列表（word_map）前10项示例：

[(‘’, 0), (‘’, 1), (‘’, 2), (‘’, 3), (‘{’, 4), (‘}’, 5), (‘[’, 6), (‘]’, 7), (‘enter’, 8), (‘\ar’, 9)]

4）对词元进行编码后，将tex码中的词元转化为其对应的索引，并修改数据类型为tensor。将tex码进行填充截断，根据分配的索引编码后转为长度125的tensor类型。

tex码转换后示例：

tensor([ 2, 21, 4, 8, 29, 13, 4, 98, 37, 5, 9, 6, 12, 7, 9, 6, 14, 7, 10, 4, 67, 10, 23, 22, 70, 22, 67, 10, 37, 5, 11, 88, 17, 9, 6, 14, 7, 19, 8, 49, 9, 6, 12, 7, 11, 88, 8, 5, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

5、数据集

经过上述操作后我们得到了latex的编码，和其对应的公式图片。接下来可以将其封装成数据集方便后续的训练。模型的输入为270*200*1大小的图像，输出为长度125的序列。

preprocess = transforms.Compose([
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize(mean=[0.5], std=[0.5]),  # 标准化
])

# 生成dataset类
class ImageTextDataset(Dataset):
    def __init__(self, image_folder):
        self.image_folder = image_folder
        with open('dataprocess/tensor_all.pkl', 'rb') as f:
            self.text = pickle.load(f)

    def __len__(self):
        return len(self.text)

    def __getitem__(self, index):
        filenames = os.listdir(self.image_folder)

        image_path = os.path.join(self.image_folder, filenames[index])
        img = Image.open(image_path).convert('L')  # 转为单通道图片
        # 正则化
        input_tensor = preprocess(img)
        text_tensor = self.text[index]
        return input_tensor, text_tensor

模型

由于本次需要实现的模型输入为图像、输出为序列，故参考图像字幕生成模型的实现。

图像字幕模型参考

1、图像编码

模型分为编码器和解码器两个部分，其中编码器既可以使用预训练模型进行编码，也可以自定义编码器。

一、使用预训练模型如resnet、vgg进行编码

由于已经有大量对图像编码的CNN模型，故可以使用已有的预训练模型对图像进行编码，而不必重新训练一个编码器。由于CNN模型中的最后一层或两层一般是是线性层和softmax激活层，故剥离预训练模型的最后两层。

使用预训练模型属于迁移学习。迁移学习指通过在新模型中使用现有模型的部分来借用现有模型；它几乎总是好于从零开始训练新模型（即什么都不知道）。

常用的预训练模型有：

• ResNet-101：将图像编码为(2048，14，14)大小的张量

• VGGnet

可以提前使用预训练模型对图片进行编码，并将编码得到的特征保存到文件中，这样在使用解码器时只需要直接加载编码后的特征就可以了。但是使用预训练模型时我遇到了OOM的问题，可能是这些模型都太大了。

# 加载预训练的ResNet模型
resnet = models.resnet18(pretrained=True)
# 最后一层为分类层，删除该层
resnet = torch.nn.Sequential(*list(resnet.children())[:-2])

送到预训练编码器的图像必须修改为维度预训练模型指定维度的Float 张量，并通过平均值和标准偏差进行归一化。

# 三通道
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
# 单通道
transforms.Normalize(mean=[0.5], std=[0.5])

参考链接

二、自定义卷积神经网络（本次使用）

使用自定义的CNN对其进行编码，编码端结构如下所示。

# 对输入图像进行编码
class Encoder(nn.Module):
    def __init__(self):  # 输入为270*200大小的图像
        super(Encoder, self).__init__()
        self.conv = Sequential(                             # (16,1,270,200)
            Conv2d(1, 6, 5, stride=1),                      # (16,6,131,96)
            ReLU(),
            MaxPool2d(2),                                   # (16,6,128,93)
            BatchNorm2d(6),

            Conv2d(6, 16, 5, 1),                            # (16,16,57,42)
            ReLU(),
            MaxPool2d(2),                                   # (16,16,54,39)
            BatchNorm2d(16),

            Conv2d(16, 20, 5, 1),                           # (16,16,57,42)
            ReLU(),
            MaxPool2d(2),                                   # (16,16,54,39)
            BatchNorm2d(20),
        )
        self.fc = Sequential(
            Linear(12600, 10000),
            ReLU(),
            Linear(10000, 1000),                               # [16,1000]
            ReLU()
        )

    # 前向过程
    def forward(self, x):
        # X的形状应该为(batch_size,num_steps,embed_size)
        x = self.conv(x)
        x = x.view(x.size(0), -1)  # 展平
        x = self.fc(x)
        return x

2、特征解码

embedding嵌入层用于将离散的整数序列（如单词）映射到连续的低维向量表示。GRU是一种循环神经网络（RNN）的变种，用于处理序列数据，通过门控机制来捕捉和传递序列中的重要信息。

解码器使用了一个GRU单元来解码编码器的输出，并通过全连接层将其映射到词表空间。

def __init__(self, input_size=1000, hidden_size=1000, vocab_size=1491, device = “cuda”):
    super(Decoder, self).__init__()
    self.hidden_size = hidden_size
    # 定义嵌入层
    self.embedding = nn.Embedding(vocab_size, hidden_size)
    # 定义GRU单元
    self.gru = nn.GRU(input_size, hidden_size)
    # 定义全连接层
    self.fc = nn.Linear(hidden_size, vocab_size)

解码器前向过程：

def forward(self, encoder_output, captions):  # encoder_output(16,1000) caption(16,125)
    # 将编码器输出进行形状变换
    batch_size = encoder_output.size(0)
    caption_length = captions.size(1)  # 125
    embedded = self.embedding(captions).permute(1, 0, 2)  # (16,125,1000)
    hidden = self.init_hidden(batch_size)  # 1，16，1000
    # hidden = encoder_output.unsqueeze(0)  # (1,16,1000)
    outputs = []
    for t in range(caption_length):  # t :时间步，从0到125
        # 将编码器输出通过GRU单元进行解码
        # embedded[:, t, :].unsqueeze(1)表示时间步t上所有样本的嵌入向量
        output, hidden = self.gru(embedded[t, :, :].unsqueeze(1), hidden)  # (16, 1, 1000)  (1,16,1000)
        # 将GRU的输出通过全连接层进行映射到词表空间
        output = self.fc(output.squeeze(1))
        outputs.append(output.unsqueeze(1))
        outputs = torch.cat(outputs, dim=1)
        return outputs  # 输出词元的分布概率

其中的隐含向量hidden则初始化为编码端生成的embedding。

解码器是循环推理、逐个单词生成结果的。最开始，将编码器提取的特征以及一个传给解码器，解码器预期会输出一个单词。然后将预测的第一个单词输入给解码器，会再预测出下一个单词，以此类推直到预测出句子终止符完成预测。

3、模型训练

对于模型的输出 output，其形状为 [batch_size, seq_len, vocab_size]，其中 batch_size 是批次大小，seq_len 是序列长度，vocab_size 是词汇表大小。

对于目标tex码 captions，其形状为 [batch_size, seq_len]，其中 batch_size 和 seq_len 与模型输出相对应。

为了计算交叉熵损失，我们需要将模型输出和目标字幕进行展平，即变形为二维张量。这可以通过 view 方法实现。view(-1) 将张量展平为一维，并保持原有数据顺序。

在代码中，output.view(-1, output.size(2)) 将模型输出展平为二维张量，形状为 [batch_size * seq_len, vocab_size]，而 captions.view(-1) 将目标字幕展平为一维张量，形状为 [batch_size * seq_len]。

这样，我们就可以将展平后的模型输出和目标字幕传递给交叉熵损失函数进行计算，确保维度的匹配。

# 划分训练集
lengths = [int(len(dataset) * 0.9), int(len(dataset) * 0.05),
           len(dataset) - int(len(dataset) * 0.9) - int(len(dataset) * 0.05)]
train_set, val_set, test_set = random_split(dataset, lengths)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True)

# 定义损失函数和优化器
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 迭代训练数据
for epoch in range(num_epochs):
    for i,(images, captions) in enumerate(train_loader):
        # 前向传播
        output = model(images, captions)  # 依据编码端和解码端构建模型model
        # 计算损失
        loss = loss_fn(output.view(-1, output.size(2)), captions.view(-1))
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

4、模型测试

将测试集中的某图像输入模型，得到如下输出。解编码后得到tex码，编译后得到对应的公式图像。

不足之处：模型评估不足。对于自然语言处理（NLP）方向的模型，可以通过计算BLEU分数对模型进行评估。