Upload 8 files

data_script/flintstones_hdf5.py

@@ -0,0 +1,51 @@
+import argparse
+import json
+import os
+import pickle
+
+import cv2
+import h5py
+import numpy as np
+from tqdm import tqdm
+
+
+def main(args):
+    splits = json.load(open(os.path.join(args.data_dir, 'train-val-test_split.json'), 'r'))
+    train_ids, val_ids, test_ids = splits["train"], splits["val"], splits["test"]
+    followings = pickle.load(open(os.path.join(args.data_dir, 'following_cache4.pkl'), 'rb'))
+    annotations = json.load(open(os.path.join(args.data_dir, 'flintstones_annotations_v1-0.json')))
+    descriptions = dict()
+    for sample in annotations:
+        descriptions[sample["globalID"]] = sample["description"]
+
+    f = h5py.File(args.save_path, "w")
+    for subset, ids in {'train': train_ids, 'val': val_ids, 'test': test_ids}.items():
+        ids = [i for i in ids if i in followings and len(followings[i]) == 4]
+        length = len(ids)
+
+        group = f.create_group(subset)
+        images = list()
+        for i in range(5):
+            images.append(
+                group.create_dataset('image{}'.format(i), (length,), dtype=h5py.vlen_dtype(np.dtype('uint8'))))
+        text = group.create_dataset('text', (length,), dtype=h5py.string_dtype(encoding='utf-8'))
+        for i, item in enumerate(tqdm(ids, leave=True, desc="saveh5")):
+            globalIDs = [item] + followings[item]
+            txt = list()
+            for j, globalID in enumerate(globalIDs):
+                img = np.load(os.path.join(args.data_dir, 'video_frames_sampled', '{}.npy'.format(globalID)))
+                img = np.concatenate(img, axis=0).astype(np.uint8)
+                img = cv2.imencode('.png', img)[1].tobytes()
+                img = np.frombuffer(img, np.uint8)
+                images[j][i] = img
+                txt.append(descriptions[globalID])
+            text[i] = '|'.join([t.replace('\n', '').replace('\t', '').strip() for t in txt])
+    f.close()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='arguments for flintstones hdf5 file saving')
+    parser.add_argument('--data_dir', type=str, required=True, help='flintstones data directory')
+    parser.add_argument('--save_path', type=str, required=True, help='path to save hdf5')
+    args = parser.parse_args()
+    main(args)

data_script/pororo_hdf5.py

