def __getitem__()

in pytorch3d/structures/pointclouds.py [0:0]

• 37 lines of code
• 24 McCabe index (conditional complexity)

    def __getitem__(self, index) -> "Pointclouds":
        """
        Args:
            index: Specifying the index of the cloud to retrieve.
                Can be an int, slice, list of ints or a boolean tensor.

        Returns:
            Pointclouds object with selected clouds. The tensors are not cloned.
        """
        normals, features = None, None
        normals_list = self.normals_list()
        features_list = self.features_list()
        if isinstance(index, int):
            points = [self.points_list()[index]]
            if normals_list is not None:
                normals = [normals_list[index]]
            if features_list is not None:
                features = [features_list[index]]
        elif isinstance(index, slice):
            points = self.points_list()[index]
            if normals_list is not None:
                normals = normals_list[index]
            if features_list is not None:
                features = features_list[index]
        elif isinstance(index, list):
            points = [self.points_list()[i] for i in index]
            if normals_list is not None:
                normals = [normals_list[i] for i in index]
            if features_list is not None:
                features = [features_list[i] for i in index]
        elif isinstance(index, torch.Tensor):
            if index.dim() != 1 or index.dtype.is_floating_point:
                raise IndexError(index)
            # NOTE consider converting index to cpu for efficiency
            if index.dtype == torch.bool:
                # advanced indexing on a single dimension
                index = index.nonzero()
                index = index.squeeze(1) if index.numel() > 0 else index
                index = index.tolist()
            points = [self.points_list()[i] for i in index]
            if normals_list is not None:
                normals = [normals_list[i] for i in index]
            if features_list is not None:
                features = [features_list[i] for i in index]
        else:
            raise IndexError(index)

        return self.__class__(points=points, normals=normals, features=features)