-
Notifications
You must be signed in to change notification settings - Fork 1
/
dataset.py
59 lines (50 loc) · 1.96 KB
/
dataset.py
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
from torch import Tensor
import torch
from PIL import Image
import numpy as np
from torch.utils.data import Dataset
from typing import Dict
from group_video_segment import GroupSegment
class EPIC(Dataset):
def __init__(self, mode, cfg, transforms=None, test_mode=None, participant_id=None):
super().__init__()
self.cfg = cfg
self.n_segments = self.cfg.n_segments
assert mode in ["train-seen", "train-unseen", "val-seen", "val-unseen", "whole-train", "test"]
self.setname = mode
self.transforms = transforms
self.test_mode = test_mode
# Define Group Segment Extractor
self.group_segment = GroupSegment(self.cfg, self.setname, self.test_mode, participant_id)
self.haveData = self.group_segment.haveData
self.verb_label = self.group_segment.verb_labels
self.noun_label = self.group_segment.noun_labels
def __len__(self) -> int:
return len(self.group_segment.items)
def __getitem__(self, idx: int) -> Dict[str, Tensor]:
self.group_segment._process_snapshot_(idx)
# Looping over # of segments
data = []
for seg_indx in range(self.n_segments):
tmp = Image.open(self.group_segment._get_image_path_(seg_indx)).convert(
"RGB"
)
data.append(tmp)
if self.transforms:
data = self.transforms(data)
# PyTorch automatically converts data to Torch Tensors
if self.setname == "test":
return {
"data": data,
"uid": self.group_segment.uid
}
else:
return {
"data": data,
"verb_id": torch.Tensor(
np.expand_dims(self.group_segment.ann_verb, 1)
).type(torch.LongTensor),
"noun_id": torch.Tensor(
np.expand_dims(self.group_segment.ann_noun, 1)
).type(torch.LongTensor),
}