luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix static cache

Browse Source
This commit is contained in:
luni
2023-06-30 23:27:53 +02:00
parent 25d9814b70
commit d3ff1c05e1
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/dataset.py
+3 -3
View File
@@ -105,15 +105,15 @@ class StaticCache():
def __getitem__(self, frame) -> dict: def __getitem__(self, frame) -> dict:
if frame not in self.store: if (frame.id, frame.garment) not in self.store:
item = self.read(frame.id, frame.garment) item = self.read(frame.id, frame.garment)
stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()} stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()}
for key, value in stored_data.items(): for key, value in stored_data.items():
if key.startswith('uv_') and value.shape[-1] == 3: if key.startswith('uv_') and value.shape[-1] == 3:
stored_data[key] = torch.permute(value, [2,0,1]) stored_data[key] = torch.permute(value, [2,0,1])
self.store[frame] = stored_data self.store[(frame.id, frame.garment)] = stored_data
data = { key: value.to(self.device) for key, value in self.store[frame].items()} data = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)].items()}
return data return data