图像语义分割——python滑窗法裁剪数据

B站：xxx
CSDN：python图像分割——滑窗法裁剪数据_百年后封笔-CSDN博客
Github：封笔
公众号：百年后封笔

一、 背景和需求

对图像分割而言，往往给的原图是非常大的，无法直接用于网络训练，因此有必要使用滑窗法进行图像的裁剪，把大图裁剪成一个个小的patch图，如下所示，当然如果有分类的需求，也可以根据要求来把裁剪的图像进行分类。

下面给出一个例子：

目标：把一个二分类（不包含背景）的细胞分割图（label标注为红色和绿色），使用滑窗法进行分割，并根据patch中两种类别的占比来进行分类保存。
注：不考虑超出像素边界的残缺patch，如需考虑可以用cv2的 copyMakeBorder解决

二、实现

2.1 代码实现

直接上代码：

import os
import cv2
import numpy as np
from tqdm import tqdm

# 根据传统视觉进行图像两类标签的mask生成（原始标签是彩色图像，需要提取绿色和红色的部分）
def get_g_r_label(label):
    b, g, r = label[..., 0], label[..., 1], label[..., 2]
    b = b.astype(np.float)
    g = g.astype(np.float)
    r = r.astype(np.float)
    green = g - b - r
    red = r - b - g

    red = np.where(red > 0, 255, 0)
    green = np.where(green > 0, 255, 0)
    #
    # cv2.imshow('label', label.astype(np.uint8))
    # cv2.imshow('green', green.astype(np.uint8))
    # cv2.imshow('red', red.astype(np.uint8))
    # cv2.waitKey(0)

    return red.astype(np.uint8), green.astype(np.uint8)

# 裁剪函数
def crop(img, label, label_g, label_r, save_dirs, save_name,
         crop_size=(50, 50), gap=(50, 50), ratio=0.7, isshow=False):
    h, w, _ = img.shape
    gp_w, gp_h = gap
    cp_w, cp_h = crop_size
    num = 0
    for j in range(0, h, gp_h):
        if j + cp_h > h: continue
        for i in range(0, w, gp_w):
            if i + cp_w > w: continue
            # print(j, i, j*gap_h, j*gap_h+cp_h, i*gap_w, i*gp_w+cp_w)
            cp_img = img[j:j+cp_h, i:i+cp_w, :]
            a_img = label_r[j:j+cp_h, i:i+cp_w]
            b_img = label_g[j:j+cp_h, i:i+cp_w]
            if np.sum(a_img.flatten()) > cp_w * cp_h * 255 * ratio:
                cv2.imwrite(os.path.join(save_dirs[0], save_name.replace('.jpg', f'_{num}.jpg')), cp_img)
                if isshow:
                    cv2.imwrite(os.path.join(save_dirs[0], save_name.replace('.jpg', f'_{num}_show.jpg')), label[j:j+cp_h, i:i+cp_w, :])

            elif np.sum(b_img.flatten()) > cp_w * cp_h * 255 * ratio:
                cv2.imwrite(os.path.join(save_dirs[1], save_name.replace('.jpg', f'_{num}.jpg')), cp_img)
                if isshow:
                    cv2.imwrite(os.path.join(save_dirs[1], save_name.replace('.jpg', f'_{num}_show.jpg')), label[j:j+cp_h, i:i+cp_w, :])

            num += 1

            # cv2.imshow('cp', cp_img)
            # cv2.imshow('ori', img)
            # cv2.imshow('a', a_img)
            # cv2.imshow('b', b_img)
            # cv2.waitKey(0)


if __name__ == '__main__':
    label_dir = r'path/to/your_label'
    img_dir = r'path/to/your_images'
    # 定义两个类别的保存路径
    save_dir1 = r'./cls_1'
    save_dir2 = r'./cls_2'
    if not os.path.isdir(save_dir1): os.makedirs(save_dir1)
    if not os.path.isdir(save_dir2): os.makedirs(save_dir2)
    crop_w, crop_h = 100, 100 # 定义裁剪图像尺寸
    gap_w, gap_h = 100, 100 # 定义滑动间隔
    ratio = 0.7 # 像素占比
    for label_name in tqdm(os.listdir(label_dir)):
        img_path = os.path.join(img_dir, label_name.replace('.v2', ''))
        label_path = os.path.join(label_dir, label_name)
        label = cv2.imread(label_path, cv2.IMREAD_COLOR)
        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
        red, green = get_g_r_label(label)  # 获取标签模板
        crop(img, label, red, green, [save_dir1, save_dir2], save_name=label_name.replace('.v2', ''),
             crop_size=(crop_w, crop_h), gap=(gap_w, gap_h), ratio=ratio, isshow=False)

2.2 根据颜色获取不同类别的mask

# 根据传统视觉进行图像两类标签的mask生成（原始标签是彩色图像，需要提取绿色和红色的部分）
def get_g_r_label(label):
    b, g, r = label[..., 0], label[..., 1], label[..., 2]
    b = b.astype(np.float)
    g = g.astype(np.float)
    r = r.astype(np.float)
    green = g - b - r
    red = r - b - g

    red = np.where(red > 0, 255, 0)
    green = np.where(green > 0, 255, 0)
    #
    # cv2.imshow('label', label.astype(np.uint8))
    # cv2.imshow('green', green.astype(np.uint8))
    # cv2.imshow('red', red.astype(np.uint8))
    # cv2.waitKey(0)

    return red.astype(np.uint8), green.astype(np.uint8)

2.3 滑窗法裁剪 crop

# 裁剪函数
def crop(img, label, label_g, label_r, save_dirs, save_name,
         crop_size=(50, 50), gap=(50, 50), ratio=0.7, isshow=False):
    h, w, _ = img.shape
    gp_w, gp_h = gap
    cp_w, cp_h = crop_size
    num = 0
    for j in range(0, h, gp_h):
        if j + cp_h > h: continue
        for i in range(0, w, gp_w):
            if i + cp_w > w: continue
            # print(j, i, j*gap_h, j*gap_h+cp_h, i*gap_w, i*gp_w+cp_w)
            cp_img = img[j:j+cp_h, i:i+cp_w, :]
            a_img = label_r[j:j+cp_h, i:i+cp_w]
            b_img = label_g[j:j+cp_h, i:i+cp_w]
            if np.sum(a_img.flatten()) > cp_w * cp_h * 255 * ratio:
                cv2.imwrite(os.path.join(save_dirs[0], save_name.replace('.jpg', f'_{num}.jpg')), cp_img)
                if isshow:
                    cv2.imwrite(os.path.join(save_dirs[0], save_name.replace('.jpg', f'_{num}_show.jpg')), label[j:j+cp_h, i:i+cp_w, :])

            elif np.sum(b_img.flatten()) > cp_w * cp_h * 255 * ratio:
                cv2.imwrite(os.path.join(save_dirs[1], save_name.replace('.jpg', f'_{num}.jpg')), cp_img)
                if isshow:
                    cv2.imwrite(os.path.join(save_dirs[1], save_name.replace('.jpg', f'_{num}_show.jpg')), label[j:j+cp_h, i:i+cp_w, :])

            num += 1

如上是分割出来的不同类别的patch图像。