DIST -> DataLoader

General DataLoader

General class for creating Datasets. It works by detecting all specified file formats in the given directory and return the index list.

class tardis_em.dist_pytorch.datasets.dataloader.BasicDataset(coord_dir=None, coord_format='.csv', patch_if=500, downscale=None, rgb=False, benchmark=False, train=True)

BasicDataset is a dataset handling class, mainly designed to manage and preprocess coordinate data for tasks requiring large-scale numerical data handling. This class initializes the dataset environment, processes input data, and provides several utility functions to save, load, and transform data efficiently.

This class is particularly useful for managing coordinate files, setting downscaled versions of files, and utility methods for temporary storage of preprocessed data.

save_temp(i: int, **kwargs)

Saves temporary data to numpy files.

This method takes an integer identifier i and saves temporary data provided in kwargs to .npy files within a specified directory. Files are named using the format <key>_<i>.npy.

Parameters:

  • i – An integer used in naming the saved files.

  • kwargs – Keyword arguments where keys represent file name prefixes, and values are the data to be saved in the numpy files. Each value is converted to a numpy array with object dtype before saving.

Returns:

None

load_temp(i: int, **kwargs) → List[ndarray]

Loads temporary numpy array files from the disk. The function can load a single file or multiple files depending on the provided keyword arguments. If one keyword argument is given, it will load the corresponding file directly. Otherwise, it will load all files corresponding to the provided keyword arguments in a list.

Parameters:

  • i (int) – Index used to identify the file(s) to be loaded.

  • kwargs (dict) – A set of keyword arguments where the keys represent the prefix of the filename(s) to be loaded. Each key corresponds to a possible file.

Returns:

A list of numpy ndarray objects if multiple keyword arguments are provided, or a single numpy ndarray object if only one keyword argument is given.

Return type:

List[np.ndarray]

static list_to_tensor(**kwargs) → List

Converts a list of numerical dataframes into a list of PyTorch tensors. This method takes variable keyword arguments, where each value is expected to be a list of dataframes. It processes each dataframe in these lists to convert them into PyTorch tensors with a float32 data type, preserving the structure of the input.

Parameters:

kwargs – Variable keyword arguments where each key corresponds to a list of pandas DataFrames. Each dataframe is processed and converted to a PyTorch tensor with dtype set to float32.

Returns:

A list of lists containing the converted PyTorch tensors. The structure of the returned list parallels the structure of the input keyword arguments.

Return type:

List

Specialized DataLoader’s

Filament structure DataSet

class tardis_em.dist_pytorch.datasets.dataloader.FilamentDataset(**kwargs)

FilamentDataset class for handling filament data processing.

This class is designed to process filament datasets, including pre-processing of coordinate data, image patch creation, and handling normalized data points. It is useful for machine learning tasks requiring patched datasets, graph representations, and customized preprocessing.

PartNet synthetic DataSet

class tardis_em.dist_pytorch.datasets.dataloader.PartnetDataset(**kwargs)

PartnetDataset is a specialized dataset class designed to work with 3D point clouds and graph representations. It integrates functionalities to preprocess, manage, and retrieve data efficiently. This includes handling downscaling, patch generation, and dynamic data loading for both training and testing.

The class leverages the PatchDataSet to create patches of data points with specified parameters like maximum points, overlap, and drop rate. Additionally, it supports temporary storage and loading for optimized access.

ScanNet V2 synthetic DataSet

class tardis_em.dist_pytorch.datasets.dataloader.ScannetDataset(**kwargs)

ScannetDataset handles loading and processing of 3D scan data, specifically for tasks involving graph and patch-based methods in the context of 3D scene understanding.

This class is a specialized dataset loader that extends the BasicDataset. It can handle a range of tasks including preprocessing of scan data, handling coordinate formats, and managing batch-level operations for training and testing workflows. The class also supports saving and retrieving preprocessed temporary data for faster data loading.

ScanNet V2 with RGB values synthetic DataSet

class tardis_em.dist_pytorch.datasets.dataloader.ScannetColorDataset(**kwargs)

