Skip to content

sdf_dataset

Module which provides SDFDataset class.

SDFDataset

Bases: Dataset

Dataset of SDF volumes stored in .npy format.

Expected dataset format

{sdf_folder}/00000.npy {sdf_folder}/00001.npy ...

Source code in sdfest/vae/sdf_dataset.py 
 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
class SDFDataset(torch.utils.data.Dataset):
    """Dataset of SDF volumes stored in .npy format.

    Expected dataset format:
        {sdf_folder}/00000.npy
        {sdf_folder}/00001.npy
        ...
    """

    def __init__(self, sdf_folder: str):
        """Construct the dataset.

        Args:
            sdf_folder: The folder containing the npy files.
        """
        self.path = sdf_folder

        self.size = len([f for f in os.listdir(sdf_folder) if f.endswith(".npy")])

    def __len__(self):
        """Return the number of images in the dataset."""
        return self.size

    def __getitem__(self, idx: int) -> torch.Tensor:
        """Return SDF volume at a specific index.

        Args:
            idx: The index of the sdf file to retrieve.
        Returns:
            The loaded SDF volume.
        """
        sdf_path = os.path.join(self.path, f"{idx:05}.npy")
        sdf_np = np.load(sdf_path)
        return torch.as_tensor(sdf_np).unsqueeze(0)

__getitem__(idx)

Return SDF volume at a specific index.

Parameters:

Name Type Description Default
idx int

The index of the sdf file to retrieve.

 required

Returns: The loaded SDF volume.

Source code in sdfest/vae/sdf_dataset.py 
30
31
32
33
34
35
36
37
38
39
40
def __getitem__(self, idx: int) -> torch.Tensor:
    """Return SDF volume at a specific index.

    Args:
        idx: The index of the sdf file to retrieve.
    Returns:
        The loaded SDF volume.
    """
    sdf_path = os.path.join(self.path, f"{idx:05}.npy")
    sdf_np = np.load(sdf_path)
    return torch.as_tensor(sdf_np).unsqueeze(0)

__init__(sdf_folder)

Construct the dataset.

Parameters:

Name Type Description Default
sdf_folder str

The folder containing the npy files.

 required
Source code in sdfest/vae/sdf_dataset.py 
16
17
18
19
20
21
22
23
24
def __init__(self, sdf_folder: str):
    """Construct the dataset.

    Args:
        sdf_folder: The folder containing the npy files.
    """
    self.path = sdf_folder

    self.size = len([f for f in os.listdir(sdf_folder) if f.endswith(".npy")])

__len__()

Return the number of images in the dataset.

Source code in sdfest/vae/sdf_dataset.py 
26
27
28
def __len__(self):
    """Return the number of images in the dataset."""
    return self.size