Class TikTecDataset

code/utils/dataloader.py:357–426

class TikTecDataset(Dataset):

    def __init__(self, path_vid):
        self.data_complete = pd.read_json('./data/data.json',orient='records',dtype=False,lines=True)

        self.vid = []
        with open(f'./data/vids/{path_vid}', "r") as fr:
            for line in fr.readlines():
                self.vid.append(line.strip())
        self.data = self.data_complete[self.data_complete['video_id'].isin(self.vid)]

    def __len__(self):
        return self.data.shape[0]

    def __getitem__(self, idx):
        item = self.data.iloc[idx]
        vid = item['video_id']

        label = 1 if item['label']=='假' else 0
        label = torch.tensor(label)

        max_K = 200  # max num of frames
        max_N = 500  # max num of ASR words

        # get caption feature
        with open('./data/caption_w2v_pad%s.pkl' % vid, 'rb') as f:
            caption_feature = pickle.load(f)  # (num_frame, 100, 300)
        if max_K / caption_feature.shape[0] >= 2:
            times = math.floor(max_K / caption_feature.shape[0])
            caption_feature = caption_feature.repeat_interleave(times, dim=0)
        elif caption_feature.shape[0] > max_K:
            times = math.ceil(caption_feature.shape[0] / max_K)
            caption_feature = caption_feature[::times][:max_K]
        actual_K = caption_feature.shape[0]
        caption_feature = torch.cat([caption_feature, torch.zeros((max_K - caption_feature.shape[0], 100, 300))], dim=0)

        # get visual feature
        with open( './data/vgg19_result%s.pkl' % vid, 'rb') as f:
            visual_feature = pickle.load(f)  # (num_frame, 45, 1000)
        if max_K / visual_feature.shape[0] >= 2:
            times = math.floor(max_K / visual_feature.shape[0])
            visual_feature = visual_feature.repeat_interleave(times, dim=0)
        elif visual_feature.shape[0] > max_K:
            times = math.ceil(visual_feature.shape[0] / max_K)
            visual_feature = visual_feature[::times][:max_K]
        visual_feature = torch.cat([visual_feature, torch.zeros((max_K - visual_feature.shape[0], 45, 1000))], dim=0)

        # get ASR feature
        with open('./data/asr_w2v+mfcc%s.pkl' % vid, 'rb') as f:
            asr_feature = pickle.load(f)  # (num_word, 300+650)
        asr_feature = asr_feature[:max_N]
        actual_N = asr_feature.shape[0]
        asr_feature = torch.cat([asr_feature, torch.zeros((max_N - asr_feature.shape[0], 300+650))], dim=0)

        # get frames mask & ASR words mask
        mask_K = torch.zeros(max_K, dtype=torch.int)
        mask_K[:actual_K] = 1
        mask_N = torch.zeros(max_N, dtype=torch.int)