在进行深度学习网络学习时遇到问题，报错说batch size中的图像size不一致，虽然模型中自带的数据处理模块中有统一尺寸这个功能，还是单独写个python代码处理一下。一共包含4个功能，同意尺寸，随机划分训练集/验证集/测试集，计算图像的均值/方差，图像标签二值化处理。

我的数据集是开源的毕节市滑坡公开数据集，长这样

1.统一图像尺寸

import cv2
import os

#提取目录下所有图片,更改尺寸后保存到另一目录
from PIL import Image
import os.path
import glob
def convertjpg(jpgfile,outdir,width=224,height=224):
    img=Image.open(jpgfile)
    try:
        new_img=img.resize((width,height),Image.BILINEAR)
        new_img.save(os.path.join(outdir,os.path.basename(jpgfile)))
    except Exception as e:
        print(e)
for jpgfile in glob.glob((r'F:\data\bijie\val\image\*.png')):   # 处理签图像的路径
    convertjpg(jpgfile, r"F:\data\bijie1\val\image")            # 处理后图像的路径

# 查看更改前后图像的size
# infile = r'F:\data\bijie\train\image\df027.png'
# infile_1 =r'F:\data\bijie1\train\image\df027.png'
# im = Image.open(infile)
# im_1 = Image.open(infile_1)
# (x, y) = im.size  # read image size
# (x1,y1) = im_1.size
# print ('original size: ', x, y)
# print ('current size: ', x1, y1)

2.随机划分训练集/验证集/测试集

我的数据，图像都存放在一个文件夹中，label也存放在一个文件夹中。

import os
import shutil
import random
from sklearn.model_selection import train_test_split

# 文件夹路径
image_folder = r'F:\Bijie_landslide_dataset\chuli\image'
mask_folder = r'F:\Bijie_landslide_dataset\chuli\mask'
train_folder = r'F:\mmsegmentation-main\data\Bijie4\train'
val_folder = r'F:\mmsegmentation-main\data\Bijie4\val'
test_folder = r'F:\mmsegmentation-main\data\Bijie_landslide_dataset\test'

# 创建目标文件夹，, test_folder
for folder in [train_folder, val_folder]:
    os.makedirs(os.path.join(folder, 'image'), exist_ok=True)
    os.makedirs(os.path.join(folder, 'mask'), exist_ok=True)

# 获取所有图片和mask文件
image_files = sorted([f for f in os.listdir(image_folder) if f.endswith('.png')])
mask_files = sorted([f for f in os.listdir(mask_folder) if f.endswith('.png')])

# 确保图片和mask数量一致
assert len(image_files) == len(mask_files), "Images and masks count mismatch"


# 划分数据集，先8/2划分训练集/temp，再将temp以1/1划分，得到训练集/验证集/测试集=8/1/1
train_imgs, val_imgs = train_test_split(image_files, test_size=0.2, random_state=42)
val_imgs, test_imgs = train_test_split(temp_imgs, test_size=0.5, random_state=42)

# 移动文件到相应的文件夹
def move_files(files, src_folder, dest_folder):
    for file in files:
        shutil.copy(os.path.join(src_folder, file), os.path.join(dest_folder, file))

# 训练集
move_files(train_imgs, image_folder, os.path.join(train_folder, 'image'))
move_files(train_imgs, mask_folder, os.path.join(train_folder, 'mask'))

# 验证集
move_files(val_imgs, image_folder, os.path.join(val_folder, 'image'))
move_files(val_imgs, mask_folder, os.path.join(val_folder, 'mask'))

# 测试集
move_files(test_imgs, image_folder, os.path.join(test_folder, 'image'))
move_files(test_imgs, mask_folder, os.path.join(test_folder, 'mask'))

3.计算输入图像均值/方差

深度学习网络模型中有时候需要根据图像的均值和方差进一步处理，可以通过一下代码批量计算。

import os
from PIL import Image
import numpy as np
import tqdm

def main(path):
    # 数据集通道数
    img_channels = 3
    img_names = os.listdir(path)
    cumulative_mean = np.zeros(img_channels)
    cumulative_std = np.zeros(img_channels)

    for img_name in tqdm.tqdm(img_names, total=len(img_names)):
        img_path = os.path.join(path, img_name)
        img = np.array(Image.open(img_path)) / 255.
        # 对每个维度进行统计，Image.open打开的是HWC格式，最后一维是通道数
        for d in range(3):
            cumulative_mean[d] += img[:, :, d].mean()
            cumulative_std[d] += img[:, :, d].std()

    mean = cumulative_mean / len(img_names)
    std = cumulative_std / len(img_names)
    print(f"mean: {mean}")
    print(f"std: {std}")

if __name__ == '__main__':
    main(r"F:\data\bijie1\train\image")   #数据集路径

输出如图：

4.图像二值化

我的分类任务是二分类，label是0和255，分别对应正样本与负样本，但我的网络模型只能识别出0和1这两个值，所以这段代码的目的是将0和255两个值的图像转换为0和1两个值。

import os
import numpy as np
from PIL import Image

# 源文件夹路径和目标文件夹路径
source_folder = r'F:\Bijie_landslide_dataset\chuli\mask'   # 源文件夹路径
destination_folder = r'F:\Bijie_landslide_dataset\mask'   # 目标文件夹的路径

# 确保目标文件夹存在
os.makedirs(destination_folder, exist_ok=True)

# 遍历源文件夹中的所有图像文件
for filename in os.listdir(source_folder):
    if filename.endswith('.png') or filename.endswith('.jpg'):  # 根据你的图像格式调整
        # 构建完整的文件路径
        file_path = os.path.join(source_folder, filename)

        # 读取图像
        img = Image.open(file_path).convert('L')
        img_array = np.array(img)

        # 转换为二值图像
        binary_array = np.where(img_array == 255, 1, 0)

        # 保存二值图像到目标文件夹
        binary_img = Image.fromarray(binary_array.astype(np.uint8))
        binary_img.save(os.path.join(destination_folder, filename))