def __init__(self,
n_channels,
n_classes,
court_img,
target_size,
court_poi=None,
bilinear=True,
resnet_name='resnetreg50',
resnet_pretrained=None,
warp_with_nearest=False,
img2input=False):
super(Reconstructor, self).__init__()
self.n_channels = n_channels
self.n_classes = n_classes
self.bilinear = bilinear
self.img2input = img2input
# UNet:
self.inc = DoubleConv(n_channels, 64)
self.down1 = Down(64, 128)
self.down2 = Down(128, 256)
self.down3 = Down(256, 512)
factor = 2 if bilinear else 1
self.down4 = Down(512, 1024 // factor)
self.up1 = Up(1024, 512 // factor, bilinear)
self.up2 = Up(512, 256 // factor, bilinear)
self.up3 = Up(256, 128 // factor, bilinear)
self.up4 = Up(128, 64, bilinear)
self.outc = OutConv(64, n_classes)
# UNet regressor that outputs the first 3x3 transformation matrix:
# self.conv_top = nn.Conv2d(1024 // factor, n_classes, kernel_size=1)
# self.unet_reg = Reconstructor.make_regressor(n_classes)
# ResNet regressor that outputs the second 3x3 transformation matrix:
in_classes = n_classes+3 if img2input else n_classes
self.resnet_reg = resnet(resnet_name, resnet_pretrained, in_classes)
# The court template image and court points of interest.
# This court template will be warped by the learnt transformation matrix:
self.court_img = court_img
self.court_poi = court_poi
# STN warper:
h, w = target_size[1], target_size[0]
if warp_with_nearest is False:
self.warper = kornia.HomographyWarper(h, w)
else:
# It seems mode='nearest' has a bug when used during training
self.warper = kornia.HomographyWarper(h, w, mode='nearest')
def __init__(self, args, direc, N, scale_h, scale_w):
super(warp_kornia, self).__init__()
self.warper = dgm.HomographyWarper(int(1208 * scale_h),
int(1920 * scale_w)).to(args.device)
self.MyHomography = Homography(direc, N, scale_h,
scale_w).to(args.device)
def __init__(self, n_channels, n_classes, template, target_size, bilinear=True):
super(CourtReconstruction, self).__init__()
self.n_channels = n_channels
self.n_classes = n_classes
self.bilinear = bilinear
self.inc = DoubleConv(n_channels, 64)
self.down1 = Down(64, 128)
self.down2 = Down(128, 256)
self.down3 = Down(256, 512)
factor = 2 if bilinear else 1
self.down4 = Down(512, 1024 // factor)
self.up1 = Up(1024, 512 // factor, bilinear)
self.up2 = Up(512, 256 // factor, bilinear)
self.up3 = Up(256, 128 // factor, bilinear)
self.up4 = Up(128, 64, bilinear)
self.outc = OutConv(64, n_classes)
# Regressor for the 3x3 transformation matrix:
self.conv_reg = nn.Conv2d(1024 // factor, 8, kernel_size=1)
self.reg = nn.Sequential(
nn.Linear(8 * 22 * 40, 32),
nn.ReLU(True),
nn.Linear(32, 3 * 3)
)
# Initialize the weights/bias with identity transformation:
self.reg[-1].weight.data.zero_()
self.reg[-1].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0, 0, 0, 1], dtype=torch.float))
# Court template that will be warped by the learnt transformation matrix:
self.template = template
# Warper:
h, w = target_size[1], target_size[0]
self.warper = kornia.HomographyWarper(h, w)#, mode='nearest')
def test_gradcheck(self, batch_shape, device, dtype):
# generate input data
eye_size = 3 # identity 3x3
# create checkerboard
patch_src = torch.rand(batch_shape, device=device, dtype=dtype)
patch_src = utils.tensor_to_gradcheck_var(patch_src) # to var
# create base homography
batch_size, _, height, width = patch_src.shape
dst_homo_src = utils.create_eye_batch(batch_size,
eye_size,
device=device,
dtype=dtype)
dst_homo_src = utils.tensor_to_gradcheck_var(
dst_homo_src, requires_grad=False) # to var
# instantiate warper
warper = kornia.HomographyWarper(height, width, align_corners=True)
# evaluate function gradient
assert gradcheck(warper, (
patch_src,
dst_homo_src,
),
raise_exception=True)
def test_smoke(self, device):
img_src_t: torch.Tensor = torch.rand(1, 3, 120, 120).to(device)
img_dst_t: torch.Tensor = torch.rand(1, 3, 120, 120).to(device)
init_homo: torch.Tensor = torch.from_numpy(
np.array([[0.0415, 1.2731, -1.1731], [-0.9094, 0.5072, 0.4272],
[0.0762, 1.3981, 1.0646]])).float()
height, width = img_dst_t.shape[-2:]
warper = kornia.HomographyWarper(height, width)
dst_homo_src = MyHomography(init_homo=init_homo).to(device)
learning_rate = self.lr
optimizer = optim.Adam(dst_homo_src.parameters(), lr=learning_rate)
for _ in range(self.num_iterations):
# warp the reference image to the destiny with current homography
img_src_to_dst = warper(img_src_t, dst_homo_src())
# compute the photometric loss
loss = F.l1_loss(img_src_to_dst, img_dst_t)
optimizer.zero_grad()
loss.backward()
optimizer.step()
assert not bool(torch.isnan(dst_homo_src.h**o.grad).any())
def test_rotation(self, device, batch_shape):
# create input data
batch_size, channels, height, width = batch_shape
patch_src = torch.rand(batch_size, channels, height, width).to(device)
# rotation of 90deg
dst_homo_src = torch.eye(3).to(device)
dst_homo_src[..., 0, 0] = 0.0
dst_homo_src[..., 0, 1] = 1.0
dst_homo_src[..., 1, 0] = -1.0
dst_homo_src[..., 1, 1] = 0.0
dst_homo_src = dst_homo_src.expand(batch_size, -1, -1)
# instantiate warper and warp from source to destination
warper = kornia.HomographyWarper(height, width, align_corners=True)
patch_dst = warper(patch_src, dst_homo_src)
# check the corners
assert_allclose(
patch_src[..., 0, 0], patch_dst[..., 0, -1])
assert_allclose(
patch_src[..., 0, -1], patch_dst[..., -1, -1])
assert_allclose(
patch_src[..., -1, 0], patch_dst[..., 0, 0])
assert_allclose(
patch_src[..., -1, -1], patch_dst[..., -1, 0])
def test_translation(self, shape):
# create input data
offset = 2. # in pixel
height, width = shape
patch_src = torch.rand(1, 1, height, width)
dst_homo_src = utils.create_eye_batch(batch_size=1, eye_size=3)
dst_homo_src[..., 0, 2] = offset / (width - 1) # apply offset in x
# instantiate warper and from source to destination
warper = kornia.HomographyWarper(height, width)
patch_dst = warper(patch_src, dst_homo_src)
assert_allclose(patch_src[..., 1:], patch_dst[..., :-1])
def test_identity(self):
# create input data
height, width = 2, 5
patch_src = torch.rand(1, 1, height, width)
dst_homo_src = utils.create_eye_batch(batch_size=1, eye_size=3)
# instantiate warper
warper = kornia.HomographyWarper(height, width)
# warp from source to destination
patch_dst = warper(patch_src, dst_homo_src)
assert_allclose(patch_src, patch_dst)
def test_homography_warper(self, batch_size, device, dtype):
# generate input data
height, width = 128, 64
eye_size = 3 # identity 3x3
# create checkerboard
board = utils.create_checkerboard(height, width, 4)
patch_src = torch.from_numpy(board).to(
device=device,
dtype=dtype).view(1, 1, height,
width).expand(batch_size, 1, height, width)
# create base homography
dst_homo_src = utils.create_eye_batch(batch_size,
eye_size,
device=device,
dtype=dtype)
# instantiate warper
warper = kornia.HomographyWarper(height, width, align_corners=True)
for i in range(self.num_tests):
# generate homography noise
homo_delta = torch.rand_like(dst_homo_src) * 0.3
dst_homo_src_i = dst_homo_src + homo_delta
# transform the points from dst to ref
patch_dst = warper(patch_src, dst_homo_src_i)
patch_dst_to_src = warper(patch_dst,
_torch_inverse_cast(dst_homo_src_i))
# same transform precomputing the grid
warper.precompute_warp_grid(_torch_inverse_cast(dst_homo_src_i))
patch_dst_to_src_precomputed = warper(patch_dst)
assert (patch_dst_to_src_precomputed == patch_dst_to_src).all()
# projected should be equal as initial
error = utils.compute_patch_error(patch_src, patch_dst_to_src,
height, width)
assert error.item() < self.threshold
# check functional api
patch_dst_to_src_functional = kornia.homography_warp(
patch_dst,
_torch_inverse_cast(dst_homo_src_i), (height, width),
align_corners=True)
assert_allclose(patch_dst_to_src,
patch_dst_to_src_functional,
atol=1e-4,
rtol=1e-4)
def test_translation(self, shape, device, dtype):
# create input data
offset = 2.0 # in pixel
height, width = shape
patch_src = torch.rand(1, 1, height, width, device=device, dtype=dtype)
dst_homo_src = utils.create_eye_batch(batch_size=1, eye_size=3, device=device, dtype=dtype)
dst_homo_src[..., 0, 2] = offset / (width - 1) # apply offset in x
# instantiate warper and from source to destination
warper = kornia.HomographyWarper(height, width, align_corners=True)
patch_dst = warper(patch_src, dst_homo_src)
assert_close(patch_src[..., 1:], patch_dst[..., :-1], atol=1e-4, rtol=1e-4)
def test_identity(self, device, dtype):
# create input data
height, width = 2, 5
patch_src = torch.rand(1, 1, height, width, device=device, dtype=dtype)
dst_homo_src = utils.create_eye_batch(batch_size=1, eye_size=3, device=device, dtype=dtype)
# instantiate warper
warper = kornia.HomographyWarper(height, width, align_corners=True)
# warp from source to destination
patch_dst = warper(patch_src, dst_homo_src)
assert_close(patch_src, patch_dst)
def forward(self, img_src, direc, scale):
height, width = img_src.shape[-2:]
self.warper = dgm.HomographyWarper(height, width)
self.MyHomography = Homography(direc,
height,
width,
N=img_src.shape[0],
scale=scale)
img_src_tensor = img_src / 255.
warped_img_tensor = 255. * self.warper(img_src_tensor,
self.MyHomography())
return warped_img_tensor
def test_identity_resize(self, batch_shape):
# create input data
batch_size, channels, height, width = batch_shape
patch_src = torch.rand(batch_size, channels, height, width)
dst_homo_src = utils.create_eye_batch(batch_size, eye_size=3)
# instantiate warper warp from source to destination
warper = kornia.HomographyWarper(height // 2, width // 2)
patch_dst = warper(patch_src, dst_homo_src)
# check the corners
assert_allclose(patch_src[..., 0, 0], patch_dst[..., 0, 0])
assert_allclose(patch_src[..., 0, -1], patch_dst[..., 0, -1])
assert_allclose(patch_src[..., -1, 0], patch_dst[..., -1, 0])
assert_allclose(patch_src[..., -1, -1], patch_dst[..., -1, -1])
def __init__(self, window_size=512, border=32):
super(AlignLoss, self).__init__()
self.window_size = window_size
self.border = border
self.warper = tgm.HomographyWarper(window_size,
window_size,
normalized_coordinates=True,
mode="nearest")
self.L1_criterion = torch.nn.MSELoss()
self.L1_criterion = self.L1_criterion.cuda()
self.L2_criterion = torch.nn.L1Loss()
self.L2_criterion = self.L2_criterion.cuda()
def test_identity_resize(self, batch_shape, device, dtype):
# create input data
batch_size, channels, height, width = batch_shape
patch_src = torch.rand(batch_size, channels, height, width, device=device, dtype=dtype)
dst_homo_src = utils.create_eye_batch(batch_size, eye_size=3, device=device, dtype=dtype)
# instantiate warper warp from source to destination
warper = kornia.HomographyWarper(height // 2, width // 2, align_corners=True)
patch_dst = warper(patch_src, dst_homo_src)
# check the corners
assert_close(patch_src[..., 0, 0], patch_dst[..., 0, 0], atol=1e-4, rtol=1e-4)
assert_close(patch_src[..., 0, -1], patch_dst[..., 0, -1], atol=1e-4, rtol=1e-4)
assert_close(patch_src[..., -1, 0], patch_dst[..., -1, 0], atol=1e-4, rtol=1e-4)
assert_close(patch_src[..., -1, -1], patch_dst[..., -1, -1], atol=1e-4, rtol=1e-4)
def test_warp_grid_translation(self, shape, offset, device):
# create input data
height, width = shape
dst_homo_src = utils.create_eye_batch(batch_size=1,
eye_size=3).to(device)
dst_homo_src[..., 0, 2] = offset # apply offset in x
# instantiate warper
warper = kornia.HomographyWarper(height,
width,
normalized_coordinates=False)
flow = warper.warp_grid(dst_homo_src)
# the grid the src plus the offset should be equal to the flow
# on the x-axis, y-axis remains the same.
assert_allclose(warper.grid[..., 0].to(device) + offset, flow[..., 0])
assert_allclose(warper.grid[..., 1].to(device), flow[..., 1])
def test_rotation(self, batch_shape):
# create input data
batch_size, channels, height, width = batch_shape
patch_src = torch.rand(batch_size, channels, height, width)
# rotation of 90deg
dst_homo_src = utils.create_eye_batch(batch_size, 3)
dst_homo_src[..., 0, 0] = 0.0
dst_homo_src[..., 0, 1] = 1.0
dst_homo_src[..., 1, 0] = -1.0
dst_homo_src[..., 1, 1] = 0.0
# instantiate warper and warp from source to destination
warper = kornia.HomographyWarper(height, width)
patch_dst = warper(patch_src, dst_homo_src)
# check the corners
assert_allclose(patch_src[..., 0, 0], patch_dst[..., 0, -1])
assert_allclose(patch_src[..., 0, -1], patch_dst[..., -1, -1])
assert_allclose(patch_src[..., -1, 0], patch_dst[..., 0, 0])
assert_allclose(patch_src[..., -1, -1], patch_dst[..., -1, 0])
def test_homography_warper(self, batch_size, device_type):
# generate input data
height, width = 128, 64
eye_size = 3 # identity 3x3
device = torch.device(device_type)
# create checkerboard
board = utils.create_checkerboard(height, width, 4)
patch_src = torch.from_numpy(board).view(1, 1, height, width).expand(
batch_size, 1, height, width)
patch_src = patch_src.to(device)
# create base homography
dst_homo_src = utils.create_eye_batch(batch_size, eye_size).to(device)
# instantiate warper
warper = kornia.HomographyWarper(height, width)
for i in range(self.num_tests):
# generate homography noise
homo_delta = torch.zeros_like(dst_homo_src)
homo_delta[:, -1, -1] = 0.0
dst_homo_src_i = dst_homo_src + homo_delta
# transform the points from dst to ref
patch_dst = warper(patch_src, dst_homo_src_i)
patch_dst_to_src = warper(patch_dst, torch.inverse(dst_homo_src_i))
# projected should be equal as initial
error = utils.compute_patch_error(patch_dst, patch_dst_to_src,
height, width)
assert error.item() < self.threshold
# check functional api
patch_dst_to_src_functional = kornia.homography_warp(
patch_dst, torch.inverse(dst_homo_src_i), (height, width))
assert utils.check_equal_torch(patch_dst_to_src,
patch_dst_to_src_functional)
def photometric_loss(H_norm, img_fix, img_move):
b, c, w, h = img_fix.shape
dev = img_fix.device
# Since we assumed corner point coordinate is -1~1, Normalization will be given as follows:
# N = torch.FloatTensor([[2/(w-1),0,-1],[0,2/(h-1),-1],[0,0,1]]).unsqueeze(0)
# N_inv = torch.inverse(N)
# # To compensate the negative movement,we should translate image.
# # T = torch.FloatTensor([[1,0,w],[0,1,h],[0,0,1]]).unsqueeze(0)
# # Matrix should be have same batch size
# N = torch.cat([N]*b).to(dev)
# N_inv = torch.cat([N_inv]*b).to(dev)
# # T = torch.cat([T]*b).to(dev)
# # Denormalize
# H = torch.einsum('bij,bjk->bik',N_inv,torch.einsum('bij,bjk->bik',H_norm,N))
# # Compensate the translation
# # H = torch.einsum('bij,bjk->bik',T,H)
# # fixed image has only translation
# # img_fix_transformed = kornia.warp_perspective(img_fix,T,(h*3,w*3))
# img_fix_transformed = img_fix
# # warp the second image
# # img_move_transformed = kornia.warp_perspective(img_move,H,(h*3,w*3))
# img_move_transformed = kornia.warp_perspective(img_move,H,(h,w))
# # img_move_transformed[img_move_transformed != img_move_transformed]=0
warper = kornia.HomographyWarper(h, w)
img_move_transformed = warper(img_move, torch.inverse(H_norm))
img_fix_transformed = img_fix
# Compute L1 loss
loss = F.l1_loss(img_move_transformed,
img_fix_transformed,
reduction='sum')
return loss / (2 * h * w), img_move_transformed
def HomographyRegressionApp():
# Training settings
parser = argparse.ArgumentParser(
description='Homography Regression with photometric loss.')
parser.add_argument('--input-dir',
type=str,
required=True,
help='the path to the directory with the input data.')
parser.add_argument('--output-dir',
type=str,
required=True,
help='the path to output the results.')
parser.add_argument('--num-iterations',
type=int,
default=1000,
metavar='N',
help='number of training iterations (default: 1000)')
parser.add_argument('--lr',
type=float,
default=1e-3,
metavar='LR',
help='learning rate (default: 1e-3)')
parser.add_argument('--cuda',
action='store_true',
default=False,
help='enables CUDA training')
parser.add_argument('--seed',
type=int,
default=666,
metavar='S',
help='random seed (default: 666)')
parser.add_argument(
'--log-interval',
type=int,
default=10,
metavar='N',
help='how many batches to wait before logging training status',
)
parser.add_argument(
'--log-interval-vis',
type=int,
default=100,
metavar='N',
help='how many batches to wait before visual logging training status',
)
args = parser.parse_args()
# define the device to use for inference
use_cuda = args.cuda and torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')
torch.manual_seed(args.seed)
# load the data
img_src, _ = load_image(os.path.join(args.input_dir, 'img1.ppm'))
img_dst, _ = load_image(os.path.join(args.input_dir, 'img2.ppm'))
# instantiate the homography warper from `kornia`
height, width = img_src.shape[-2:]
warper = dgm.HomographyWarper(height, width)
# create the homography as the parameter to be optimized
dst_homo_src = MyHomography().to(device)
# create optimizer
optimizer = optim.Adam(dst_homo_src.parameters(), lr=args.lr)
# main training loop
for iter_idx in range(args.num_iterations):
# send data to device
img_src, img_dst = img_src.to(device), img_dst.to(device)
# warp the reference image to the destiny with current homography
img_src_to_dst = warper(img_src, dst_homo_src())
# compute the photometric loss
loss = F.l1_loss(img_src_to_dst, img_dst, reduction='none')
# propagate the error just for a fixed window
w_size = 100 # window size
h_2, w_2 = height // 2, width // 2
loss = loss[..., h_2 - w_size:h_2 + w_size, w_2 - w_size:w_2 + w_size]
loss = torch.mean(loss)
# compute gradient and update optimizer parameters
optimizer.zero_grad()
loss.backward()
optimizer.step()
if iter_idx % args.log_interval == 0:
print(
f'Train iteration: {iter_idx}/{args.num_iterations}\tLoss: {loss.item():.6}'
)
print(dst_homo_src.h**o)
def draw_rectangle(image, dst_homo_src):
height, width = image.shape[:2]
pts_src = torch.FloatTensor(
[[[-1, -1], [1, -1], [1, 1],
[-1,
1]]] # top-left # bottom-left # bottom-right # top-right
).to(dst_homo_src.device)
# transform points
pts_dst = dgm.transform_points(torch.inverse(dst_homo_src),
pts_src)
def compute_factor(size):
return 1.0 * size / 2
def convert_coordinates_to_pixel(coordinates, factor):
return factor * (coordinates + 1.0)
# compute conversion factor
x_factor = compute_factor(width - 1)
y_factor = compute_factor(height - 1)
pts_dst = pts_dst.cpu().squeeze().detach().numpy()
pts_dst[...,
0] = convert_coordinates_to_pixel(pts_dst[..., 0],
x_factor)
pts_dst[...,
1] = convert_coordinates_to_pixel(pts_dst[..., 1],
y_factor)
# do the actual drawing
for i in range(4):
pt_i, pt_ii = tuple(pts_dst[i % 4]), tuple(pts_dst[(i + 1) %
4])
image = cv2.line(image, pt_i, pt_ii, (255, 0, 0), 3)
return image
if iter_idx % args.log_interval_vis == 0:
# merge warped and target image for visualization
img_src_to_dst = warper(img_src, dst_homo_src())
img_vis = 255.0 * 0.5 * (img_src_to_dst + img_dst)
img_vis_np = dgm.utils.tensor_to_image(img_vis)
image_draw = draw_rectangle(img_vis_np, dst_homo_src())
# save warped image to disk
file_name = os.path.join(args.output_dir, f'warped_{iter_idx}.png')
cv2.imwrite(file_name, image_draw)
