在PyTorch中从一组.mat文件创建数据集类

ia2d9nvy  于 2022-12-13  发布在  其他
关注(0)|答案(1)|浏览(293)

我正在处理一个LIDAR问题,我试图将一组3D直方图测量值(大小[64,64,128])Map到原始深度图(大小[64,64])。我有一个充满.mat文件的文件夹,每个文件都有自己的记录直方图和地面真实深度测量值。
在示例图片中,我有一个示例测量的.mat文件。现在重要的是SPAD和深度。
.mat文件示例:

我想在PyTorch中创建一个具有SPAD度量的Dataset类,但在给定的文件夹下,我被困在如何创建该类上。
您可以在PyTorch中使用以下命令读取.mat文件:

mat = scipy.io.loadmat('drive/My Drive/Colab Notebooks/LIDAR/SPAD_NYU/SPAD_Counts/spad_measurements_dining_room_0001a_0051.mat')
spad = mat['spad']
depth = mat['depth']

我知道您可以使用以下方法创建自定义数据集:

from torch.utils.data import Dataset, IterableDataset, Dataloader
class CTSet(Dataset):
  def __init__(self):
    
  def __len__(self):
    return len(self.dataset)
  def __getitem__(self,index):
    return self.dataset[idx],self.labels[idx]

然后使用Dataloader进行训练。我在Google Colab中将文件夹SPAD_Counts安装到Drive中。我希望:
1.循环文件夹
1.对于每个.mat文件,获取spad和depth文件。
1.将这些np数组转换为TorchTensor。(假设我是在init中完成的)
1.当我使用Dataloader进行训练时,以可用于迭代的方式返回特定测量的spad和深度(假定在getitem中完成)
我已经看过其他自定义数据集的实现,但是他们都使用CSV文件,使过程更容易。如果我不能为这类信息创建CSV文件,因为它是一个Map到2D数组的3D数组,我该怎么办?
谢谢你,谢谢你

yr9zkbsy

yr9zkbsy1#

实际上,您应该为这些.mat文件编写自己的Dataset
您可以使用os.listdir来列出子文件夹中的所有文件。torchvision库有几个非常有用的转换/扩充功能可供您使用。具体来说,torchvision.transforms.ToTensor可以将np.array转换为torch.tensors
总的来说,您的自定义Dataset看起来如下所示:

from torch.utils.data import Dataset, IterableDataset, Dataloader

class CTSet(Dataset):
  def __init__(self, base_dir, transforms):
    super(CTSet, self).__init__()
    self.transforms = transforms  # make sure transforms has at leat ToTensor()
    self.data = []  # store all data here
    # go over all files in base_dir
    for file in os.listdir(base_dir):
      if file.endswith('.mat'):
        mat = scipy.io.loadmat(os.path.join(base_dir, file))
        self.data.append((mat['spad'], mat['depth']))

  def __len__(self):
    return len(self.data)

  def __getitem__(self,index):
    spad, depth = self.data[index]
    return self.transforms(spad), self.transform(depth)

相关问题