Segmentación semántica de imágenes: método de ventana deslizante de Python para recortar datos

Estación B: xxx
CSDN: segmentación de imágenes de Python: método de ventana deslizante para recortar datos_Cubierto después de cien años - Blog de CSDN
Github: limitado
después de cien años Cuenta pública: limitado después de cien años

1. Antecedentes y necesidades

Para la segmentación de imágenes, la imagen original suele ser muy grande y no se puede usar directamente para el entrenamiento de la red, por lo que es necesario utilizar el método de ventana deslizante para recortar la imagen y cortar la imagen grande en imágenes de parches pequeños, como se muestra a continuación. Por supuesto, si es necesario realizar una clasificación, las imágenes recortadas también se pueden clasificar según los requisitos.

A continuación se da un ejemplo:

Objetivo: utilizar el método de ventana deslizante para segmentar un mapa de segmentación de celdas de dos categorías (excluido el fondo) (etiquetado como rojo y verde), clasificarlo y guardarlo de acuerdo con la proporción de las dos categorías en el parche.
Nota: Los parches incompletos más allá del límite de píxeles no se consideran. Si es necesario, puede usar copyMakeBorder de cv2 para resolver el problema.

2. Realización

2.1 Implementación del código

Vaya directamente al código:

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 Obtener máscaras de diferentes categorías según el color

# 根据传统视觉进行图像两类标签的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 Recorte con el método de ventana deslizante

# 裁剪函数
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

Las anteriores son imágenes de parches segmentadas de diferentes categorías.