MIC-DKFZ / batchgenerators

A framework for data augmentation for 2D and 3D image classification and segmentation
Apache License 2.0
1.09k stars 221 forks source link

Random rotations are not uniformly sampled #84

Open Grieverheart opened 3 years ago

Grieverheart commented 3 years ago

Due to the use of Euler angles (yaw-pitch-roll), the rotations are not uniformly distributed. I'm not sure what the motivations is for choosing this convention. This choice can lead to various problems, for example in nnUnet, it is difficult to calculate the patch_size correctly. Presumably, orientations closer to the original one are sampled more, so you need much higher sampling to sample more different orientations. I'm not sure what the impact of this is on learning, but it might be worth checking out.

FabianIsensee commented 3 years ago

Indeed. How would you do this instead? I am not an expert in this at all :-)

maxpietsch commented 3 years ago

There are multiple methods, see for instance: https://mathworld.wolfram.com/SpherePointPicking.html

Since scipy is in your dependencies, you might want to use its implementation of uniformly distributed rotations:

from scipy.spatial.transform import Rotation as R
R.random().as_euler('zxy', degrees=False) # assuming you want to stick to Euler angles and the convention is zxy
FabianIsensee commented 3 years ago

so if we use r = R.random().as_matrix() we get a 3x3 rotation matrix that is uniformly distributed? How can we restrict the angles of rotation around each axis? And how do we get a 2D rotation from that? Sorry if those are stupid questions. I am in the middle of multiple projects and can't find the time right now to do a deep dive into this. I really appreciate your help! Best, Fabian

maxpietsch commented 3 years ago

Yes, according to the documentation, the resulting rotation matrices are uniformly distributed and 10k random rotations of a single fixed vector look uniformly distributed to me.

I don't have a solution that preserves uniformity and works if your bounds are defined in Euler angles. For the 2D case, it is sufficient to produce random angles uniformly in the interval of interest. For 3D,Miles 1965 doi:10.2307/2333716 might contain an answer. You could convert R to Euler angles and randomly scale them so that they fall within the bounds but that might break the uniformity.