@@ -0,0 +1,83 @@
+import argparse
+import os
+
+import cv2
+import h5py
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+
+def main(args):
+    # 使用numpy库的load函数来加载名为descriptions.npy的文件。该文件是一个Python字典对象，因此我们使用item()方法将其转换为字典对象。
+    # ——os.path.join函数用于连接文件路径
+    # ——args.data_dir作为基础目录，将'descriptions.npy'添加到该目录中
+    # ——指定allow_pickle=True，表示允许加载包含Python对象的文件
+    # ——指定encoding='latin1'，表示使用拉丁字符编码加载该文件
+    descriptions = np.load(os.path.join(args.data_dir, 'descriptions.npy'), allow_pickle=True, encoding='latin1').item()
+    # imgs_list包含一组图像文件的路径，
+    # followings_list包含每个图像的一些附加信息
+    imgs_list = np.load(os.path.join(args.data_dir, 'img_cache4.npy'), encoding='latin1')
+    followings_list = np.load(os.path.join(args.data_dir, 'following_cache4.npy'))
+    # 使用numpy库的load函数来加载名为train_seen_unseen_ids.npy的文件
+    # 该文件包含三个numpy数组：train_ids、val_ids和test_ids，分别代表训练集、验证集和测试集的ID列表。
+    # 使用元组来一次性加载这三个数组，并将它们赋值给相应的变量。
+    train_ids, val_ids, test_ids = np.load(os.path.join(args.data_dir, 'train_seen_unseen_ids.npy'), allow_pickle=True)
+    # 按照ID的顺序逐一排序
+    train_ids = np.sort(train_ids)
+    val_ids = np.sort(val_ids)
+    test_ids = np.sort(test_ids)
+
+    # 创建一个新的HDF5文件，并指定文件名为args.save_path。
+    # 使用h5py库的File函数来创建文件对象，指定打开方式为写模式("w")。
+    # 在这个文件中存储处理后的图像和文本数据。
+    f = h5py.File(args.save_path, "w")
+    for subset, ids in {'train': train_ids, 'val': val_ids, 'test': test_ids}.items():
+        length = len(ids)
+
+        # 为每个数据集（train、val和test）创建一个组
+        # 针对每个数据集都创建了5个数据集，名为'image0'、'image1'、'image2'、'image3'、'image4'，分别对应于当前图像及其相关联的4个图像。
+        # 目的：将每个图像及其相关联的图像数据保存到同一个HDF5文件中，并按照一定的组织方式存储，方便后续的数据读取和处理。
+        group = f.create_group(subset)
+        # 创建一个长度为ids列表长度的空列表images，按照image0-4顺序添加了5个HDF5数据集对象
+        images = list()
+        # 为当前数据集中的每个图像创建了五个数据集。
+        # 每个数据集都使用vlen_dtype(np.dtype('uint8'))作为数据类型，并将其添加到当前组group中。
+        # ——vlen_dtype(np.dtype('uint8'))表示可变长度的无符号8位整数数组。
+        for i in range(5):
+            images.append(
+                group.create_dataset('image{}'.format(i), (length,), dtype=h5py.vlen_dtype(np.dtype('uint8'))))
+        # 创建一个数据集text，用于存储与当前数据集中图像相关的文本描述。该数据集的数据类型为字符串，编码方式为utf-8，并将其添加到当前组group中。
+        text = group.create_dataset('text', (length,), dtype=h5py.string_dtype(encoding='utf-8'))
+        # 遍历当前数据集中的每个图像，并将相关数据保存到HDF5文件中
+        for i, item in enumerate(tqdm(ids, leave=True, desc="saveh5")):
+            # 获取与当前图像相关的所有图像的路径，存储到列表img_paths中。
+            # ——imgs_list是一个字典，存储了所有图像的路径
+            # ——followings_list是一个字典，存储了与每个图像相关的四张图像的路径
+            img_paths = [str(imgs_list[item])[2:-1]] + [str(followings_list[item][i])[2:-1] for i in range(4)]
+            # 打开img_paths列表中的每个图像，并将其转换为RGB格式的PIL图像对象。
+            imgs = [Image.open(os.path.join(args.data_dir, img_path)).convert('RGB') for img_path in img_paths]
+            # 将每个PIL图像对象转换为numpy数组
+            for j, img in enumerate(imgs):
+                img = np.array(img).astype(np.uint8)
+                # 使用OpenCV将其编码为png格式的二进制数据
+                img = cv2.imencode('.png', img)[1].tobytes()
+                # 将该二进制数据转换为numpy数组
+                img = np.frombuffer(img, np.uint8)
+                # 将其存储到images列表中与当前图像相关的数据集中
+                images[j][i] = img
+            # 获取与当前图像相关的所有图像的文件名，并将其存储到列表tgt_img_ids中
+            tgt_img_ids = [str(img_path).replace('.png', '') for img_path in img_paths]
+            # 根据目标图像的文件名，获取其对应的文本描述，并将其存储到列表txt中。
+            txt = [descriptions[tgt_img_id][0] for tgt_img_id in tgt_img_ids]
+            # 将txt列表中的所有文本描述合并为一个字符串，并将其中的"\n"、"\t"等无关字符替换为空格。然后，将该字符串存储到数据集text中
+            text[i] = '|'.join([t.replace('\n', '').replace('\t', '').strip() for t in txt])
+    f.close()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='arguments for flintstones pororo file saving')
+    parser.add_argument('--data_dir', type=str, required=True, help='pororo data directory')
+    parser.add_argument('--save_path', type=str, required=True, help='path to save hdf5')
+    args = parser.parse_args()
+    main(args)

data_script/vist_hdf5.py

