Class MultiViewPipeline

embodiedscan/datasets/transforms/multiview.py:10–109

Multiview data processing pipeline. The transform steps are as follows: 1. Select frames. 2. Re-ororganize the selected data structure. 3. Apply transforms for each selected frame. 4. Concatenate data to form a batch. Args: transforms (list[dict | c

@TRANSFORMS.register_module()
class MultiViewPipeline(BaseTransform):
    """Multiview data processing pipeline.

    The transform steps are as follows:

        1. Select frames.
        2. Re-ororganize the selected data structure.
        3. Apply transforms for each selected frame.
        4. Concatenate data to form a batch.

    Args:
        transforms (list[dict | callable]):
            The transforms to be applied to each select frame.
        n_images (int): Number of frames selected per scene.
        ordered (bool): Whether to put these frames in order.
            Defaults to False.
    """

    def __init__(self, transforms, n_images, ordered=False):
        super().__init__()
        self.transforms = Compose(transforms)
        self.n_images = n_images
        self.ordered = ordered

    def transform(self, results: dict) -> dict:
        """Transform function.

        Args:
            results (dict): Result dict from loading pipeline.

        Returns:
            dict: output dict after transformation.
        """
        imgs = []
        img_paths = []
        points = []
        intrinsics = []
        extrinsics = []
        ids = np.arange(len(results['img_path']))
        replace = True if self.n_images > len(ids) else False
        if self.ordered:
            step = (len(ids) - 1) // (self.n_images - 1
                                      )  # TODO: BUG, fix from branch fbocc
            if step > 0:
                ids = ids[::step]
                # sometimes can not get the accurate n_images in this way
                # then take the first n_images one
                ids = ids[:self.n_images]
            else:  # the number of images < pre-set n_images
                # randomly select n_images ids to enable batch-wise inference
                # In practice, can directly use the original ids to avoid
                # redundant computation
                ids = np.random.choice(ids, self.n_images, replace=replace)
        else:
            ids = np.random.choice(ids, self.n_images, replace=replace)
        for i in ids.tolist():
            _results = dict()
            _results['img_path'] = results['img_path'][i]