语义分割数据集增强增广

"""
@author: jialin_li
@data: 20230829
@version: v1.0
"""

import os
import cv2
import numpy as np
from skimage.util import random_noise
import shutil

# 原始数据集路径 (修改为自己的路径)
dataset_path = 'data'
images_dir = os.path.join(dataset_path, 'images')  #原图文件夹
annotations_dir = os.path.join(dataset_path, 'annotations') #单通道彩图文件夹

# 增强后的数据集路径
enhanced1_dir = os.path.join(dataset_path, 'enhanced1')  #增强后图片保存路径
enhanced2_dir = os.path.join(dataset_path, 'enhanced2')  #标签保存路径
os.makedirs(enhanced1_dir, exist_ok=True)
os.makedirs(enhanced2_dir, exist_ok=True)

# 获取原始图像和标注文件路径列表
image_files = sorted(os.listdir(images_dir))
annotation_files = sorted(os.listdir(annotations_dir))

# 遍历每个图像文件并进行增强
for img_file, anno_file in zip(image_files, annotation_files):
    # 读取原始图像和标注文件
    img_path = os.path.join(images_dir, img_file)
    img = cv2.imread(img_path)

    anno_path = os.path.join(annotations_dir, anno_file)
    anno = cv2.imread(anno_path)

    #提供了两种增强方法 可以在此基础上自行添加
    # 随机对比度变化
    alpha = np.random.uniform(0.5, 1.5)
    beta = np.random.uniform(-50, 50)
    contrast_img = cv2.convertScaleAbs(img, alpha=alpha, beta=beta)

    # 随机添加噪声
    noise = random_noise(img, mode='gaussian', var=0.02)
    noise_img = np.array(255 * noise, dtype=np.uint8)

    # 保存增强后的原图像
    enhanced_img_path1 = os.path.join(enhanced1_dir, f"contrast_{img_file}")
    cv2.imwrite(enhanced_img_path1, contrast_img)

    # 复制标注文件到增强后的目录并改名（对比度增强后的标签图像与原始标签图像不一样）
    enhanced_anno_dst1 = os.path.join(enhanced2_dir, f"contrast_{anno_file}")
    shutil.copy(anno_path, enhanced_anno_dst1)

    # 保存增强后的原图像
    enhanced_img_path2 = os.path.join(enhanced1_dir, f"noise_{img_file}")
    cv2.imwrite(enhanced_img_path2, noise_img)

    # 保存增强后的标注文件
    enhanced_anno_path2 = os.path.join(enhanced2_dir, f"noise_{anno_file}")
    shutil.copy(anno_path, enhanced_anno_path2)

    # 显示处理进度
    print(f'图像 {img_file} 处理完成')

print('数据集增强完毕。')