-
Notifications
You must be signed in to change notification settings - Fork 8
/
dataset.py
executable file
·69 lines (59 loc) · 2.31 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
58
59
60
61
62
63
64
65
66
67
68
69
from datasets.Brain4cars import Brain4cars_Inside, Brain4cars_Outside
def get_training_set(opt, spatial_transform, horizontal_flip, temporal_transform,
target_transform):
assert opt.dataset in ['Brain4cars_Inside', 'Brain4cars_Outside']
if opt.dataset == 'Brain4cars_Inside':
training_data = Brain4cars_Inside(
opt.video_path,
opt.annotation_path,
'training',
opt.n_fold,
opt.end_second,
1,
spatial_transform=spatial_transform,
horizontal_flip=horizontal_flip,
temporal_transform=temporal_transform,
target_transform=target_transform)
elif opt.dataset == 'Brain4cars_Outside':
training_data = Brain4cars_Outside(
opt.video_path,
opt.annotation_path,
'training',
opt.n_fold,
opt.end_second,
10,
spatial_transform=spatial_transform,
horizontal_flip=horizontal_flip,
target_transform=target_transform,
temporal_transform=temporal_transform)
return training_data
def get_validation_set(opt, spatial_transform, temporal_transform,
target_transform):
assert opt.dataset in ['Brain4cars_Inside', 'Brain4cars_Outside']
if opt.dataset == 'Brain4cars_Inside':
validation_data = Brain4cars_Inside(
opt.video_path,
opt.annotation_path,
'validation',
opt.n_fold,
opt.end_second,
opt.n_val_samples,
spatial_transform,
None,
temporal_transform,
target_transform,
sample_duration=opt.sample_duration)
elif opt.dataset == 'Brain4cars_Outside':
validation_data = Brain4cars_Outside(
opt.video_path,
opt.annotation_path,
'validation',
opt.n_fold,
opt.end_second,
opt.n_val_samples,
spatial_transform=spatial_transform,
horizontal_flip=None,
temporal_transform=temporal_transform,
target_transform=target_transform,
sample_duration=opt.sample_duration)
return validation_data