augmentation.py 文件源码

def create_affine_transform_augmentation(img, random_limits=(0.8, 1.1)):
        '''
        Creates an augmentation by computing a homography from three
        points in the image to three randomly generated points
        '''
        y, x = img.shape[:2]
        fx = float(x)
        fy = float(y)
        src_point = np.float32([[fx/2, fy/3,],
                                [2*fx/3, 2*fy/3],
                                [fx/3, 2*fy/3]])
        random_shift = (np.random.rand(3,2) - 0.5) * 2 * (random_limits[1]-random_limits[0])/2 + np.mean(random_limits)
        dst_point = src_point * random_shift.astype(np.float32)
        transform = cv2.getAffineTransform(src_point, dst_point)
        borderValue = 0
        if img.ndim == 3:
            borderValue = np.median(np.reshape(img, (img.shape[0]*img.shape[1],-1)), axis=0)
        else:
            borderValue=np.median(img)
        warped_img = cv2.warpAffine(img, transform, dsize=(x,y), borderValue=borderValue)
        return warped_img