kornia 之一 基础和 filter_from kornia import resize

起因

夏天时课题组一个小哥推荐了kornia这个库，传说是把一些典型的cv操作基于PyTorch设计成了differentiable的实现，辅助进行深度学习模型的训练。但是一直也没有仔细看过或者用过这个库，最近一个模型中需要用到一些常见的filter操作，于是想到了这个库，尝试了一下，把所见所感写在这里作为记录。

经过

当前kornia的版本是0.4.0. kornia 也在不断更新中。其实PyTorch更新也挺勤的。。。

常规import

若没有另外描述，那么如下import是默认存在的

import cv2
import matplotlib.pyplot as plt
import numpy as np

import torch
import kornia

def show_image(img, name='NoName', flagBGR=True):
    if ( flagBGR ):
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    fig = plt.figure(num=name)
    ax = fig.add_subplot(111)
    ax.imshow(img)
    ax.axis('off')
    ax.set_title(name)

def show_images(imgs, name='NoName', cols=2):
    '''
    imgs (list of dicts): {"img": img, "title": title, "bgr": True}
    '''
    
    N = len( imgs )
    assert ( N > 0 )
    
    if ( 1 == N ):
        d = imgs[0]
        show_image( d['img'], d['title'], d['bgr'] )
        return
    
    rows = ( N - 1 ) // cols + 1
    
    fig = plt.figure(num=name)
    fig.tight_layout()
    
    for i in range(N):
        img = imgs[i]['img']
        if ( imgs[i]['bgr'] ):
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        ax = fig.add_subplot( rows, cols, i+1 )
        ax.axis('off')
        ax.imshow(img)
        ax.set_title(imgs[i]['title'])
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

基本操作

简单看了一下tutorial，kornia提供了一些基础操作用于python环境下图像和torch Tesnor之间的沟通/转换。主要是image_to_tensor()和tensor_to_image()函数。这两个函数都定义在kornia/utils/image.py.

image_to_tensor() 贴心地提供了检查维度是否为single channel并且是否需要增加batch维度的服务，但是其默认行为是输出一个未经normalize过的tensor，tensor的dtype将为torch.uint8，这与PyTorch的torchvision.transforms.ToTensor的行为不太一样。

# Test image_to_tensor().
t = kornia.image_to_tensor(img, keepdim=False)
print('t.shape = {}'.format(t.shape))
# out: t.shape = torch.Size([1, 3, 1154, 1732])
print(t.dtype)
# out: torch.uint8
1
2
3
4
5
6

所以我写了一个nice litte wrapper.

def image_2_float_tensor(img, keepdim=True, dtype=torch.float32):
    return kornia.image_to_tensor(img, keepdim=keepdim).to(dtype) / 255
1
2

ft = image_2_float_tensor(img, keepdim=False, dtype=torch.float32)
print('ft.shape = {}'.format(ft.shape))
print('ft.dtype = {}'.format(ft.dtype))
print('ft.min() = {}, ft.max() = {}'.format( ft.min(), ft.max() ))
# out: ft.shape = torch.Size([1, 3, 1154, 1732])
#      ft.dtype = torch.float32
#      ft.min() = 0.0, ft.max() = 1.0
1
2
3
4
5
6
7

类似的，tensor_to_image() 贴心地替我们做好了维度转换。

• 如果tensor的维度是2，则维持维度不变。
• 如果tesnor的维度是3，则permute((1, 2, 0))。进一步若channel是1，那么squeeze该channel。
• 如果tensor的维度是4，则permute((0, 2, 3, 1))。并且根据batch和channel的情况，将他们squeeze掉。仍然是贴心地替我们消除basth size 1 和single channel的情况。
• 如果tensor是在GPU上，则传送到CPU。
• 返回的是NumPy array。

虽然上述功能自己都有实现，但是偷懒用一用也是可以的。然后这是我的wrapper

def float_tensor_2_image(t):
    img = kornia.tensor_to_image(t)
    img = np.clip( img, 0.0, 1.0 ) * 255
    return img.astype(np.uint8)
1
2
3
4

我希望np.uint8进，np.uint8出。在我的实际应用中，多数情况从tensor还原image是为了显示，所以丢失一些数值精度是没问题的。

感觉kornia的作者很对我的胃口，kornia/utils/中另一个文件grid.py提供了两个用于生成sample坐标的函数，很方便，又不用自己写了，虽然都已经有了，偷懒用一下也挺好。

对我而言，用于2D数据的create_meshgrid()比较有用。但是我一般都是这么用

def grid_like(t):
    assert( 4 == t.ndim )
    H, W = t.shape[2:4]
    return kornia.create_meshgrid( H, W, 
    	normalized_coordinates=True, device=t.device )
1
2
3
4
5

kornia返回给我门[-1, 1] normalized过的BxHxWx2 grid，就是torch.grid_sample()接受的那种形式。

