批处理对图像添加椒盐噪声和高斯噪声

 多任务学习目标识别和语义分割数据增广，批处理对图像添加校验椒盐和高斯噪声，同时xml，label不做更改，重新编号写入文件夹

"添加椒盐噪声"
def sp_noise(image,prob):      # prob:噪声比例

    output = np.zeros(image.shape,np.uint8)
    thres = 1 - prob
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            rdn = random.random()
            if rdn < prob:
                output[i][j] = 0
            elif rdn > thres:
                output[i][j] = 255
            else:
                output[i][j] = image[i][j]
    return output

"添加高斯噪声"
def gasuss_noise(image, mean=0, var=0.001):     #  mean : 均值  var : 方差

    image = np.array(image/255, dtype=float)
    noise = np.random.normal(mean, var ** 0.5, image.shape)
    out = image + noise
    if out.min() < 0:
        low_clip = -1.
    else:
        low_clip = 0.
    out = np.clip(out, low_clip, 1.0)
    out = np.uint8(out*255)
    return out

def add_noise(image):
    l = np.random.uniform(0, 1)
    if l > 0.6:
        img_noise = sp_noise(image,0.02)
    else:
        img_noise = gasuss_noise(image, mean=0, var=0.001)

    return img_noise

头文件：

import numpy as np
import cv2
import math
import random
import os
import xml.etree.ElementTree as ET
from PIL import Image

主函数：

if __name__ == '__main__':
    resize_jpg_path = r'F:\数据集\image'    # 原始的图像路径
    resize_png_path = r'F:\数据集\label'
    resize_xml_path = r'F:\数据集\annotation'

    orginal_jpg_files =os.listdir(resize_jpg_path)
    orginal_png_files = os.listdir(resize_png_path)
    orginal_xml_files = os.listdir(resize_xml_path)

    # new_size = (224, 224)
    # original_size = (256, 256)

    new_jpg_path='F:/数据集/image'
    new_png_path = 'F:/数据集/label'
    new_xml_path = 'F:/数据集/annotation'

    num =1
    
    "添加椒盐噪声，高斯噪声"
    for i in orginal_jpg_files:
        orginal_image_id = i
        orginal_png_id = i.split('.')[0] + '.png'
        orginal_xml_id = i.split('.')[0] + '.xml'

        new_jpg_id = str(num) + '.jpg'
        new_jpg_name = os.path.join(new_jpg_path, new_jpg_id)
        new_png_id = str(num) + '.png'
        new_png_name = os.path.join(new_png_path, new_png_id)
        new_xml_id = str(num) + '.xml'
        new_xml_name = os.path.join(new_xml_path, new_xml_id)

        orginal_jpg_name = os.path.join(resize_jpg_path, orginal_image_id)
        image = Image.open(orginal_jpg_name)
        img = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
        new_jpg_noise = add_noise(img)

        png = Image.open(os.path.join(resize_png_path, orginal_png_id))
        tree = ET.parse(os.path.join(resize_xml_path, orginal_xml_id))
        # print(new_jpg_noise)

        cv2.imencode('.jpg', new_jpg_noise)[1].tofile(new_jpg_name)
        png.save(new_png_name)
        tree.write(new_xml_name)
        num = num + 1