/
p3d_dataset.py
56 lines (43 loc) · 1.46 KB
/
p3d_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
import torch
import torch.utils.data as data
from PIL import Image
import os
import os.path
import numpy as np
class ClipRecord(object):
def __init__(self, row):
self._data = row
@property
def path(self):
return self._data[0]
@property
def start_frames(self):
return int(self._data[1])
@property
def label(self):
return int(self._data[2])
class P3DDataSet(data.Dataset):
def __init__(self, list_file,length=16, modality='RGB',image_tmpl='frame{:06d}.jpg', transform=None):
self.list_file = list_file
self.length = length
self.modality = modality
self.image_tmpl = image_tmpl
self.transform = transform
self._parse_list()
def _load_image(self, directory, idx):
if self.modality == 'RGB':
return Image.open(os.path.join(directory, self.image_tmpl.format(idx)))
def _parse_list(self):
self.clip_list = [ClipRecord(x.strip().split(' ')) for x in open(self.list_file)]
def __getitem__(self, index):
record = self.clip_list[index]
return self.get(record)
def get(self, record):
clip = list()
for i in range(self.length):
img = self._load_image(record.path,i+record.start_frames)
clip.append(img)
clip = self.transform(clip)
return clip, record.label
def __len__(self):
return len(self.clip_list)