处理Flickr8k数据集

处理Flickr8k数据集

数据集下载链接：https://www.kaggle.com/datasets/adityajn105/flickr8k?resource=download
包括8091张图像，1个文本文件包含对每张图像的5个描述。

目的：只是借此数据集记录一下文本处理基础方法，如果开展相应的项目，应该参考官方的数据集处理和数据集分割方式。

流程如下：

0、数据集情况

文本存储在 captions.txt 中，记录方式如下： 第一行给出的表示方式为 “image, caption”。

图像数据：

1、读取文本并将每张图像的caption group到一起

import os
import json
import random
from collections import defaultdict, Counter
from PIL import Image
from matplotlib import pyplot as plt
import string

### 读取文件为list
def read_txt_as_list(txt_path):
    f = open(txt_path)
    f_list = []
    for line in f:
        f_list.append(line.strip())
    return f_list

caption_path = './flickr8k/captions.txt'
caption_txt = read_txt_as_list(caption_path)

### 创建一个字典，key为img_name，value为5个对应的caption
flickr8k_dict = defaultdict(list)   # 用defaultdict的好处是值默认初始化
for i in range(1, len(caption_txt)): # 从1开始因为第一行为标识表示
    comma_pos = caption_txt[i].find(',')  # 第一个逗号分隔name和caption
    img_name = caption_txt[i][:comma_pos]
    img_caption = caption_txt[i][comma_pos+1:]
    flickr8k_dict[img_name].append(img_caption)

2、随机分割数据集

### 获得所有的图像名称即keys
all_img_names = list(flickr8k_dict.keys())
### shuffle
random.shuffle(all_img_names)
### 按照一定比例划分train val test set
train_img_names = all_img_names[:int(len(all_img_names) * 0.7)]
val_img_names = all_img_names[int(len(all_img_names) * 0.7):int(len(all_img_names) * 0.9)]
test_img_names = all_img_names[int(len(all_img_names) * 0.9):]

3、构造词典

（1）将句子分割成单词list，单词都小写，且忽略标点。

### 定义提取token的函数
def get_token_from_sent(sentence):
    words = sentence.translate(str.maketrans('', '', string.punctuation)).split()   # 前两个参数没有定义映射表，第三个参数将标点映射到None；然后分割单词
    for i in range(len(words)):  # 将单词都存为小写
        words[i] = words[i].lower()
    return words

（2）用训练集构造词典（只将至少出现min_word_count次的次加入词典）

def build_vocab(data_dict, train_keys, min_word_count=5):
    """
    :param data_dict: key=imgname value=caption的字典
    :param train_keys: train_img_names list
    :param min_word_count: 只将出现了多次的词加入词典
    :return: vocab 词典
    """
    vocab = Counter()
    for i in range(len(train_keys)):
        img_name = train_keys[i]
        sents = data_dict[img_name]
        for j in range(len(sents)):
            words = get_token_from_sent(sents[j])
            vocab.update(words)
    words = [w for w in vocab.keys() if vocab[w] >= min_word_count]  # 保存出现了最小次数以上的单词
    vocab = {k: v + 1 for v, k in enumerate(words)}  # 为每个词分配一个编号
    vocab['<pad>'] = 0  # 增加占位标识符<pad>
    vocab['<unk>'] = len(vocab)  # 未登录词标识符<unk>
    vocab['<start>'] = len(vocab) # 句子首尾标识符<start>和<end>
    vocab['<end>'] = len(vocab)
    return vocab

train_vocab = build_vocab(flickr8k_dict, train_img_names)
# 存储词典
# with open('vocab.json', 'w') as fw:
#     json.dump(train_vocab, fw)

4、对数据集中的样本的描述用词典编码

结构化数据，关键在于“多则截取，少则补全”。

def get_imgname_captioncode(data_dict, set_keys, vocab, captions_per_image=5, max_len=30):
    """
    :param data_dict: key=imgname value=caption的字典
    :param set_keys: 数据集的keys
    :param captions_per_image: 每张图像包含的主题句子描述的数目
    :param max_len: 文本描述包含的最大单词数，如果文本描述超过该值，则截断
    :return: ret_imgname, ret_captioncode  图像名称list，文本编码list
    """
    ### 初始化
    ret_imgname = []
    ret_captioncode = []
    ### 对每个样本执行
    for i in range(len(set_keys)):  
        img_name = set_keys[i]  # 图像名字
        sents = data_dict[img_name]  # 相应的文本描述们
        ### 将sents转为words tokens
        tokens = []
        for j in range(len(sents)):  # 提取样本的每个描述为token
            tok = get_token_from_sent(sents[j])
            if(len(tok) <= max_len):  # 如果token数目过多，则截断
                tokens.append(tok)
            else:
                tokens.append(tok[:max_len])
        ### 如果该图片对应的描述数量不足，则补足
        if len(tokens) < captions_per_image:
            captions = tokens + [random.choice(tokens) for _ in range(captions_per_image - len(tokens))]
        ### 如果该图片对应的描述数量多了，则随机采样
        else:
            captions = random.sample(tokens, k=captions_per_image)
        assert len(captions) == captions_per_image
        ### 用词典对文本进行编码
        enc_captions = []
        for k, c in enumerate(captions):
            enc_c = [vocab['<start>']] + [vocab.get(word, vocab['<unk>']) for word in c] + [vocab['<end>']]  # list
            enc_c = enc_c + [vocab['<pad>']] * (max_len + 2 - len(enc_c))  # 不足的位置用pad补全
            enc_captions.append(enc_c)
        # 存储到列表中
        ret_imgname.append(img_name)
        ret_captioncode.append(enc_captions)
    return ret_imgname, ret_captioncode

# 对训练集进行：
train_ret_names, train_ret_captioncode = get_imgname_captioncode(flickr8k_dict, train_img_names, train_vocab)

检查一下是否正确：
选取图片：1813266419_08bf66fe98.jpg
其中一个描述：‘A child in a striped shirt gleefully plays among some water fountains .’

sent = 'A child in a striped shirt gleefully plays among some water fountains .'
word = get_token_from_sent(sent)
# ['a', 'child', 'in', 'a', 'striped', 'shirt', 'gleefully', 'plays', 'among', 'some', 'water', 'fountains']
word_code = [train_vocab.get(w, train_vocab['<unk>']) for w in word]
# [1, 2, 3, 1, 4, 5, 2208, 6, 7, 8, 9, 10]

查找编码好的列表中可得（因为shuffle过和random sample过，所以不是原来的顺序了，只能手动查找）

# [2209, 1, 2, 3, 1, 4, 5, 2208, 6, 7, 8, 9, 10, 2210, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

对比可知，二者是匹配的，2209对应’<start>‘，2210对应’<end>‘，0对应’<pad>'。

参考：飞桨多模态大模型学习课程https://aistudio.baidu.com/projectdetail/7758527