grid = grid_like(ft.cuda())
print('grid.shape = {}'.format( grid.shape ))
print('grid.is_cuda = {}'.format( grid.is_cuda ))
# out:
# grid.shape = torch.Size([1, 1154, 1732, 2])
# grid.is_cuda = True
1
2
3
4
5
6

最后，kornia/utils/pointcloud_io.py内有读写PLY点云的实现，功能比较单一，建议自己另行实现。另外默认写的是double类型的浮点数。

Filter

最开始关心的是median filter的实现，kornia.filters.MedianBlur。kornia/filters/median.py即为MedianBlur的实现，挺有意思，之前还真没思考过在PyTorch里面如何“选择”特定的feature。MedianBlur的实现最终使用了torch.median()函数，所以它要做的就是根据指定的卷积核尺寸选出feature。选feature，其实是手工指定了一个kernel，在下面的示例代码里展示了这个kernel的值。kornia.filters.MedianBlur是分channel进行median filter的，也就是每个channel单独filter，并且自动进行zero-padding。测试代码

def isotropic_resize(img, f, interpolation=cv2.INTER_CUBIC):
    assert( f > 0 )
    return cv2.resize( img, ( 0, 0 ), fx=f, fy=f, interpolation=interpolation )

plt.close('all')

# Load the image.
img = cv2.imread('../Basics/IMG_2417_Resized.JPG', cv2.IMREAD_UNCHANGED)

# Resize the image so we could see the difference before and after the filter.
img = isotropic_resize(img, 0.25)
print('img.shape = {}'.format(img.shape))
H, W = img.shape[:2]
1
2
3
4
5
6
7
8
9
10
11
12
13

img.shape = (288, 433, 3)

t = image_2_float_tensor(img, keepdim=False).cuda()
t.requires_grad = True

print('t.shape = {}'.format(t.shape))
1
2
3
4

t.shape = torch.Size([1, 3, 288, 433])

# Median blur module.
median = kornia.filters.MedianBlur([3, 3])

print('median.kernel.shape = {}'.format(median.kernel.shape))
print('median.kernel = {}'.format(median.kernel))
1
2
3
4
5

median.kernel.shape = torch.Size([9, 1, 3, 3])
median.kernel = tensor([[[[1., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]],
[[[0., 1., 0.],
[0., 0., 0.],
[0., 0., 0.]]],
[[[0., 0., 1.],
[0., 0., 0.],
[0., 0., 0.]]],
[[[0., 0., 0.],
[1., 0., 0.],
[0., 0., 0.]]],
[[[0., 0., 0.],
[0., 1., 0.],
[0., 0., 0.]]],
[[[0., 0., 0.],
[0., 0., 1.],
[0., 0., 0.]]],
[[[0., 0., 0.],
[0., 0., 0.],
[1., 0., 0.]]],
[[[0., 0., 0.],
[0., 0., 0.],
[0., 1., 0.]]],
[[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 1.]]]])

tm = median(t)

tImg = float_tensor_2_image(tm)

imgs = [ {'img': img, 'title': 'Original @ %dx%d' % (W, H), 'bgr': True}, 
        {'img': tImg, 'title': 'Blurred @ %dx%d' % (W, H), 'bgr': True} ]
show_images(imgs, 'Original & blurred')

y = tm * 2.0 + 1.0

y.backward(torch.ones_like(y))

print(t.grad)
1
2
3
4
5
6
7
8
9
10
11
12
13

tensor([[[[0., 0., 0., …, 2., 4., 0.],
[0., 2., 2., …, 0., 0., 4.],
[2., 4., 0., …, 8., 0., 0.],
…,
[4., 0., 2., …, 0., 8., 0.],
[2., 6., 6., …, 0., 0., 2.],
[0., 0., 0., …, 2., 0., 2.]],
[[0., 0., 0., …, 2., 2., 0.],
[0., 2., 2., …, 2., 2., 2.],
[2., 4., 0., …, 2., 4., 0.],
…,
[2., 0., 2., …, 0., 2., 2.],
[2., 6., 2., …, 0., 6., 2.],
[2., 0., 0., …, 4., 0., 0.]],
[[0., 0., 0., …, 2., 0., 2.],
[0., 2., 2., …, 2., 2., 2.],
[2., 4., 0., …, 2., 4., 0.],
…,
[2., 0., 2., …, 8., 0., 0.],
[2., 6., 4., …, 0., 0., 4.],
[2., 0., 0., …, 4., 0., 2.]]]], device=‘cuda:0’)

t.grad.max()
1

tensor(18., device=‘cuda:0’)

t.grad.min()
1

tensor(0., device=‘cuda:0’)

至少从gradients上看是make sense的。

这是测试用图像

比较神奇的是，kornia提供一个motion blur的实现。

结果

作者懂我，用着方便。