多任务学习目标检测和语义分割，数据处理汇总

原始图像，目标检测的标签框和语义分割的标签水平翻转处理：https://blog.csdn.net/hejunran/article/details/117483650

原始图像加雾化，目标检测的标签框和语义分割的标签重命名处理：https://blog.csdn.net/hejunran/article/details/117513626

原始图像加椒盐噪声和高斯噪声，目标检测的标签框和语义分割的标签重命名处理：

https://blog.csdn.net/hejunran/article/details/117534930

原始图像resize，边缘填充成原始图像大小，目标检测的标签框xml也resize，和语义分割的标签resize重命名处理：https://blog.csdn.net/hejunran/article/details/117566416

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

name_classes = ['hqc']  # 类别名，可以更改为对应的voc类别名称即可


"图像水平翻转处理"
def flip_horizontal(jpg_file,new_jpg,png_file,new_png):
    img = Image.open(jpg_file)
    png =Image.open(png_file)

    out_img = img.transpose(Image.FLIP_LEFT_RIGHT)
    out_png = png.transpose(Image.FLIP_LEFT_RIGHT)

    out_img.save(new_jpg)
    out_png.save(new_png)

def flip_xml(xml_file,new_xml):
    tree = ET.parse(xml_file)
    objs = tree.findall('object')
    for ix, obj in enumerate(objs):
        name = obj.find('name').text
        if name == 'hqc':
            print(xml_file)
            obj_new = obj.find('bndbox')
            xmin = str( 256 - int(obj_new.find('xmax').text))     # 因为进行水平翻转，这只改变x坐标值，y坐标值不变
            xmax = str( 256 - int(obj_new.find('xmin').text))
            obj_new.find('xmin').text = xmin
            obj_new.find('xmax').text = xmax
            tree.write(new_xml)


