bioemu-cv/data.py at main · seonghyun26/bioemu-cv · GitHub

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
78
79
80
81
82
83
84
85
import torch
from torch.utils.data import Dataset

class TimelagDataset(Dataset):
    def __init__(
        self,
        cfg_data,
        device,
    ):
        self.cfg = cfg_data
        self.sequence = cfg_data.sequence
        self.representation = cfg_data.representation
        self.time_lag = cfg_data.time_lag
        self.dataset_size = cfg_data.dataset_size
        self.system_id = cfg_data.system_id
        self.data_dir = cfg_data.data_dir
        self.device = device

        self._load_data()

    def _load_data(self):
        self.current_cad_path = f"{self.data_dir}/{self.system_id}-{self.dataset_size}/current-cad.pt"
        self.current_pos_path = f"{self.data_dir}/{self.system_id}-{self.dataset_size}/current-pos.pt"
        self.current_orientation_path = f"{self.data_dir}/{self.system_id}-{self.dataset_size}/current-orientation.pt"

        if self.time_lag ==0:
            self.timelag_cad_path = f"{self.data_dir}/{self.system_id}-{self.dataset_size}/current-cad.pt"
            self.timelag_pos_path = f"{self.data_dir}/{self.system_id}-{self.dataset_size}/current-pos.pt"
            self.timelag_orientation_path = f"{self.data_dir}/{self.system_id}-{self.dataset_size}/current-orientation.pt"
        else:
            self.timelag_cad_path = f"{self.data_dir}/{self.system_id}-{self.dataset_size}/lag{self.time_lag}-cad.pt"
            self.timelag_pos_path = f"{self.data_dir}/{self.system_id}-{self.dataset_size}/lag{self.time_lag}-pos.pt"
            self.timelag_orientation_path = f"{self.data_dir}/{self.system_id}-{self.dataset_size}/lag{self.time_lag}-orientation.pt"

        self.current_cad = torch.load(self.current_cad_path, map_location=self.device)
        self.current_pos = torch.load(self.current_pos_path, map_location=self.device)
        self.timelag_cad = torch.load(self.timelag_cad_path, map_location=self.device)
        self.timelag_pos = torch.load(self.timelag_pos_path, map_location=self.device)
        self.current_orientation = torch.load(self.current_orientation_path, map_location=self.device)
        self.timelag_orientation = torch.load(self.timelag_orientation_path, map_location=self.device)

    def __len__(self):
        if self.representation == "cad":
            return len(self.current_cad)
        elif self.representation == "pos":
            return len(self.current_pos)
        elif self.representation == "cad-pos":
            return len(self.current_cad)
        elif self.representation == "cad-pos-ac":
            return len(self.current_cad)
        else:
            raise ValueError(f"Invalid representation: {self.representation}")

    def __getitem__(self, idx):
        if self.representation == "cad":
            return {
                "current_data": self.current_cad[idx],
                "timelagged_data": self.timelag_cad[idx],
                "current_orientation": self.current_orientation[idx],
                "timelagged_orientation": self.timelag_orientation[idx]
            }
        elif self.representation == "pos":
            return {
                "current_data": self.current_pos[idx],
                "timelagged_data": self.timelag_pos[idx],
                "current_orientation": self.current_orientation[idx],
                "timelagged_orientation": self.timelag_orientation[idx]
            }
        elif self.representation == "cad-pos":
            return {
                "current_data": self.current_cad[idx],
                "timelagged_data": self.timelag_pos[idx],
                "current_orientation": self.current_orientation[idx],
                "timelagged_orientation": self.timelag_orientation[idx]
            }
        elif self.representation == "cad-pos-ac":
            return {
                "current_data": self.current_cad[idx],
                "timelagged_data": self.timelag_pos[idx],
                "current_orientation": self.current_orientation[idx],
                "timelagged_orientation": self.timelag_orientation[idx],
                "timelagged_cad": self.timelag_cad[idx],
            }
        else:
            raise ValueError(f"Invalid representation: {self.representation}")