Represents a dataset for Scannet, inheriting functionalities from the BasicDataset. It supports handling of color information, coordination files, and patch-based data processing.

This dataset class processes 3D point cloud data and corresponding color information from Scannet using various custom functionalities, including coordinate scaling, patch-based segmentation, and preprocessed data caching for efficient usage. It also incorporates graph-based dataset representations.

Stanford S3DIS DataSet

class tardis_em.dist_pytorch.datasets.dataloader.Stanford3DDataset(**kwargs)

This class represents the Stanford 3D dataset, which handles dataset processing, room identification, and patch data preparation. It is responsible for organizing room data from directories and subdirectories, preprocessing 3D point cloud data for machine learning tasks, and optimizing computational resources by leveraging patches. The class integrates with the PatchDataSet for patch extraction and manipulation.

Helper Functions

Build dataloader

tardis_em.dist_pytorch.datasets.dataloader.build_dataset(dataset_type: str | list, dirs: list, max_points_per_patch: int, downscale=None, benchmark=False)

Build dataset objects and corresponding data loaders based on the specified dataset type and configurations. Supports various dataset types for training and testing such as simulated datasets, Filament datasets, PartNet, ScanNet, Stanford datasets, and more. The function also handles benchmark mode for testing datasets only. Depending on the dataset_type, appropriate dataset classes are instantiated, configured with training and testing attributes, and wrapped in data loaders for model consumption.

Parameters:

  • dataset_type – Specifies the type of dataset to load. Can be either a string or a list. If it is a list, it must define simulation-related metadata. For strings, supported values include “filament”, “MT”, “Mem”, “partnet”, “scannet”, “stanford”, and their variations (e.g., “scannet_rgb”, “stanford_rgb”).

  • dirs – A list containing directory paths to load training and testing dataset files from. The directories depend on the specific type of dataset being used.

  • max_points_per_patch – The maximum number of points that a single patch within the dataset can contain. This value is used to configure dataset-specific patch loading and processing.

  • downscale – Defines optional downscaling operations to be applied during dataset loading to alter point cloud resolution. If not provided, defaults are used based on the dataset type.

  • benchmark – If set to True, indicates benchmark testing mode. In this mode, only the test dataset loader is returned, and the training loader setup is bypassed.

Returns:

For non-benchmark mode, returns a tuple where the first element is the training dataset loader instance and the second element is the testing dataset loader instance. For benchmark mode, only the testing dataset loader is returned.

Point Cloud Augmentation

tardis_em.dist_pytorch.datasets.augmentation.preprocess_data(coord: str | ndarray, image: str | None = None, size: int | None = None, include_label=True, normalization='simple') → Tuple[ndarray, ndarray] | Tuple[ndarray, ndarray, ndarray]

Preprocesses coordinate and image data for further analysis.

This function reads and processes geometric coordinate data and associated image data, applying normalization and cropping based on specified parameters. It supports various file formats for input coordinates (.csv, .npy, .am) and optionally uses the image data aligned with these coordinates. The processed output can include labels, graphs, or specific image subsets, depending on the provided parameters.

Parameters:

  • coord (Union[str, np.ndarray]) – Required input coordinates in either raw array format or file paths (strings) pointing to supported formats like .csv, .npy, or .am files. Depending on the format, data is loaded and processed for subsequent use.

  • image (Optional[str]) – Optional input image path corresponding to the given coordinates. If provided, image patches are extracted and normalized around coordinate positions. It can also be a .am file for specific processing.

  • size (Optional[int]) – Optional size of image patches to be cropped. If not provided, a default size of 1 is used. Determines whether the cropping is applied in 2D or 3D, based on the dimensions of the provided coordinates.

  • include_label (bool) – Indicates whether to return labels along with coordinates and image patches. If False, an additional graph is created from coordinates.

  • normalization (str) – Specifies the normalization method for image data. Accepts “simple” or “minmax” as methods or None for no normalization. Default is “simple”.

Returns:

Returns a tuple that may consist of coordinates with labels, images, or a graph depending on the value of include_label. If include_label is True, a tuple of coordinates, labels, and images is returned. Otherwise, a graph of coordinates and images is included in the returned tuple.

Return type:

