emmi_inference.models.pipelines.abupt

Classes

AbuptCollator

A Collator that processes the list of samples into a batch in multiple stages:

Module Contents

class emmi_inference.models.pipelines.abupt.AbuptCollator(num_geometry_points, num_surface_anchor_points, num_volume_anchor_points, num_geometry_supernodes, dataset_statistics, use_query_positions=False, seed=None)

Bases: ksuit.data.pipeline.MultiStagePipeline

A Collator that processes the list of samples into a batch in multiple stages:

  • sample_processors: Processing the data before collation on a per-sample level.

  • collators: Conversion from a list of samples into a batch (dict of usually tensors).

  • batch_processors: Processing after collation on a batch-level.

Most of the work is usually done by the sample_processors. One or two collators, and batch processors are often not needed. However this depends on the use case. .. rubric:: Example

>>> sample_processors = [MySampleProcessor1(), MySampleProcessor2()]
>>> collators = [MyCollator1(), MyCollator2()]
>>> batch_processors = [MyBatchProcessor1(), MyBatchProcessor2()]
>>> multistage_pipeline = MultiStagePipeline(
>>>     sample_processors=sample_processors,
>>>     collators=collators,
>>>     batch_processors=batch_processors
>>> )
>>> batch = multistage_pipeline(samples)
Parameters:

  • sample_processors – A list of callables that will be applied sequentially to pre-process on a per-sample level (e.g., subsample a pointcloud).

  • collators – A list of callables that will be applied sequentially to convert the list of individual samples into a batched format. If None, the default PyTorch collator will be used.

  • batch_processors – A list of callables that will be applied sequentially to process on a per-batch level.

  • num_geometry_points (int)

  • num_surface_anchor_points (int)

  • num_volume_anchor_points (int)

  • num_geometry_supernodes (int)

  • dataset_statistics (object)

  • use_query_positions (bool)

  • seed (int | None)

preprocess_inputs_only(samples)