MultiModalSA/datasets.py at master · Ighina/MultiModalSA · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
# -*- coding: utf-8 -*-
"""
Created on Tue Aug  4 13:02:35 2020

@author: Iacopo
"""
import torch
import numpy as np
from torch.utils.data import Dataset

class VideosSentenceDataset(Dataset):
    def __init__(self, lines):
        self.src_dataset = [line for line in lines[0]]
        self.src_sizes = np.array([len(tokens) for tokens in self.src_dataset])
        self.text = [line for line in lines[2]]
        self.tgt_dataset = [line for line in lines[1]]
        self.video = [line for line in lines[3]]


    def __getitem__(self, index):
        return {
            'id': torch.tensor(index),
            'source': torch.tensor(self.src_dataset[index]),
            'target': self.tgt_dataset[index],
            'text': torch.tensor(self.text[index]).float(),
            'video': torch.tensor(self.video[index])
        }

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

    def collater(self, samples):
        """Merge a list of samples to form a mini-batch."""
        if len(samples) == 0:
            return {}
        def merge(values, continuous=False):
            if len(values[0].shape)<2:
              return torch.stack(values)
            else:
              max_length = max(v.size(0) for v in values)
              result = torch.zeros((len(values),max_length, values[0].shape[1]))
              for i, v in enumerate(values):
                  result[i, :len(v)] = v
              return result

        id = torch.tensor([s['id'] for s in samples])
        # src_tokens = merge([s['source'] for s in samples])
        src_tokens = merge([s['source'].float() for s in samples])
        tgt_tokens = torch.tensor([s['target'] for s in samples])
        text = torch.stack([s['text'] for s in samples])
        src_lengths = torch.LongTensor([s['source'].shape[0] for s in samples])
        video = merge([s['video'].float() for s in samples])

        return {
            'id': id,
            'src_tokens': src_tokens,
            'src_lengths': src_lengths,
            'tgt_tokens': tgt_tokens,
            'text': text,
            'video': video
        }

class VideosWordsDataset(Dataset):
    def __init__(self, lines):
        self.src_dataset = [line for line in lines[0]]
        self.src_sizes = np.array([len(tokens) for tokens in self.src_dataset])
        self.text = [line for line in lines[2]]
        self.tgt_dataset = [line for line in lines[1]]
        self.video = [line for line in lines[3]]


    def __getitem__(self, index):
        return {
            'id': torch.tensor(index),
            'source': torch.tensor(self.src_dataset[index]),
            'target': self.tgt_dataset[index],
            'text': torch.tensor(self.text[index]).float(),
            'video': torch.tensor(self.video[index])
        }

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

    def collater(self, samples):
        """Merge a list of samples to form a mini-batch."""
        if len(samples) == 0:
            return {}
        def merge(values, continuous=False):
            max_length = max(v.size(0) for v in values)
            result = torch.zeros((len(values),max_length, values[0].shape[1]))
            for i, v in enumerate(values):
                result[i, :len(v)] = v
            return result

        id = torch.tensor([s['id'] for s in samples])
        # src_tokens = merge([s['source'] for s in samples])
        src_tokens = merge([s['source'] for s in samples])
        tgt_tokens = torch.tensor([s['target'] for s in samples])
        text = torch.stack([s['text'] for s in samples])
        src_lengths = torch.LongTensor([s['source'].shape[0] for s in samples])
        video = merge([s['video'] for s in samples])

        return {
            'id': id,
            'src_tokens': src_tokens,
            'src_lengths': src_lengths,
            'tgt_tokens': tgt_tokens,
            'text': text,
            'video': video
        }