Union[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, np.ndarray]]

class tardis_em.dist_pytorch.datasets.augmentation.BuildGraph(K=2, mesh=False)

Creates a representation of points as a graph where each edge represents the connectivity between points based on their proximity or contour sequence. Depending on the configuration, the graph is constructed either through Nearest Neighbors for mesh-based structures or explicit connections for points belonging to the same contour. The graph is represented as a binary matrix, with optional handling of self-connections.

class tardis_em.dist_pytorch.datasets.augmentation.Crop2D3D(image: ndarray, size: tuple, normalization)

Class for 2D or 3D image cropping based on a center position and size.

This class facilitates cropping sections of an image (both 2D and 3D depending on the input) by specifying a center position and dimensions for the crop. It supports normalization of the cropped section if a normalization function is provided.

static get_xyz_position(center_point: int, size: int, max_size: int) → Tuple[int, int]

Calculate the adjusted x-coordinate positions for a frame, ensuring that the frame fits within the bounds of a specified maximum size.

This method adjusts the lower and upper bounds of the frame’s position based on its size, ensuring it remains within the allowed range (max_size). If the computed position exceeds boundaries, it shifts the frame appropriately.

Parameters:

  • center_point – The center point of the frame, around which the lower and upper x-coordinate positions are calculated.

  • size – The size of the frame on the x-axis, which will be distributed evenly around the center point.

  • max_size – The maximum allowable boundary in the x-axis that the frame must stay within.

Returns:

A tuple containing the adjusted lower and upper x-coordinate positions, ensuring the frame lies within the bounds.

Return type:

Tuple[int, int]

tardis_em.dist_pytorch.datasets.augmentation.upsample_pc(org_coord: ndarray, sampled_coord: ndarray)

Upsample the point cloud by associating each original coordinate (org_coord) with its nearest neighbor in the sampled coordinate set (sampled_coord). The function finds the nearest neighbor using a KDTree for efficient spatial querying and computes the new coordinates based on these associations.

Parameters:

  • org_coord – A numpy array of shape (n, m), where n is the number of points and m represents the dimensionality of the space. The first column represents unique identifiers, and the remaining columns represent the coordinates of the points.

  • sampled_coord – A numpy array of shape (k, m), where k is the number of sampled points and m represents the dimensionality of the space. The first column represents unique identifiers, and the remaining columns represent the coordinates of the points.

Returns:

A numpy array of shape (n, m-1) representing the upsampled point cloud. The first column contains the identifiers of the nearest neighbors from the sampled_coord array, and the remaining columns replicate the coordinate dimensions of the input org_coord.

Patch dataset

class tardis_em.dist_pytorch.datasets.patches.PatchDataSet(max_number_of_points=500, overlap=0.15, drop_rate=None, graph=True, tensor=True)

Represents a class designed for configuring and processing point cloud data. This includes functionalities for down-sampling points, defining overlapping feature patches, and specifying output formats such as tensor representation or graph structure. It also initializes and manages patch size configurations.

boundary_box(coord, offset=None) → ndarray

Calculate the boundary box for the given coordinates, optionally applying an offset. The function computes the minimum and maximum dimensions along each axis (x, y, z if applicable). If the offset is provided, the boundary box is expanded proportionally around the center. The result is a NumPy array representing the boundary box with two points: the lower and upper bounds.

Parameters:

  • coord (np.ndarray) – The input array of coordinates where rows represent points, and columns represent dimensional axes (x, y, and optionally z).

  • offset (float, optional) – An optional expansion multiplier for the bounding box dimensions. If provided, it scales the boundary box’s size relative to its center.

Returns:

A NumPy array representing the boundary box with two rows: the minimum (x, y, z) coordinates and the maximum (x, y, z) coordinates.

Return type:

np.ndarray

static center_patch(bbox, voxel_size=1) → ndarray

Calculates the coordinates for the centers of voxels within a given bounding box. The bounding box is divided into a three-dimensional grid using the specified voxel size. The method ensures a minimum of 2 voxels along each axis to avoid degenerate cases. The voxel centers are returned as a NumPy array containing their coordinates.

