Spaces:
Runtime error
Runtime error
| import torch | |
| import torchvision as tv | |
| import numpy as np | |
| import glob | |
| import os | |
| import random | |
| from torch.utils.data import Dataset | |
| class MultiviewDataset(Dataset): | |
| def __init__(self, data_folder, image_size, num_view): | |
| super(MultiviewDataset, self).__init__() | |
| self.image_size = image_size | |
| self.num_view = num_view | |
| self.loader = tv.datasets.folder.default_loader | |
| self.transform = tv.transforms.Compose([tv.transforms.Resize(self.image_size), tv.transforms.ToTensor()]) | |
| self.folders = sorted(glob.glob(os.path.join(data_folder, '*'))) | |
| self.camera_ids = ['220700191', '221501007', '222200036', '222200037', '222200038', '222200039', '222200040', '222200041', | |
| '222200042', '222200043', '222200044', '222200045', '222200046', '222200047', '222200048', '222200049'] | |
| def get_item(self, index): | |
| data = self.__getitem__(index) | |
| return data | |
| def __getitem__(self, index): | |
| images = torch.stack([self.transform(self.loader(self.folders[index] + '/image_%s.jpg' % self.camera_ids[v])) for v in range(self.num_view)]) | |
| intrinsics = torch.stack([torch.from_numpy(np.load(self.folders[index] + '/camera_%s.npz' % self.camera_ids[v])['intrinsic']) for v in range(self.num_view)]).float() | |
| extrinsics = torch.stack([torch.from_numpy(np.load(self.folders[index] + '/camera_%s.npz' % self.camera_ids[v])['extrinsic']) for v in range(self.num_view)]).float() | |
| return {'images': images, | |
| 'intrinsics': intrinsics, | |
| 'extrinsics': extrinsics, | |
| 'exp_id': int(self.folders[index][-4:])} | |
| def __len__(self): | |
| return len(self.folders) | |