@@ -0,0 +1,111 @@
+import argparse
+import json
+import os
+
+import cv2
+import h5py
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+
+def main(args):
+    train_data = json.load(open(os.path.join(args.sis_json_dir, 'train.story-in-sequence.json')))
+    val_data = json.load(open(os.path.join(args.sis_json_dir, 'val.story-in-sequence.json')))
+    test_data = json.load(open(os.path.join(args.sis_json_dir, 'test.story-in-sequence.json')))
+
+    prefix = ["train", "val", "test"]
+    whole_album = {}
+    for i, data in enumerate([train_data, val_data, test_data]):
+        album_mapping = {}
+        for annot_new in data["annotations"]:
+            annot = annot_new[0]
+            assert len(annot_new) == 1
+            if annot['story_id'] not in album_mapping:
+                album_mapping[annot['story_id']] = {"flickr_id": [annot['photo_flickr_id']],
+                                                    "sis": [annot['original_text']],
+                                                    "length": 1}
+            else:
+                album_mapping[annot['story_id']]["flickr_id"].append(annot['photo_flickr_id'])
+                album_mapping[annot['story_id']]["sis"].append(
+                    annot['original_text'])
+                album_mapping[annot['story_id']]["length"] += 1
+        whole_album[prefix[i]] = album_mapping
+
+    for p in prefix:
+        deletables = []
+        for story_id, story in whole_album[p].items():
+            if story['length'] != 5:
+                print("deleting {}".format(story_id))
+                deletables.append(story_id)
+                continue
+            d = [os.path.exists(os.path.join(args.img_dir, "{}.jpg".format(_))) for _ in story["flickr_id"]]
+            if sum(d) < 5:
+                print("deleting {}".format(story_id))
+                deletables.append(story_id)
+            else:
+                pass
+        for i in deletables:
+            del whole_album[p][i]
+
+    train_data = json.load(open(os.path.join(args.sis_json_dir, 'train.description-in-isolation.json')))
+    val_data = json.load(open(os.path.join(args.sis_json_dir, 'val.description-in-isolation.json')))
+    test_data = json.load(open(os.path.join(args.sis_json_dir, 'test.description-in-isolation.json')))
+
+    flickr_id2text = {}
+    for i, data in enumerate([train_data, val_data, test_data]):
+        for l in data['annotations']:
+            assert len(l) == 1
+            if l[0]['photo_flickr_id'] in flickr_id2text:
+                flickr_id2text[l[0]['photo_flickr_id']] = \
+                    max([flickr_id2text[l[0]['photo_flickr_id']], l[0]['original_text']], key=len)
+            else:
+                flickr_id2text[l[0]['photo_flickr_id']] = l[0]['original_text']
+
+    for p in prefix:
+        deletables = []
+        for story_id, story in whole_album[p].items():
+            story['dii'] = []
+            for i, flickr_id in enumerate(story['flickr_id']):
+                if flickr_id not in flickr_id2text:
+                    print("{} not found in story {}".format(flickr_id, story_id))
+                    deletables.append(story_id)
+                    break
+                story['dii'].append(flickr_id2text[flickr_id])
+        for i in deletables:
+            del whole_album[p][i]
+
+    f = h5py.File(args.save_path, "w")
+    for p in prefix:
+        group = f.create_group(p)
+        story_dict = whole_album[p]
+        length = len(story_dict)
+        images = list()
+        for i in range(5):
+            images.append(
+                group.create_dataset('image{}'.format(i), (length,), dtype=h5py.vlen_dtype(np.dtype('uint8'))))
+        sis = group.create_dataset('sis', (length,), dtype=h5py.string_dtype(encoding='utf-8'))
+        dii = group.create_dataset('dii', (length,), dtype=h5py.string_dtype(encoding='utf-8'))
+        for i, (story_id, story) in enumerate(tqdm(story_dict.items(), leave=True, desc="saveh5")):
+            imgs = [Image.open('{}/{}.jpg'.format(args.img_dir, flickr_id)).convert('RGB') for flickr_id in
+                    story['flickr_id']]
+            for j, img in enumerate(imgs):
+                img = np.array(img).astype(np.uint8)
+                img = cv2.imencode('.png', img)[1].tobytes()
+                img = np.frombuffer(img, np.uint8)
+                images[j][i] = img
+            sis[i] = '|'.join([t.replace('\n', '').replace('\t', '').strip() for t in story['sis']])
+            txt_dii = [t.replace('\n', '').replace('\t', '').strip() for t in story['dii']]
+            txt_dii = sorted(set(txt_dii), key=txt_dii.index)
+            dii[i] = '|'.join(txt_dii)
+    f.close()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='arguments for vist hdf5 file saving')
+    parser.add_argument('--sis_json_dir', type=str, required=True, help='sis json file directory')
+    parser.add_argument('--dii_json_dir', type=str, required=True, help='dii json file directory')
+    parser.add_argument('--img_dir', type=str, required=True, help='json file directory')
+    parser.add_argument('--save_path', type=str, required=True, help='path to save hdf5')
+    args = parser.parse_args()
+    main(args)

