File size: 2,925 Bytes
d4cbafd | 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 | from torch.utils import data
import numpy as np
from .preprocessing import get_node_timestep_data
class EnvironmentDataset(object):
def __init__(self, env, state, pred_state, node_freq_mult, scene_freq_mult, hyperparams, **kwargs):
self.env = env
self.state = state
self.pred_state = pred_state
self.hyperparams = hyperparams
self.max_ht = self.hyperparams['maximum_history_length']
self.max_ft = kwargs['min_future_timesteps']
self.node_type_datasets = list()
self._augment = False
for node_type in env.NodeType:
if node_type not in hyperparams['pred_state']:
continue
self.node_type_datasets.append(NodeTypeDataset(env, node_type, state, pred_state, node_freq_mult,
scene_freq_mult, hyperparams, **kwargs))
@property
def augment(self):
return self._augment
@augment.setter
def augment(self, value):
self._augment = value
for node_type_dataset in self.node_type_datasets:
node_type_dataset.augment = value
def __iter__(self):
return iter(self.node_type_datasets)
class NodeTypeDataset(data.Dataset):
def __init__(self, env, node_type, state, pred_state, node_freq_mult,
scene_freq_mult, hyperparams, augment=False, **kwargs):
self.env = env
self.state = state
self.pred_state = pred_state
self.hyperparams = hyperparams
self.max_ht = self.hyperparams['maximum_history_length']
self.max_ft = kwargs['min_future_timesteps']
self.augment = augment
self.node_type = node_type
self.index = self.index_env(node_freq_mult, scene_freq_mult, **kwargs)
self.len = len(self.index)
self.edge_types = [edge_type for edge_type in env.get_edge_types() if edge_type[0] is node_type]
def index_env(self, node_freq_mult, scene_freq_mult, **kwargs):
index = list()
for scene in self.env.scenes:
present_node_dict = scene.present_nodes(np.arange(0, scene.timesteps), type=self.node_type, **kwargs)
for t, nodes in present_node_dict.items():
for node in nodes:
index += [(scene, t, node)] *\
(scene.frequency_multiplier if scene_freq_mult else 1) *\
(node.frequency_multiplier if node_freq_mult else 1)
return index
def __len__(self):
return self.len
def __getitem__(self, i):
(scene, t, node) = self.index[i]
if self.augment:
scene = scene.augment()
node = scene.get_node_by_id(node.id)
return get_node_timestep_data(self.env, scene, t, node, self.state, self.pred_state,
self.edge_types, self.max_ht, self.max_ft, self.hyperparams)
|