"图像加雾化处理"
def AddHaze1(img):
    # img_f = img
    img_f = img / 255.0
    (row, col, chs) = img.shape
    A = 0.6  # 亮度
    beta = 0.08  # 雾的浓度
    size = math.sqrt(max(row, col))  # 雾化尺寸
    center = (row // 2, col // 2)  # 雾化中心
    for j in range(row):
        for l in range(col):
            d = -0.04 * math.sqrt((j - center[0]) ** 2 + (l - center[1]) ** 2) + size
            td = math.exp(-beta * d)
            img_f[j][l][:] = img_f[j][l][:] * td + A * (1 - td)

    img_f = np.uint8(img_f * 255)
    return img_f

def AddHaze2(img):
    img_h =img/255.0
    # img_h = img
    A = np.random.uniform(0.6, 0.95)
    t = np.random.uniform(0.30, 0.60)
    img_h = img_h * t + A * (1 - t)
    img_h = np.uint8(img_h * 255)

    return img_h

def AddHaze(img):
    l = np.random.uniform(0, 1)
    if l > 0.7:
        img_T = AddHaze1(img)
    else:
        img_T = AddHaze2(img)

    return img_T

"添加椒盐噪声"
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


def resize_xml(xml_file_name,new_xml_name,original_jpg_name,original_png_name,new_jpg_name,new_png_name,new_size=(224,224),original_size=(256,256)):

    original_jpg =Image.open(original_jpg_name)
    cv_original_jpg = cv2.cvtColor(np.asarray(original_jpg), cv2.COLOR_RGB2BGR)
    cv_original_jpg= cv2.resize(cv_original_jpg,(224,224))
    # 16是边缘填充的长度，图片先resize为224大小，则再变回256，则四边都需要填充16
    new_jpg =cv2.copyMakeBorder(cv_original_jpg, 16, 16, 16, 16, cv2.BORDER_CONSTANT, value=(0, 0, 0))   # 填充保持图片的大小不变。
    cv2.imencode('.jpg', new_jpg)[1].tofile(new_jpg_name)

    original_png =Image.open(original_png_name)
    cv_original_png = cv2.cvtColor(np.asarray(original_png), cv2.COLOR_RGB2BGR)
    cv_original_png = cv2.resize(cv_original_png, (224, 224))
    new_png = cv2.copyMakeBorder(cv_original_png, 16, 16, 16, 16, cv2.BORDER_CONSTANT, value=(0, 0, 0))
    cv2.imencode('.png', new_png)[1].tofile(new_png_name)

    scale_x =new_size[0]/original_size[0]
    scale_y = new_size[1] / original_size[1]

    tree = ET.parse(xml_file_name)         # 因为我对边缘进行了填充，故不改变原图的大小size，只是对xml的边界框坐标进行更改。
    objs = tree.findall('object')
    for ix, obj in enumerate(objs):
        name = obj.find('name').text
        if name in name_classes:            # name_classes的里面放的是你含有多少种类的框的名字，列表
            # print(xml_file)
            obj_new = obj.find('bndbox')
            xmin= str(int(float(obj_new.find('xmin').text) * scale_x)+16)
            xmax = str(int(float(obj_new.find('xmax').text) * scale_x)+16)
            ymin = str(int(float(obj_new.find('ymin').text) * scale_y)+16)
            ymax = str(int(float(obj_new.find('ymax').text) * scale_y)+16)
            obj_new.find('xmin').text = xmin
            obj_new.find('xmax').text = xmax
            obj_new.find('ymin').text = ymin
            obj_new.find('ymax').text = ymax
            tree.write(new_xml_name)


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'
        # print(orginal_image_id)
        # print(os.path.join(resize_jpg_path,orginal_image_id))
        # print(os.path.join(resize_png_path, orginal_png_id))
        # print(os.path.join(resize_xml_path, orginal_xml_id))
        image=Image.open(os.path.join(resize_jpg_path,orginal_image_id))
        png = Image.open(os.path.join(resize_png_path, orginal_png_id))
        tree = ET.parse(os.path.join(resize_xml_path, orginal_xml_id))

        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)

        image.save(new_jpg_name)
        png.save(new_png_name)
        tree.write(new_xml_name)
        num =num +1
        print(new_jpg_name)

    "批处理原始的数据集，进行水平的翻转"
    for i in orginal_jpg_files:
        orginal_image_id = i
        orginal_png_id = i.split('.')[0] + '.png'
        orginal_xml_id = i.split('.')[0] + '.xml'
        # print(orginal_image_id)
        # print(os.path.join(resize_jpg_path, orginal_image_id))
        # print(os.path.join(resize_png_path, orginal_png_id))
        # print(os.path.join(resize_xml_path, orginal_xml_id))
        # image = Image.open(os.path.join(resize_jpg_path, orginal_image_id))
        # png = Image.open(os.path.join(resize_png_path, orginal_png_id))
        # tree = ET.parse(os.path.join(resize_xml_path, orginal_xml_id))

        orginal_jpg_name = os.path.join(resize_jpg_path, orginal_image_id)
        orginal_png_name = os.path.join(resize_png_path, orginal_png_id)
        orginal_xml_name=os.path.join(resize_xml_path, orginal_xml_id)
        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)

        flip_xml(orginal_xml_name, new_xml_name)
        flip_horizontal(orginal_jpg_name, new_jpg_name, orginal_png_name, new_png_name)
        num = 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'
        # image = Image.open(os.path.join(resize_jpg_path, orginal_image_id))
        img = cv2.imdecode(np.fromfile(os.path.join(resize_jpg_path, orginal_image_id), dtype=np.uint8), flags=-1)
        # image = Image.open(os.path.join(resize_jpg_path, orginal_image_id))
        # img = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
        png = Image.open(os.path.join(resize_png_path, orginal_png_id))
        tree = ET.parse(os.path.join(resize_xml_path, orginal_xml_id))

        "此处调用雾化操作"
        new_img = AddHaze(img)

        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)

        cv2.imencode('.jpg', new_img)[1].tofile(new_jpg_name)
        png.save(new_png_name)
        tree.write(new_xml_name)
        num = 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

    # 批处理resize图像，边缘填充处理
    for i in orginal_jpg_files:
        orginal_image_id = i
        orginal_png_id = i.split('.')[0] + '.png'
        orginal_xml_id = i.split('.')[0] + '.xml'
        orginal_jpg_name = os.path.join(resize_jpg_path, orginal_image_id)
        orginal_png_name = os.path.join(resize_png_path, orginal_png_id)
        orginal_xml_name = os.path.join(resize_xml_path, orginal_xml_id)


        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)

        resize_xml(xml_file_name=orginal_xml_name,new_xml_name=new_xml_name,original_jpg_name=orginal_jpg_name,
                   new_jpg_name=new_jpg_name,original_png_name=orginal_png_name,new_png_name=new_png_name)
        num=num+1