data_script/vist_img_download.py

@@ -0,0 +1,61 @@
+import json
+import requests
+from io import BytesIO
+from PIL import Image
+from tqdm import tqdm
+from multiprocessing import Process
+import os
+import argparse
+
+
+def download_subprocess(dii, save_dir):
+    for image in tqdm(dii):
+        key, value = image.popitem()
+        try:
+            img_data = requests.get(value).content
+            img = Image.open(BytesIO(img_data)).convert('RGB')
+            h = img.size[0]
+            w = img.size[1]
+            if min(h, w) > 512:
+                img = img.resize((int(h / (w / 512)), 512) if h > w else (512, int(w / (h / 512))))
+            img.save('{}/{}.jpg'.format(save_dir, key))
+        except:
+            print(key, value)
+
+
+def main(args):
+    train_data = json.load(open(os.path.join(args.json_dir, 'train.description-in-isolation.json')))
+    val_data = json.load(open(os.path.join(args.json_dir, 'val.description-in-isolation.json')))
+    test_data = json.load(open(os.path.join(args.json_dir, 'test.description-in-isolation.json')))
+    dii = []
+    for subset in [train_data, val_data, test_data]:
+        for image in subset["images"]:
+            try:
+                dii.append({image['id']: image['url_o']})
+            except:
+                dii.append({image['id']: image['url_m']})
+
+    dii = [image for image in dii if not os.path.exists('{}/{}.jpg'.format(args.save_dir, list(image)[0]))]
+    print('total images: {}'.format(len(dii)))
+
+    def splitlist(inlist, chunksize):
+        return [inlist[x:x + chunksize] for x in range(0, len(inlist), chunksize)]
+
+    dii_splitted = splitlist(dii, int((len(dii) / args.num_process)))
+    process_list = []
+    for dii_sub_list in dii_splitted:
+        p = Process(target=download_subprocess, args=(dii_sub_list,))
+        process_list.append(p)
+        p.Daemon = True
+        p.start()
+    for p in process_list:
+        p.join()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='arguments for vist images downloading')
+    parser.add_argument('--json_dir', type=str, required=True, help='dii json file directory')
+    parser.add_argument('--img_dir', type=str, required=True, help='images saving directory')
+    parser.add_argument('--num_process', type=int, default=32)
+    args = parser.parse_args()
+    main(args)

datasets/flintstones.py

@@ -0,0 +1,93 @@
+import random
+
+import cv2
+import h5py
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+from transformers import CLIPTokenizer
+
+from models.blip_override.blip import init_tokenizer
+
+
+class StoryDataset(Dataset):
+    """
+    A custom subset class for the LRW (includes train, val, test) subset
+    """
+
+    def __init__(self, subset, args):
+        super(StoryDataset, self).__init__()
+        self.args = args
+
+        self.h5_file = args.get(args.dataset).hdf5_file
+        self.subset = subset
+
+        self.augment = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Resize([512, 512]),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5])
+        ])
+        self.dataset = args.dataset
+        self.max_length = args.get(args.dataset).max_length
+        self.clip_tokenizer = CLIPTokenizer.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="tokenizer")
+        self.blip_tokenizer = init_tokenizer()
+        msg = self.clip_tokenizer.add_tokens(list(args.get(args.dataset).new_tokens))
+        print("clip {} new tokens added".format(msg))
+        msg = self.blip_tokenizer.add_tokens(list(args.get(args.dataset).new_tokens))
+        print("blip {} new tokens added".format(msg))
+
+        self.blip_image_processor = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Resize([224, 224]),
+            transforms.ToTensor(),
+            transforms.Normalize([0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711])
+        ])
+
+    def open_h5(self):
+        h5 = h5py.File(self.h5_file, "r")
+        self.h5 = h5[self.subset]
+
+    def __getitem__(self, index):
+        if not hasattr(self, 'h5'):
+            self.open_h5()
+
+        images = list()
+        for i in range(5):
+            im = self.h5['image{}'.format(i)][index]
+            im = cv2.imdecode(im, cv2.IMREAD_COLOR)
+            idx = random.randint(0, 4)
+            images.append(im[idx * 128: (idx + 1) * 128])
+
+        source_images = torch.stack([self.blip_image_processor(im) for im in images])
+        images = images[1:] if self.args.task == 'continuation' else images
+        images = torch.stack([self.augment(im) for im in images]) \
+            if self.subset in ['train', 'val'] else torch.from_numpy(np.array(images)).permute(0, 3, 1, 2)
+
+        texts = self.h5['text'][index].decode('utf-8').split('|')
+
+        # tokenize caption using default tokenizer
+        tokenized = self.clip_tokenizer(
+            texts[1:] if self.args.task == 'continuation' else texts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=False,
+            return_tensors="pt",
+        )
+        captions, attention_mask = tokenized['input_ids'], tokenized['attention_mask']
+
+        tokenized = self.blip_tokenizer(
+            texts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=False,
+            return_tensors="pt",
+        )
+        source_caption, source_attention_mask = tokenized['input_ids'], tokenized['attention_mask']
+        return images, captions, attention_mask, source_images, source_caption, source_attention_mask
+
+    def __len__(self):
+        if not hasattr(self, 'h5'):
+            self.open_h5()
+        return len(self.h5['text'])