Parameters:

  • bbox (np.ndarray) – The bounding box defined as a 2x3 NumPy array where the first row contains the minimum x, y, z coordinates and the second row contains the maximum x, y, z coordinates.

  • voxel_size (float) – The size of each voxel along all dimensions. Defaults to 1.

Returns:

A NumPy array with the coordinates of all voxel centers as rows.

Return type:

np.ndarray

points_in_patch(coord: ndarray, patch_center: ndarray) → bool

Determines if specific points represented by their coordinates are located within a defined patch based on the patch’s center and pre-defined size. The patch is modeled as a cuboid, and each coordinate is verified to be within the patch bounds along all axes.

Parameters:

  • coord – A numpy array of shape (N, 3), where N is the number of points, and each row represents the x, y, z coordinates of a point.

  • patch_center – A numpy array of shape (3,), representing the x, y, z coordinates of the patch center.

Returns:

A boolean array of shape (N,), where each element signifies whether the corresponding point in coord is located within the defined patch limits.

Return type:

np.ndarray

optimal_patches(coord: ndarray, random=False) → List[bool]

The optimal_patches function calculates and returns a list of boolean values representing whether specific patches within a given coordinate space contain points that meet certain criteria. It operates in two modes: random patch selection or systematic grid-based patch calculation.

This function first calculates a bounding box for the input coordinates. It determines voxel sizes dynamically based on the bounding box dimensions, applies strategies for patch creation, and evaluates patches for spatial distribution of points. The two modes allow either randomly selected patches or systematic patches using voxel downscaling until a threshold is satisfied.

Parameters:

  • coord (np.ndarray) – The spatial coordinate input as a numpy array.

  • random (bool, optional) – Determines if the patch selection should follow a random process or systematic method. Default is False.

Returns:

A list of boolean values indicating which patches meet the specific criteria.

Return type:

List[bool]

static normalize_idx(coord_with_idx: ndarray) → ndarray

Normalizes the index values of coordinates in the given array. The method updates the array such that indices in the first column are replaced with normalized indices, starting from zero. The normalization process maps unique indices to a consecutive range of indices.

Parameters:

coord_with_idx (np.ndarray) – A numpy array where the first column contains coordinate indices to be normalized.

Returns:

The input array with normalized indices in its first column.

Return type:

np.ndarray

output_format(data: ndarray) → ndarray | Tensor

Converts the given data into the desired output format. If the TORCH_OUTPUT flag is set, the input NumPy array is converted to a PyTorch tensor of type torch.float32.

Parameters:

data (np.ndarray) – Input data as a NumPy array.

Returns:

The input data converted to the desired format. Returns a PyTorch tensor if TORCH_OUTPUT is enabled, otherwise returns the original NumPy array.

Return type:

Union[np.ndarray, torch.Tensor]

patched_dataset(coord: ndarray, label_cls=None, rgb=None, mesh=6, random=False, voxel_size=None) → Tuple[List, List, List, List, List] | Tuple[List, List, List, List]

Generates patches of spatial data and optional associated properties (class labels, graph structures, and RGB values) for deep learning tasks. This function processes input spatial coordinates to divide them into smaller, manageable patches, and conditionally builds graph structures or assigns additional information like class labels and RGB values for each patch. This method supports either a 2D-to-3D transformation or segmented and raw data processing, depending on the provided input configurations.

Parameters:

  • coord – An array of spatial coordinates, potentially in 2D, 3D, or segmented 4D form.

  • label_cls – Optional array of class labels for spatial coordinates.

  • rgb – Optional array of RGB values for spatial coordinates.

  • mesh – An integer indicating the scale or structure of graph connectivity (e.g., K in a K-Nearest Neighbors graph).

  • random – Boolean indicating whether to randomly segment patches or determine them based on an optimization algorithm.

  • voxel_size – Optional numeric value specifying the desired voxel size for downsampling spatial data.

Returns:

A tuple of lists with various outputs segmented into patches: - Point cloud patches for input coordinates. - RGB data patches if available. - Graph patches if graph output is configured. - Indices of original data points corresponding to each patch. - Optional class label patches if label_cls is provided.