prepkit.py 文件源码

def randomrotate(angle=90, randomstate=None, invert=False, padding=0, extrapadding=0):
    ignorechannels=True

    if isinstance(randomstate, int):
        rng = np.random.RandomState(randomstate)
    elif isinstance(randomstate, np.random.RandomState):
        rng = randomstate
    else:
        rng = np.random.RandomState(None)

    def _randomrot90(im):
        k = rng.randint(0, 4)
        if not invert:
            return np.rot90(im, k=k)
        else:
            return np.rot90(im, k=(4-k))

    def _randomrot45(im):
        assert im.shape[0] == im.shape[1], "45 degree rotations are tested only for square images."

        k = int(rng.choice(a=[0, 1, 3, 5, 7, 8], size=1))
        # Rotation angle (in degrees)
        rotangle = 45 * k

        if not invert:
            if k == 0 or k == 8:
                im = np.pad(im, (padding + extrapadding), mode='reflect') if padding > 0 else im
                return im
            else:
                im = _rotate(im, angle=rotangle, resize=True, mode='reflect')
                im = np.pad(im, padding, mode='reflect') if padding > 0 else im
                return im
        else:
            if k == 0 or k == 8:
                im = im[(padding + extrapadding):-(padding + extrapadding),
                     (padding + extrapadding):-(padding + extrapadding)] if padding > 0 else im
                return im
            else:
                im = im[padding:-padding, padding:-padding] if padding > 0 else im
                # For some reason, _rotate doesn't like if it's values are larger than +1 or smaller than -1.
                # Scale
                scale = np.max(np.abs(im))
                im *= (1./scale)
                # Process
                im = _rotate(im, angle=(360 - rotangle), resize=True, mode='reflect')
                # Rescale
                im *= scale
                # Edges of im are now twice as large as they were in the original image. Crop.
                cropstart = im.shape[0]/4
                cropstop = cropstart * 3
                im = im[cropstart:cropstop, cropstart:cropstop]
                return im

    if angle == 45:
        return image2batchfunc(_randomrot45, ignorechannels=ignorechannels)
    elif angle == 90:
        return image2batchfunc(_randomrot90, ignorechannels=ignorechannels)
    else:
        raise NotImplementedError("Curently implemented rotation angles are 45 and 90 degrees.")

# Function for random flips of the image