-
Notifications
You must be signed in to change notification settings - Fork 39
/
dataset.py
43 lines (30 loc) · 1.08 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
from typing import Optional
from torch.utils.data import Dataset
from slam.common.utils import assert_debug
class WrapperDataset(Dataset):
"""
A Wrapper Dataset which applies a transform on the items returned by a torch.Dataset
"""
def __init__(self, dataset: Dataset, transform: Optional[callable] = None):
self.dataset = dataset
self.transform = transform
def __len__(self):
return len(self.dataset)
def __getitem__(self, item):
data = self.dataset[item]
transformed = self.transform(data)
return transformed
class WindowDataset(Dataset):
"""
A Window datasets wraps a dataset_config, and limit it to a defined window
"""
def __init__(self, dataset: Dataset, start: int = 0, length: int = -1):
self.dataset = dataset
assert_debug(start < len(dataset))
assert_debug(start + length <= len(dataset))
self.length = length
self.start = start
def __len__(self):
return self.length
def __getitem__(self, item):
return self.dataset[item + self.start]