preprocessing.py 文件源码

def rotate_image(image, subjets):
    """Return rotated and repixelised image array.

    Rotation puts subleading subjet or first principle component at -pi/2.
    Repixelisation interpolates with cubic spline.
    """
    # Use subleading subject information to rotate
    if len(subjets) > 1:
        theta = np.arctan2(subjets['phi'][1], subjets['eta'][1])
        theta = -90.0-(theta*180.0/np.pi)
        return transform.rotate(image, theta, order=3)

    # Use principle component of image intensity to rotate
    width, height = image.shape
    pix_coords = np.array([[i, j] for i in range(-width+1, width, 2)
                           for j in range(-height+1, height, 2)])
    covX = np.cov(pix_coords, aweights=np.reshape(image, (width*height)),
                  rowvar=0, bias=1)
    e_vals, e_vecs = np.linalg.eigh(covX)
    pc = e_vecs[:,-1]
    theta = np.arctan2(pc[1], pc[0])
    theta = -90.0-(theta*180.0/np.pi)
    t_image = transform.rotate(image, theta, order=3)
    # Check orientation of principle component
    pix_bot = np.sum(t_image[:, :-(-height//2)])
    pix_top = np.sum(t_image[:, (height//2):])
    if pix_top > pix_bot:
        t_image = transform.rotate(t_image, 180.0, order=3)
        theta += 180.0
    return t_image