datasets/pororo.py

@@ -0,0 +1,144 @@
+import copy
+import os
+import random
+from PIL import Image
+import cv2
+import h5py
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+from transformers import CLIPTokenizer
+
+from models.blip_override.blip import init_tokenizer
+
+
+class StoryDataset(Dataset):
+    """
+    A custom subset class for the LRW (includes train, val, test) subset
+    """
+    # StoryDataset 类的构造函数
+    def __init__(self, subset, args):
+        # 用来调用父类 Dataset 的初始化函数，确保该类能够继承 Dataset 类的所有方法和属性。
+        super(StoryDataset, self).__init__()
+        # args 则是该类的其他参数，是一个命名空间（namespace）对象
+        self.args = args
+        # 一个 HDF5 文件的路径，存储了训练、验证和测试集的图像和文本数据。
+        # ——args.get(args.dataset)表示从命名空间对象args中获取指定数据集（训练集、验证集或测试集）的参数。
+        self.h5_file = args.get(args.dataset).hdf5_file
+        # 初始化函数中 subset 表示要读取的子集的类型（如训练集、验证集、测试集）
+        self.subset = subset
+
+        # 一个图像变换函数序列（transform），用来对图像进行预处理，包括将图像转化为 PIL 格式，调整图像大小，将图像转换为 Tensor，并进行归一化。
+        self.augment = transforms.Compose([
+            transforms.ToPILImage(),
+           # transforms.Resize([256, 256]),
+            transforms.Resize([512, 512]),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5])
+        ])
+        # 表示当前数据集的类型（训练集、验证集或测试集）
+        self.dataset = args.dataset
+        # 最大的 caption 长度,在进行tokenize操作时，caption中的单词数量将被填充到该长度。
+        self.max_length = args.get(args.dataset).max_length
+        # 一个使用CLIP模型进行tokenize的tokenizer
+        self.clip_tokenizer = CLIPTokenizer.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="tokenizer")
+        # 一个自定义的tokenizer，用于处理文本输入
+        self.blip_tokenizer = init_tokenizer()
+        msg = self.clip_tokenizer.add_tokens(list(args.get(args.dataset).new_tokens))
+        print("clip {} new tokens added".format(msg))
+        msg = self.blip_tokenizer.add_tokens(list(args.get(args.dataset).new_tokens))
+        print("blip {} new tokens added".format(msg))
+
+        # 一个用于对输入的图像进行处理的函数序列，包括转换为PIL图像、重置图像大小、转换为tensor、归一化等。
+        self.blip_image_processor = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Resize([224, 224]),
+            transforms.ToTensor(),
+            transforms.Normalize([0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711])
+        ])
+
+    #  打开与数据集对应的h5文件
+    def open_h5(self):
+        h5 = h5py.File(self.h5_file, "r")
+        self.h5 = h5[self.subset]
+
+    # 用于按索引获取数据。
+
+    # 对于每个图像，都进行数据增强操作，以进行数据增强。
+    # 然后，将文本输入的caption进行tokenize操作，
+    # 使用CLIP tokenizer和自定义tokenizer分别进行tokenize。
+    # 最后，将处理好的图像、caption和attention mask返回
+    def __getitem__(self, index):
+        # 首先调用open_h5()打开数据集的h5文件
+        if not hasattr(self, 'h5'):
+            self.open_h5()
+        #index = 1
+        images = list()
+        for i in range(5):
+            # 从h5文件中读取一组图像和对应的文本。
+            im = self.h5['image{}'.format(i)][index]
+            # print(im)
+            # pil_img = Image.fromarray(im)
+            # # 保存图像
+            # pil_img.save(os.path.join('/root/lihui/StoryVisualization/ori_test_images', '{:04d}.png'.format(i)))
+            # 对每个图像解码
+            im = cv2.imdecode(im, cv2.IMREAD_COLOR)
+            # 随机选择一个128像素的图像切片
+            idx = random.randint(0, im.shape[0] / 128 - 1)
+            # 将切片后的图像加到images列表中
+            images.append(im[idx * 128: (idx + 1) * 128])
+        # 深拷贝，后续不随images变化
+        ori_images = copy.deepcopy(images)
+        # 保存test原始图像
+
+        # for i, im in enumerate(images):
+        #     file_path = '/root/lihui/StoryVisualization/ori_test_images/group{:02d}_image{:02d}.png'.format(index + 1,
+        #                                                                                                     i + 1)
+        #     cv2.imwrite(file_path, im)
+        # 将图像转换为张量
+        source_images = torch.stack([self.blip_image_processor(im) for im in images])
+        # 如果为continuation任务，将列表中的第一个图像从images中移除
+        images = images[1:] if self.args.task == 'continuation' else images
+        # 如果subset的值为train/val，则使用augment方法对images列表中的所有图像进行数据增强，并将其转换为张量
+        # 否则使用numpy.array方法将images列表转换为张量，并进行转置操作
+        images = torch.stack([self.augment(im) for im in images]) \
+            if self.subset in ['train', 'val'] else torch.from_numpy(np.array(images)).permute(0, 3, 1, 2)
+        ######################
+        # 读取当前索引处的文本，并使用decode方法将其解码为UTF-8
+        texts = self.h5['text'][index].decode('utf-8').split('|')
+        # print(f"index: {index}")
+        # for text in texts:
+        #     print(f"texts: {text}")
+
+        # tokenize caption using default tokenizer
+        tokenized = self.clip_tokenizer(
+            texts[1:] if self.args.task == 'continuation' else texts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=False,
+            return_tensors="pt",
+        )
+        captions, attention_mask = tokenized['input_ids'], tokenized['attention_mask']
+
+        tokenized = self.blip_tokenizer(
+            texts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=False,
+            return_tensors="pt",
+        )
+        source_caption, source_attention_mask = tokenized['input_ids'], tokenized['attention_mask']
+        return images, captions, attention_mask, source_images, source_caption, source_attention_mask, texts, ori_images
+
+    # 返回数据集中样本的数量
+    # 如果是测试集，则返回100，否则返回对应的数据集中的样本数量
+    def __len__(self):
+        if not hasattr(self, 'h5'):
+            self.open_h5()
+        if self.subset == 'test':
+            #print('')
+            return 1
+        # if self.subset == 'test':
+        #     return 100
+        return len(self.h5['text'])

