import volumentations as V
import numpy as np

augmentation = V.Compose(
    [
        V.Scale3d(scale_limit=(0.2, 0.2, 0.1), p=0.75),
        V.OneOrOther(
            V.Compose(
                [
                    V.RotateAroundAxis3d(
                        rotation_limit=np.pi, axis=(0, 0, 1), always_apply=True
                    ),
                    V.RotateAroundAxis3d(
                        rotation_limit=np.pi / 3, axis=(0, 1, 0), always_apply=True
                    ),
                    V.RotateAroundAxis3d(
                        rotation_limit=np.pi / 3, axis=(1, 0, 0), always_apply=True
                    ),
                ],
                p=1,
            ),
            V.Flip3d(axis=(0, 0, 1)),
        ),
        V.OneOf(
            [
                V.RandomDropout3d(dropout_ratio=0.2, p=0.75),
                V.RandomDropout3d(dropout_ratio=0.3, p=0.5),
            ]
        ),
    ]
)

augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)