def load_img(path: str,
height: int = 256,
width: int = 256,
bits: int = 8,
plot: bool = False,
crop_mode: str = "centre-crop",
save_gt: bool = False,
**kwargs) -> torch.Tensor:
img = cv2.imread(path, -1)[:, :, ::-1] / (2**bits - 1)
img = torch.from_numpy(img.copy()).float().permute(2, 0, 1)
if crop_mode == "resize-crop":
# Resize such that shorter side matches corresponding target side
smaller_side = min(height, width)
img = kornia.resize(img.unsqueeze(0),
smaller_side,
align_corners=False).squeeze(0)
img = kornia.center_crop(img.unsqueeze(0), (height, width),
align_corners=False)
img = img.squeeze(0).permute(1, 2, 0)
if plot:
plt.imshow(img)
plt.show()
if save_gt:
cv2.imwrite("gt.png", img.numpy()[:, :, ::-1] * 255.0)
# H x W x 3
return img
def forward(self, images):
batch_size = images.shape[0]
data = []
for i in range(6):
img_warp = kornia.warp_perspective(
images[:, i, :, :, :],
self.M_matrices[i].unsqueeze(0).repeat(batch_size, 1, 1),
dsize=(204, 306))
img_warp = kornia.center_crop(img_warp, (192, 288))
out = self.darknet(img_warp)
data.append(out.unsqueeze(0))
agg = torch.cat(data, dim=0)
agg = torch.max(agg, dim=0)[0]
agg = agg.view(agg.size(0), 1024, -1)
agg = self.lin1(agg)
boxes = agg.view(agg.size(0), -1)
boxes = self.lin2(boxes)
boxes = self.classifier(boxes)
boxes = boxes.view(-1, self.feature_size, self.feature_size,
5 * self.num_bboxes)
return boxes
def test_center_crop_h4_w2_batch(self, device):
inp = torch.tensor([[
[1., 2., 3., 4.],
[5., 6., 7., 8.],
[9., 10., 11., 12.],
[13., 14., 15., 16.],
],
[
[1., 5., 9., 13.],
[2., 6., 10., 14.],
[3., 7., 11., 15.],
[4., 8., 12., 16.],
]]).to(device)
height, width = 4, 2
expected = torch.tensor([[
[2., 3.],
[6., 7.],
[10., 11.],
[14., 15.],
], [
[5., 9.],
[6., 10.],
[7., 11.],
[8., 12.],
]]).to(device)
out_crop = kornia.center_crop(inp, (height, width))
assert_allclose(out_crop, expected)
def test_jit_trace(self, device):
@torch.jit.script
def op_script(input: torch.Tensor,
size: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
return kornia.center_crop(input, size)
# 1. Trace op
batch_size, channels, height, width = 1, 2, 5, 4
img = torch.ones(batch_size, channels, height, width).to(device)
crop_height, crop_width = 4, 2
op_trace = torch.jit.trace(
op_script,
(img, (torch.tensor(crop_height), torch.tensor(crop_width))))
# 2. Generate new input
batch_size, channels, height, width = 2, 1, 6, 3
img = torch.ones(batch_size, channels, height, width).to(device)
# 3. Evaluate
crop_height, crop_width = 2, 3
actual = op_trace(
img, (torch.tensor(crop_height), torch.tensor(crop_width)))
expected = kornia.center_crop(img, (crop_height, crop_width))
assert_allclose(actual, expected)
def test_center_crop_h4_w2_batch(self, device, dtype):
inp = torch.tensor([
[[[1., 2., 3., 4.],
[5., 6., 7., 8.],
[9., 10., 11., 12.],
[13., 14., 15., 16.]]],
[[[1., 5., 9., 13.],
[2., 6., 10., 14.],
[3., 7., 11., 15.],
[4., 8., 12., 16.]]]
], device=device, dtype=dtype)
expected = torch.tensor([[[
[2., 3.],
[6., 7.],
[10., 11.],
[14., 15.],
]], [[
[5., 9.],
[6., 10.],
[7., 11.],
[8., 12.],
]]], device=device, dtype=dtype)
out_crop = kornia.center_crop(inp, (4, 2))
assert_allclose(out_crop, expected, rtol=1e-4, atol=1e-4)
def random_crop(*imgs):
w = random.choice(range(256, 512))
h = random.choice(range(256, 512))
results = []
for img in imgs:
img = kornia.resize(img, (max(h, w), max(h, w)))
img = kornia.center_crop(img, (h, w))
results.append(img)
return results
def random_crop(*imgs):
H_src, W_src = imgs[0].shape[2:]
W_tgt = random.choice(range(1024, 2048)) // 4 * 4
H_tgt = random.choice(range(1024, 2048)) // 4 * 4
scale = max(W_tgt / W_src, H_tgt / H_src)
results = []
for img in imgs:
img = kornia.resize(img, (int(H_src * scale), int(W_src * scale)))
img = kornia.center_crop(img, (H_tgt, W_tgt))
results.append(img)
return results
def test_center_crop_h2_w4(self, device, dtype):
inp = torch.tensor(
[[[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]]]],
device=device,
dtype=dtype,
)
expected = torch.tensor([[[[5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]]]], device=device, dtype=dtype)
out_crop = kornia.center_crop(inp, (2, 4))
assert_close(out_crop, expected, rtol=1e-4, atol=1e-4)
def test_jit(self, device):
@torch.jit.script
def op_script(input: torch.Tensor,
size: Tuple[int, int]) -> torch.Tensor:
return kornia.center_crop(input, size)
batch_size, channels, height, width = 1, 2, 5, 4
img = torch.ones(batch_size, channels, height, width).to(device)
crop_height, crop_width = 4, 2
actual = op_script(img, (crop_height, crop_width))
expected = kornia.center_crop(img, (crop_height, crop_width))
assert_allclose(actual, expected)
def test_center_crop_h4_w2(self, device, dtype):
inp = torch.tensor(
[[[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]]]],
device=device,
dtype=dtype,
)
height, width = 4, 2
expected = torch.tensor([[[[2.0, 3.0], [6.0, 7.0], [10.0, 11.0], [14.0, 15.0]]]], device=device, dtype=dtype)
out_crop = kornia.center_crop(inp, (height, width))
assert_close(out_crop, expected, rtol=1e-4, atol=1e-4)
def warp_transform(imgs, M_matrices):
'''
input: tuple of tensor, each element is [6,3,256,306]
for eval: imgs is one element only
'''
data = []
for i in range(6): #loop through each view
img_batch = imgs[:, i, :, :, :].cuda()
img_batch = kornia.warp_perspective(img_batch,
M_matrices[i],
dsize=(204, 306))
img_batch = kornia.center_crop(img_batch, (192, 288))
data.append(img_batch.unsqueeze(0))
return torch.cat(data, dim=0)
def test_center_crop_h2_w4(self, device, dtype):
inp = torch.tensor([[[
[1., 2., 3., 4.],
[5., 6., 7., 8.],
[9., 10., 11., 12.],
[13., 14., 15., 16.],
]]], device=device, dtype=dtype)
expected = torch.tensor([[[
[5., 6., 7., 8.],
[9., 10., 11., 12.],
]]], device=device, dtype=dtype)
out_crop = kornia.center_crop(inp, (2, 4))
assert_allclose(out_crop, expected, rtol=1e-4, atol=1e-4)
def test_center_crop_h2_w4(self, device):
inp = torch.tensor([[
[1., 2., 3., 4.],
[5., 6., 7., 8.],
[9., 10., 11., 12.],
[13., 14., 15., 16.],
]]).to(device)
height, width = 2, 4
expected = torch.tensor([[
[5., 6., 7., 8.],
[9., 10., 11., 12.],
]]).to(device)
out_crop = kornia.center_crop(inp, (height, width))
assert_allclose(out_crop, expected)
def test_center_crop_h4_w2(self):
inp = torch.tensor([[
[1., 2., 3., 4.],
[5., 6., 7., 8.],
[9., 10., 11., 12.],
[13., 14., 15., 16.],
]])
height, width = 4, 2
expected = torch.tensor([[
[2., 3.],
[6., 7.],
[10., 11.],
[14., 15.],
]])
out_crop = kornia.center_crop(inp, (height, width))
assert_allclose(out_crop, expected)
def test_center_crop_h4_w2(self, device, dtype):
inp = torch.tensor([[[
[1., 2., 3., 4.],
[5., 6., 7., 8.],
[9., 10., 11., 12.],
[13., 14., 15., 16.],
]]], device=device, dtype=dtype)
height, width = 4, 2
expected = torch.tensor([[[
[2., 3.],
[6., 7.],
[10., 11.],
[14., 15.],
]]], device=device, dtype=dtype)
out_crop = kornia.center_crop(inp, (height, width))
assert_allclose(out_crop, expected, rtol=1e-4, atol=1e-4)
def test_center_crop_h4_w2_batch(self, device, dtype):
inp = torch.tensor(
[
[[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]]],
[[[1.0, 5.0, 9.0, 13.0], [2.0, 6.0, 10.0, 14.0], [3.0, 7.0, 11.0, 15.0], [4.0, 8.0, 12.0, 16.0]]],
],
device=device,
dtype=dtype,
)
expected = torch.tensor(
[
[[[2.0, 3.0], [6.0, 7.0], [10.0, 11.0], [14.0, 15.0]]],
[[[5.0, 9.0], [6.0, 10.0], [7.0, 11.0], [8.0, 12.0]]],
],
device=device,
dtype=dtype,
)
out_crop = kornia.center_crop(inp, (4, 2))
assert_close(out_crop, expected, rtol=1e-4, atol=1e-4)
def op_script(input: torch.Tensor,
size: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
return kornia.center_crop(input, size)
def forward(self, x: torch.Tensor) -> torch.Tensor:
assert isinstance(x, torch.Tensor)
assert len(x.shape) == 4, x.shape
crop_size: int = min(x.shape[-2:]) // 2
return K.center_crop(x, (crop_size, crop_size))