datasets/vistdii.py

@@ -0,0 +1,94 @@
+import cv2
+import h5py
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+from transformers import CLIPTokenizer
+
+from models.blip_override.blip import init_tokenizer
+
+
+class StoryDataset(Dataset):
+    """
+    A custom subset class for the LRW (includes train, val, test) subset
+    """
+
+    def __init__(self, subset, args):
+        super(StoryDataset, self).__init__()
+        self.args = args
+
+        self.h5_file = args.get(args.dataset).hdf5_file
+        self.subset = subset
+
+        self.augment = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Resize(512),
+            transforms.RandomCrop(512) if self.subset == 'train' else transforms.CenterCrop(512),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5])
+        ]) if self.subset in ['train', 'val'] else transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Resize(64),
+            transforms.CenterCrop(64)
+        ])
+
+        self.dataset = args.dataset
+        self.max_length = args.get(args.dataset).max_length
+        self.clip_tokenizer = CLIPTokenizer.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="tokenizer")
+        self.blip_tokenizer = init_tokenizer()
+
+        self.blip_image_processor = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Resize(224),
+            transforms.RandomCrop(224) if self.subset == 'train' else transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize([0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711])
+        ])
+
+    def open_h5(self):
+        h5 = h5py.File(self.h5_file, "r")
+        self.h5 = h5[self.subset]
+
+    def __getitem__(self, index):
+        if not hasattr(self, 'h5'):
+            self.open_h5()
+
+        images = list()
+        for i in range(5):
+            im = self.h5['image{}'.format(i)][index]
+            im = cv2.imdecode(im, cv2.IMREAD_COLOR)
+            images.append(im)
+
+        source_images = torch.stack([self.blip_image_processor(im) for im in images])
+        images = images[1:] if self.args.task == 'continuation' else images
+        images = [self.augment(im) for im in images]
+        images = torch.stack(images) if self.subset in ['train', 'val'] \
+            else torch.from_numpy(np.array([np.array(im) for im in images])).permute(0, 3, 1, 2)
+
+        texts = self.h5['dii'][index].decode('utf-8').split('|')
+
+        # tokenize caption using default tokenizer
+        tokenized = self.clip_tokenizer(
+            texts[1:] if self.args.task == 'continuation' else texts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=False,
+            return_tensors="pt",
+        )
+        captions, attention_mask = tokenized['input_ids'], tokenized['attention_mask']
+
+        tokenized = self.blip_tokenizer(
+            texts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=False,
+            return_tensors="pt",
+        )
+        source_caption, source_attention_mask = tokenized['input_ids'], tokenized['attention_mask']
+        return images, captions, attention_mask, source_images, source_caption, source_attention_mask
+
+    def __len__(self):
+        if not hasattr(self, 'h5'):
+            self.open_h5()
+        return len(self.h5['dii'])