Euler angles (and quaternions) are discontinuous (see https://arxiv.org/abs/1812.07035). Interpolation in SO(3) can be performed linearly using the Lie algebra (link). For instance, 100%, 30% and 10% rotations obtained via expm(s * logm(M)) (see (colab)): image

Grieverheart commented 3 years ago

Actually, the method I posted of sampling a uniform unit vector for the axis together with a uniform random angle is wrong. I actually found this out 8 years ago and completely forgot about it.

Here is the answer I got 8 years ago. It uses rejection sampling which is quite efficient for small rotation angles.

Another solution is to generate a uniform random point on the spherical cap (centered on e.g. the x-axis) like here, and take the cross product with the x-axis to get the axis and angle of rotation.

FabianIsensee commented 3 years ago

I recognize that this is an important problem and would like to improve this in batchgenerators. Unfortunately this is really going well over my head. Would it be possible for one of you to contribute a function that, given intervals for rotation around each axis will return a rotation matrix that is uniformly distributed on the unit sphere within the allowed range?

def uniformly_sample_rotation_matrix(angle_range_x: Tuple[float, float], angle_range_y: Tuple[float, float], angle_range_z: Tuple[float, float]) -> np.ndarray:
    XXXX
FabianIsensee commented 3 years ago

How certain are you that the angles of rotation generated by batchgenerators are bogus? when plotting them they look quite good to me: Screenshot from 2021-08-19 14-51-57 Screenshot from 2021-08-19 14-53-14 Screenshot from 2021-08-19 14-53-44

(90 means +- 90 degrees)

(what you see is the result of a (1, 0 ,0) point after being rotated multiple times with random angles)

Here is the script:

import numpy as np
from batchgenerators.augmentations.utils import rotate_coords_3d
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d

if __name__ == '__main__':
    fig = plt.figure()
    ax = plt.axes(projection='3d')

    coordinates = [rotate_coords_3d(np.array((1, 0, 0)),
                                    np.random.uniform(-90 / 180 * np.pi, 90 / 180 * np.pi),
                                    np.random.uniform(-90 / 180 * np.pi, 90 / 180 * np.pi),
                                    np.random.uniform(-90 / 180 * np.pi, 90 / 180 * np.pi)) for i in range(1000)] + [np.array((1, 1, 1))] + [np.array((-1, -1, -1))]
    ax.scatter3D(*np.array(coordinates).transpose(), c='green', s=10)

    coordinates = [rotate_coords_3d(np.array((1, 0, 0)),
                                    np.random.uniform(-45 / 180 * np.pi, 45 / 180 * np.pi),
                                    np.random.uniform(-45 / 180 * np.pi, 45 / 180 * np.pi),
                                    np.random.uniform(-45 / 180 * np.pi, 45 / 180 * np.pi)) for i in range(1000)]
    ax.scatter3D(*np.array(coordinates).transpose(), c='blue', s=10)

    coordinates = [rotate_coords_3d(np.array((1, 0, 0)),
                                    np.random.uniform(-15 / 180 * np.pi, 15 / 180 * np.pi),
                                    np.random.uniform(-15 / 180 * np.pi, 15 / 180 * np.pi),
                                    np.random.uniform(-15 / 180 * np.pi, 15 / 180 * np.pi)) for i in range(500)]
    ax.scatter3D(*np.array(coordinates).transpose(), c='red', s=10)

    ax.legend(['90', '45', '15'])
    plt.show()

EDIT: OK I see a larger concentration of points close to the poles :-/

Grieverheart commented 3 years ago

So, I also investigated further, and although the points are uniformly distributed in what you show above, the boundary of the points is 'square'. See below the difference with actual uniform sampling. Screenshot 2021-08-24 at 19 17 22 Red is with euler angles and blue is uniform.

Here my implementation of uniform sampling of rotation with maximum angle:


def random_axis():
    a = 0
    b = 0
    while True:
        a = 2.0 * np.random.random() - 1.0
        b = 2.0 * np.random.random() - 1.0
        if a**2 + b**2 < 1:
            break

    axis = np.array([
        2.0 * a * np.sqrt(1.0 - a**2 - b**2),
        2.0 * b * np.sqrt(1.0 - a**2 - b**2),
        1.0 - 2.0 * (a**2 + b**2)
    ])

    return axis

def random_uniform_rotation(angle_max):
    while True:
        angle = np.cbrt(np.random.random() * angle_max**3)
        if np.random.random() < (np.sin(angle) / angle)**2:
            break

    axis = random_axis()

    return axis * angle

which returns a rotation vector.

FabianIsensee commented 3 years ago

Hey, thank you so much! I think that the rectangular shape is actually the intended behavior. If the user specifies a max angle of rotation of 15 degrees along all axes then the corners are exactly that maximum.

But I am happy to investigate the sampling of euler angles as well. Maybe I am just dumb - but how do I get a rotation matrix out of the rotation vectors you are returning?

Grieverheart commented 3 years ago

I don't quite understand why anyone would want to specify separate angles. Perhaps something to do with how the patient can be oriented on the table? If that's the intended behaviour, there is probably also an issue with the order of rotations. If you change the order this rectangular shape will be rotated. Note that the shape is not symmetric, the top and bottom is flat, while the sides are not.

About your second question, if you use scipy, you can get the rotation matrix like this:

from scipy.spatial.transform import Rotation as R
rotation_matrix = R.from_rotvec(rotvec).as_matrix()
FabianIsensee commented 3 years ago

Thanks!

batchgenerators is not just for medical images, it is supposed to be a general purpose tool. I agree that having the rectangular is not ideal - it's just how we initially designed it and as you said: the design has flaws. That is why I am happy to look into your approach.

Some use cases might want to be able to specify different angles around different axes. You are of course right that the order of the rotations matters. Not sure what the best solution is here