33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
import numpy as np
|
|
import cv2
|
|
|
|
def get_random_patch(image, patch_size):
|
|
h, w = image.shape[1:3]
|
|
x = np.random.randint(0, w - patch_size)
|
|
y = np.random.randint(0, h - patch_size)
|
|
return image[:, y:y+patch_size, x:x+patch_size]
|
|
|
|
def get_random_homography(max_rotation=30, max_translation=20):
|
|
theta = np.random.uniform(-max_rotation, max_rotation) * np.pi / 180.0
|
|
tx = np.random.uniform(-max_translation, max_translation)
|
|
ty = np.random.uniform(-max_translation, max_translation)
|
|
cos_theta = np.cos(theta)
|
|
sin_theta = np.sin(theta)
|
|
H = np.array([
|
|
[cos_theta, -sin_theta, tx],
|
|
[sin_theta, cos_theta, ty],
|
|
[0, 0, 1]
|
|
])
|
|
return H
|
|
|
|
def apply_homography_to_image(image, H, output_size):
|
|
warped = np.zeros_like(image)
|
|
for c in range(image.shape[0]):
|
|
warped[c] = cv2.warpPerspective(image[c], H, output_size, flags=cv2.INTER_NEAREST)
|
|
return warped
|
|
|
|
def generate_training_pair(image, patch_size):
|
|
patch = get_random_patch(image, patch_size)
|
|
H = get_random_homography()
|
|
warped_patch = apply_homography_to_image(patch, H, (patch_size, patch_size))
|
|
return patch, warped_patch, H |