datasets/vistsis.py

@@ -0,0 +1,94 @@
+import cv2
+import h5py
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+from transformers import CLIPTokenizer
+
+from models.blip_override.blip import init_tokenizer
+
+
+class StoryDataset(Dataset):
+    """
+    A custom subset class for the LRW (includes train, val, test) subset
+    """
+
+    def __init__(self, subset, args):
+        super(StoryDataset, self).__init__()
+        self.args = args
+
+        self.h5_file = args.get(args.dataset).hdf5_file
+        self.subset = subset
+
+        self.augment = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Resize(512),
+            transforms.RandomCrop(512) if self.subset == 'train' else transforms.CenterCrop(512),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5])
+        ]) if self.subset in ['train', 'val'] else transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Resize(64),
+            transforms.CenterCrop(64)
+        ])
+
+        self.dataset = args.dataset
+        self.max_length = args.get(args.dataset).max_length
+        self.clip_tokenizer = CLIPTokenizer.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="tokenizer")
+        self.blip_tokenizer = init_tokenizer()
+
+        self.blip_image_processor = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Resize(224),
+            transforms.RandomCrop(224) if self.subset == 'train' else transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize([0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711])
+        ])
+
+    def open_h5(self):
+        h5 = h5py.File(self.h5_file, "r")
+        self.h5 = h5[self.subset]
+
+    def __getitem__(self, index):
+        if not hasattr(self, 'h5'):
+            self.open_h5()
+
+        images = list()
+        for i in range(5):
+            im = self.h5['image{}'.format(i)][index]
+            im = cv2.imdecode(im, cv2.IMREAD_COLOR)
+            images.append(im)
+
+        source_images = torch.stack([self.blip_image_processor(im) for im in images])
+        images = images[1:] if self.args.task == 'continuation' else images
+        images = [self.augment(im) for im in images]
+        images = torch.stack(images) if self.subset in ['train', 'val'] \
+            else torch.from_numpy(np.array([np.array(im) for im in images])).permute(0, 3, 1, 2)
+
+        texts = self.h5['sis'][index].decode('utf-8').split('|')
+
+        # tokenize caption using default tokenizer
+        tokenized = self.clip_tokenizer(
+            texts[1:] if self.args.task == 'continuation' else texts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=False,
+            return_tensors="pt",
+        )
+        captions, attention_mask = tokenized['input_ids'], tokenized['attention_mask']
+
+        tokenized = self.blip_tokenizer(
+            texts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=False,
+            return_tensors="pt",
+        )
+        source_caption, source_attention_mask = tokenized['input_ids'], tokenized['attention_mask']
+        return images, captions, attention_mask, source_images, source_caption, source_attention_mask
+
+    def __len__(self):
+        if not hasattr(self, 'h5'):
+            self.open_h5()
+        return len(self.h5['sis'])