图像分块与合并

本文提供了一种将大图分块、将小图合并为大图的方法，并给出了实现代码。

目录

0、前言

1、方案

2、结果验证

---

0、前言

        在处理图像时，有时候一个图太大，难以一次性处理（比如模型输入太大会导致GPU显存溢出），这时通常采用图像分块的策略：先将图像分为若干小块patch，每块之间有一定重叠，然后分别处理每个小的patch，最后将处理结果合并为一个大图的结果。

        本文提供了一套图像分块与合并的方法。

1、方案

        本方案主要包含两部分内容：大图分块、小图合并。 

• 对于图像分块，相邻两个块之间需要有一定的重叠，这主要是为了避免在处理图像时边缘信息处理不到位的情况；
• 对于图像合并，需要考虑相邻两个小块之间重叠部分的加权融合，这样做的目的主要是考虑到每个小块处理后，其重叠部分的结果会有不一致的情况，使用加权融合可以使合并结果更为平滑。

        分块与合并的代码如下：

import cv2
import math
import numpy as np


def img2patchs(img, patch_size=(448, 448), overlap_size=(20, 20)):

    h, w, c = img.shape
    ph, pw = patch_size
    oh, ow = overlap_size

    r_h = (h - ph) % (ph - oh)
    r_w = (w - pw) % (pw - ow)


    target_w, target_h = w, h

    if not (h >= ph > oh and w >= pw > ow):
        return [[img]], (target_h, target_w), (0, 0)

    N = math.ceil((target_h-ph)/(ph-oh)) + 1
    M = math.ceil((target_w-pw)/(pw-ow)) + 1

    patchs_all = []
    for n in range(N):
        patchs_row = []
        for m in range(M):

            if n == N-1:
                ph_start = target_h - ph
            else:
                ph_start = n*(ph-oh)

            if m == M-1:
                pw_start = target_w - pw
            else:
                pw_start = m*(pw-ow)
            patch = img[ph_start:(ph_start+ph), pw_start:(pw_start+pw), :]
            patchs_row.append(patch)
        patchs_all.append(patchs_row)

    return patchs_all, (target_h, target_w), (r_h, r_w)


def patchs2img(patchs, r_size, overlap_size=(20, 20)):
    N = len(patchs)
    M = len(patchs[0])

    # print("N:{}, M:{}".format(N, M))
    oh, ow = overlap_size

    patch_shape = patchs[0][0].shape
    ph, pw = patch_shape[:2]
    r_h, r_w = r_size

    mode = 'GRAY' if len(patch_shape) == 2 else 'RGB'
    c = 3

    if N == 1 and M == 1:
        return_img = patchs[0][0]
        return return_img if mode == 'RGB' else cv2.cvtColor(return_img, cv2.COLOR_GRAY2RGB)

    row_imgs = []
    for n in range(N):
        row_img = patchs[n][0] if mode == 'RGB' else cv2.cvtColor(patchs[n][0], cv2.COLOR_GRAY2RGB)
        for m in range(1, M):
            if m == M - 1 and r_w != 0:
                ow_new = pw - r_w
            else:
                ow_new = ow
            # ow_new = ow

            patch = patchs[n][m] if mode == 'RGB' else cv2.cvtColor(patchs[n][m], cv2.COLOR_GRAY2RGB)
            # print(mode, patch.shape)
            h, w = row_img.shape[:2]
            new_w = w + pw - ow_new
            big_row_img = np.zeros((h, new_w, c), dtype=np.uint8)
            big_row_img[:, :w-ow_new, :] = row_img[:, :w-ow_new, :]
            big_row_img[:, w:, :] = patch[:, ow_new:, :]
            overlap_row_01 = row_img[:, w-ow_new:, :]
            overlap_row_02 = patch[:, :ow_new, :]

            # get weight
            weight = vertical_grad(overlap_row_01.shape, 0, 255, mode='w') / 255
            overlap_row = (overlap_row_01 * (1 - weight)).astype(np.uint8) + (overlap_row_02 * weight).astype(np.uint8)
            big_row_img[:, w-ow_new:w, :] = overlap_row

            row_img = big_row_img

        row_imgs.append(row_img)

    column_img = row_imgs[0]
    for i in range(1, N):
        if i == N - 1 and r_h != 0:
            oh_new = ph - r_h
        else:
            oh_new = oh
        # oh_new = oh
        row_img = row_imgs[i]
        h, w = column_img.shape[:2]
        new_h = h + ph - oh_new
        big_column_img = np.zeros((new_h, w, c), dtype=np.uint8)
        big_column_img[:h-oh_new, :, :] = column_img[:h-oh_new, :, :]
        big_column_img[h:, :, :] = row_img[oh_new:, :, :]
        overlap_column_01 = column_img[h-oh_new:, :, :]
        overlap_column_02 = row_img[:oh_new, :, :]

        # get weight
        weight = vertical_grad(overlap_column_01.shape, 0, 255, mode='h') / 255
        overlap_column = (overlap_column_01 * (1 - weight)).astype(np.uint8) + (overlap_column_02 * weight).astype(np.uint8)
        big_column_img[h-oh_new:h, :, :] = overlap_column

        column_img = big_column_img

    return column_img


def vertical_grad(shape, color_start, color_end, mode):
    h, w = shape[0], shape[1]
    L = h if mode == 'h' else w

    grad_img = np.ndarray(shape, dtype=np.uint8)

    # grad
    grad = float(color_end - color_start) / L

    for i in range(L):
        if mode == 'h':
            grad_img[i, :] = color_start + i * grad
        else:
            grad_img[:, i] = color_start + i * grad

    return grad_img

2、结果验证

        为了验证，我们先将上述代码保存为一个文件：img2patchs.py；然后，写一个脚本调用上述方法，来对一张图像进行分块与合并：

import cv2
import os
from img2patchs import img2patchs, patchs2img


if __name__ == '__main__':
    img_path = "图片/vlcsnap-2022-09-14-17h24m48s612.png"
    visual_path = "visuals_merge"
    os.makedirs(visual_path, exist_ok=True)
    os.makedirs(os.path.join(visual_path, 'patchs'), exist_ok=True)
    overlap_size = (10, 10)

    img = cv2.imread(img_path)
    cv2.imwrite(os.path.join(visual_path, 'raw.jpg'), img)
    patchs, _, r_size = img2patchs(img, patch_size=(224, 224), overlap_size=overlap_size)

    # save patchs
    N = len(patchs)
    M = len(patchs[0])
    print("N*M: {}*{}".format(N, M))
    for n in range(N):
        for m in range(M):
            patch = patchs[n][m]
            patch_filename = "patch_{}_{}.jpg".format(str(n).zfill(2), str(m).zfill(2))
            cv2.imwrite(os.path.join(visual_path, 'patchs', patch_filename), patch)

    # patchs2img
    recovery_img = patchs2img(patchs, r_size, overlap_size)
    cv2.imwrite(os.path.join(visual_path, "recovery_img.jpg"), recovery_img)

        运行上述代码后，会生成一个文件夹，如下：

        其中，patchs是分出来的小块图像，里面形如：

        raw是原图， recovery_img是利用小块合并恢复出来的大图。

        在本次示例中，由于分块后，没有进一步的处理，所以体现不出来我们的合并方法的优势。如果有进一步处理的话，在合并时，若没有我们所用的加权融合，就会出现很明显的拼接痕迹。大家有兴趣的话也可以试试。