Python中将图像切分为小的patch

王云峰

发布于 2023-10-23 10:31:03

2990

发布于 2023-10-23 10:31:03

文章被收录于专栏：Yunfeng's Simple BlogYunfeng's Simple Blog

问题定义

假如有张1000x1000的图像，我们要将它切成20x20的小patch，该怎么处理呢？最简单的方法就是采用两重for循环，每次计算小patch对应的下标，在原图上进行crop:

import numpy as np

size = 1000
ncols = 20
nrows = 20
img = np.random.rand(size, size)

patches = []

for i in range(size//ncols):
	for j in range(size//nrows):
		patch = img[ncols*i:ncols*(i+1), nrows*j:nrows*(j+1)]
		patches.append(patch)

patches = np.array(patches)

但这样总共需要循环50*x50=2500次，而我们知道 Python 的 for 循环比较慢，因此整体开销还是比较大的，有没有更快的方式呢？

reshape + swapaxes

搜索发现可以使用 reshape + swapaxes函数的组合来完成这个功能:

import numpy as np

size = 1000
ncols = 20
nrows = 20
img = np.random.rand(size, size)

patches = img.reshape(size // ncols, ncols, -1, nrows).swapaxes(1, 2).reshape(-1, ncols, nrows)

完整对比代码如下：

import time

import numpy as np


size = 1000
ncols = 20
nrows = 20

# 统计100次耗时
for i in range(100):
    img = np.random.rand(size, size)
    t0 = time.time()
    patches0 = img.reshape(size // ncols, ncols, -1, nrows).swapaxes(1, 2).reshape(-1, ncols, nrows)

    t1 = time.time()
    d1 = t1 - t0

    patches = []
    for i in range(size//ncols):
        for j in range(size//nrows):
            patch = img[ncols*i:ncols*(i+1), nrows*j:nrows*(j+1)]
            patches.append(patch)

    patches1 = np.array(patches)
    t2 = time.time()
    d2 = t2 - t1

    print('time ratio:', d2/d1)
    print('diff:', (patches0-patches1).sum())

实际测试对于1000x1000的图像，采用reshape + swapaxes 要比循环快大约4倍。

time ratio: 4.684571428571428
diff: 0.0
time ratio: 4.806614785992218
diff: 0.0
time ratio: 4.696482035928144
diff: 0.0
time ratio: 3.00382226469183
diff: 0.0
time ratio: 3.710854363028276
diff: 0.0
...

Pytorch中的实现？

Pytorch相比numpy，又增加了许多操作tensor的函数，因此实现方式会更多，这里大概列一下几种实现，具体函数可以查询 Pytorch 的文档:

patches1 = img.unfold(0, ncols, nrows).unfold(1, ncols, nrows).reshape(-1, ncols, nrows)
patches2 = img.reshape(size//ncols, ncols, -1, nrows).swapaxes(1, 2).reshape(-1, ncols, nrows)
patches3 = img.reshape(size//ncols, ncols, -1, nrows).permute(0, 2, 1, 3).reshape(-1, ncols, nrows)