处理数据集的python代码

image与label

缩放一个文件夹中的所有图片

import numpy as np
import glob
import os
from PIL import Image
#image_path表示图片所在文件夹的路径
#new_image_path 表示存新图片的文件夹路径
#new_size 表示长度为2的列表，两元素分别是宽和高
def img_resize(image_path,new_image_path,new_size):
    # 找出image_path下的所有.jpg/png/jpeg图片路径，并排序
    imgs = sorted(glob.glob(image_path + "*.jpg") )
    print(len(imgs))

    for img in imgs:
        image = Image.open(img)
        image=image.resize(new_size,Image.BILINEAR)
        print(image.size)

        # os.path.splitext() 将文件名和扩展名分开
        # os.path.split() 返回文件的路径和文件名
        img_name=os.path.split(img)[1]
        print(img_name)
        image.save(new_image_path+img_name) #保存图片

使用一个阈值将文件夹中的标签图变为二值图

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


#该函数用于把只有一个类别的标签图变为二值图（因为用ps弄出的不是二值图，边缘部分会有灰度值过渡的情况，除非用labelme制作标签图）
#max_value表示二值图的大的那个值，默认为1  , samll_value表示二值图的小的那个值，默认为0
#new_size表示是否缩放图片，默认为False,若要缩放图片就传一个长度为2的列表（宽和高）
def annotation_prepare(ann_path,new_ann_path,threshold,new_size=False,max_value=1,samll_value=0): #用来将标图变为1或255的二值图
    anns = sorted(glob.glob(ann_path + "*.jpg"))
    for ann in anns:
        ann_name = os.path.split(ann)[1] #得到带扩展名的图片名字用于后面的图片命名
        print(ann_name)
        ann = Image.open(ann)
        ann = ann.convert("L") #变为灰度图
        if (new_size!=False):  #new_size!=False表示需要缩放图片，
            w, h = new_size
            ann=ann.resize((w,h),Image.BILINEAR)

        else:
            w,h=ann.size

        ann=np.array(ann)
        #进行阈值转化
        new_ann=(ann>threshold)*max_value
        new_ann=new_ann+(ann>=threshold)*samll_value

        new_ann = new_ann.astype(np.uint8)  # 再还原成uint8格式，图片需要以这种格式来保存
        new_ann = Image.fromarray(new_ann)  # 从数组形式转为图片形式，这样可以使用save函数
        new_ann.save(new_ann_path+ann_name) #保存图片

给标签图重命名

#用于标签图比原图名字（不算扩展名）的基础上多了一些后缀的情况，且一张原图对应1或多张标签图。可将标签图的多余的一些后缀去掉，用统一的_0N代替（如原图名为10.jpg  而其对应的标签图为10_water_1,10_w_2 ， 可将标签图重命名为10_01,10_02）。
#把多余的原图也也去掉，因为数据集比较乱的话，有的原图可能对应0张标签图
#img_path 原图所在文件夹的路径
#ann_path 标签图所在文件夹的路径
def annNameChange(img_path,ann_path,new_img_path,new_ann_path): #用来将标图变为1或255的二值图new_ann_path,
    imgs = sorted(glob.glob(img_path+"*.jpg"))

    for img in imgs:
        img_save=Image.open(img)  #读取img图片用于待会转移路径保存
        img = os.path.split(img)[1] #得到带后缀的图片名字
        img_name=os.path.splitext(img)[0]  #去掉扩展名

        print(img_name)

        ann_list=glob.glob(ann_path+img_name+"[_.]*")   #找出这张图片对应的所有标签图。 [ ] 表示匹配[]中的任意字符 。 这里需要根据自己图片的名字特征重改匹配规则
        print(len(ann_list))

        i=1
        for ann in ann_list:

            ann = Image.open(ann) #读取图片是为了下面的重新保存

            new_img_name=new_img_path+'/'+img_name+'.jpg'
            new_ann_name=new_ann_path+'/'+img_name+"_0{}.jpg".format(i) #这也是重点。


            print(new_ann_name)
            i =i+1
    
            img_save.save(new_img_name) #重新保存图片
            ann.save(new_ann_name)

给原图与标签图都重命名并保存

#和上面类似，但是这里是把原图与标签图全部重新命名，从1开始，第n张原图命名为n.jpg，而其对应的标签图命名为n_0x.jpg
def annNameChange_new(img_path,ann_path,new_img_path,new_ann_path): #用来将标图变为1或255的二值图new_ann_path,
    imgs = sorted(glob.glob(img_path+"*.jpg"))
    i_name=0    #表示正处理第i_name张图，命名是也以i_name命名
    for img in imgs:
        i_name = i_name+1
        img_save=Image.open(img)  #读取img图片用于待会转移路径保存
        img = os.path.split(img)[1] #得到带后缀的图片名字
        img_name=os.path.splitext(img)[0]  #去掉扩展名
        print(img_name)

        ann_list=glob.glob(ann_path+img_name+"[_.]*")   #找出这张图片对应的所有标签图。 [ ] 表示匹配[]中的任意字符
        print(len(ann_list))

        i=1
        for ann in ann_list:
            ann = Image.open(ann) #读取图片是为了下面的重新保存


            new_img_name=new_img_path+'/{}.jpg'.format(i_name)
            new_ann_name=new_ann_path+'/{}_0{}.jpg'.format(i_name,i) #这也是重点。

            print(new_img_path)
            print(new_ann_name)
            i =i+1

            img_save.save(new_img_name) #重新保存图片
            ann.save(new_ann_name)

取出数据集的图片、对应label图以及对应的名字分别放置于list中

#下面是我自定义的函数，用于只取出图片用测试，并且把对应的名字取出来用保存最后的结果。用于water数据集
#下面是我自定义的函数，用于只取出图片用测试，并且把对应的名字取出来用保存最后的结果。用于water数据集
def get_image(img_path,ann_path):
    img_list = glob.glob(img_path+r'\*')
    mask_list=glob.glob(ann_path+r'\*')
    print(len(img_list))
    print('mask_list',len(mask_list))
    imgs = []
    names=[]
    masks=[]

    for img_path in img_list:
        names.append(os.path.basename(img_path))
        img=Image.open(img_path)
        imgs.append(img)
    for mask_path in mask_list:
        mask = Image.open(mask_path)  # 读取图片，读取
        masks.append(mask)
    return names,imgs,masks

同时把img与对应的mask resize 到统一尺寸并保存

"""用来同时把img与对应的mask resize 到统一尺寸"""
def resize_image(args):
        img_names,imgs=get_image(args.imgs_dir) #get_image在上文
        mask_names,masks=get_image(args.masks_dir)
        assert len(imgs)==len(masks), 'len(imgs_dir) and len(masks_dir) must be the same'
        for i in range(len(imgs)):
            re_img = imgs[i].resize((args.w, args.h), Image.BICUBIC) # args.w,args.h指的是宽和高
            re_mask= masks[i].resize((args.w, args.h), Image.BICUBIC) #插值以后还需要再经过阈值处理
            re_label_1,re_label_255=make_label(re_mask,args.Threshold)  # make_label 在下文
            re_img.save(os.path.join(args.re_img_dir, img_names[i]))
            re_label_1.save(os.path.join(args.re_mask_dir_new, mask_names[i]))
            re_label_255.save(os.path.join(args.re_mask_dir, mask_names[i]))


#这个函数用于把mask图经过阈值处理，然后分别返回灰度值为1 和 灰度值为255的图
def make_label(mask,Threshold):
    mask = mask.convert("L")  # 把打开的图片变为灰度图
    mask = np.array(mask).astype(int)  # 把打开的图片变成数组，并且把类型从uint8变为了int32,因为下面要进行像素的相加，使用uint8会超出范围的


    label_1 = (mask > Threshold) *1  # 选取一个阈值，把像素变为1，如果有多个类，其他的应该变为2,3,4等等
    label_255=label_1*255
    # label2 = np.zeros((512,512))
    # label_1=label1+label2
    label_1 = label_1.astype(np.uint8)  # 再还原成uint8格式，图片需要以这种格式来保存
    label_255 = label_255.astype(np.uint8)  # 再还原成uint8格式，图片需要以这种格式来保存
    # label_1=np.where(mask > Threshold, 1, 0)
    # label_255 = label_1 * 255
    label_1 = Image.fromarray(label_1)  # 从数组形式转为图片形式，这样可以使用save函数
    label_255= Image.fromarray(label_255)
    return label_1,label_255

用于把一个数据集划分为训练集、验证集、测试集存起来（按比例随机）

import random
def split_dataset(img_path,ann_path,train_dir,val_dir,test_dir=None,is_split_test=False):
    img_names,imgs,masks=get_image(img_path,ann_path)  # get_image在上文中有

	#创建一些目录
    train_imgs_dir=os.path.join(train_dir,'imgs')
    train_anns_dir = os.path.join(train_dir,'anns')
    val_imgs_dir=os.path.join(val_dir,'imgs')
    val_anns_dir = os.path.join(val_dir,'anns')
    mkdir(train_imgs_dir)
    mkdir(train_anns_dir)
    mkdir(val_imgs_dir)
    mkdir(val_anns_dir)
    if test_dir:
        test_imgs_dir = os.path.join(train_dir, 'imgs')
        test_anns_dir = os.path.join(train_dir, 'anns')
        mkdir(test_imgs_dir)
        mkdir(test_anns_dir)

    assert len(masks) == len(imgs), 'len(imgs_dir) and len(masks_dir) must be the same'
    indexs = [i for i in range(len(imgs))]
    print(len(imgs))
    random.shuffle(indexs)  # 打乱索引
    print(indexs)
    print(type(indexs))


    if is_split_test:  # 若是要划分测试集，则比例为6：2：2
        cnt_train = int(round(len(imgs) * 0.6, 0))  # 四舍五入，0表示保留0个小数点
        print(cnt_train)
        cnt_val = int(round(len(imgs) * 0.2, 0))  # 因为在需要做索引，所以强制转换为整型
        print(cnt_val)
        cnt_test = int(round(len(imgs) * 0.2, 0))

        # 下面使用列表切片得到train/val/test数据的索引
        train_indexs = indexs[0:cnt_train]
        print(train_indexs)
        val_indexs = indexs[cnt_train:(cnt_train + cnt_val):1]
        test_indexs = indexs[(cnt_train + cnt_val):]

    else:  # 只划分训练集与验证集, 这里使用了 7:3
        cnt_train = int(round(len(imgs) * 0.7, 0))  # 四舍五入，0表示保留0个小数点
        cnt_val = int(round(len(imgs) * 0.3, 0))  # 因为在需要做索引，所以强制转换为整型
        train_indexs = indexs[0:cnt_train]
        print(train_indexs)
        val_indexs = indexs[cnt_train:]

	# 保存图片
    for i in train_indexs:
        imgs[i].save(os.path.join(train_imgs_dir, img_names[i]))
        # masks[i].save(os.path.join(args.train_dir, 'annotation', mask_names[i]))
        masks[i].save(os.path.join(train_anns_dir,img_names[i]))
    for i in val_indexs:
        imgs[i].save(os.path.join(val_imgs_dir,img_names[i]))
        # masks[i].save(os.path.join(args.val_dir, 'annotation', mask_names[i]))
        masks[i].save(os.path.join(val_anns_dir, img_names[i]))
    if is_split_test:
        for i in test_indexs:
            imgs[i].save(os.path.join(test_imgs_dir, img_names[i]))
            # masks[i].save(os.path.join(args.val_dir, 'annotation', mask_names[i]))
            masks[i].save(os.path.join(test_anns_dir, img_names[i]))


def mkdir(path):   #用于创建目录
    # 引入模块
    # 去除首位空格
    path = path.strip()
    # 去除尾部 \ 符号
    path = path.rstrip("\\")
    # 判断路径是否存在
    # 存在     True
    # 不存在   False
    isExists = os.path.exists(path)
    # 判断结果
    if not isExists:
        # 如果不存在则创建目录
        # 创建目录操作函数
        os.makedirs(path)
        return True
    else:
        # 如果目录存在则不创建，并提示目录已存在
        return False

语义分割数据集数据增强

import Augmentor
from glob import glob
import os
import random
from PIL import Image
import cv2
import numpy as np

def augmentor(image_path,label_path,num):
    #num表示想要生成的图片张数
    p = Augmentor.Pipeline(image_path)
    # Point to a directory containing ground truth data.
    # Images with the same file names will be added as ground truth data
    # and augmented in parallel to the original data.
    p.ground_truth(label_path)
    # Add operations to the pipeline as normal:
    p.rotate(probability=0.5, max_left_rotation=10, max_right_rotation=10)
    p.flip_left_right(probability=0.5)
    p.zoom_random(probability=0.5, percentage_area=0.8) #随机放大放小，但保持原尺寸
    p.flip_top_bottom(probability=0.5)
    p.rotate_random_90(probability=0.5)
    p.random_distortion(probability=0.2, grid_width=4, grid_height=4, magnitude=8)  #表示随机扭曲与些地方
    p.sample(num)

合并label图片

'''对于一个Groundtruth中把每一个对象都单独标注到一个图片中。该函数用于合并mask到一张图片中'''

#下面是我自定义的函数，用于只取出图片用测试，并且把对应的名字取出来用保存最后的结果。用于water数据集
import glob
import os
import numpy as np
from PIL import Image

def merge(imgpath_list):
    mask=Image.open(imgpath_list[0])
    mask = mask.convert('L')  #让标签图变为1个通道
    mask=np.array(mask)
    # print(mask.shape)
    for _mask in imgpath_list[1:]:
        _mask = Image.open(_mask)
        _mask = _mask.convert('L')
        _mask = np.array(_mask)

        mask = np.maximum(mask, _mask)  #合并两张图片
    mask=Image.fromarray(mask)
    return mask

def merge_mask(img_path,ann_path,new_annpath):
    img_list = glob.glob(img_path+'\*')
    print(len(img_list))

    for img_path in img_list:
        name=os.path.basename(img_path)
        ann_img=os.path.join(ann_path,name)
        mask_list=glob.glob(ann_img[0:-4]+'*')  #制定匹配规则，得到img_path所对应的所有Mask图
        print(name,mask_list)
        if(len(mask_list)>1):
            mask=merge(mask_list)
        else:
            mask=Image.open(mask_list[0])

        mask.save(os.path.join(new_annpath,name))

使用PIL库判断图片尺寸并旋转

img=Image.open(img_path)
h,w=img.size
if h>w:
	img=img.transpose(Image.ROTATE_90)

# 如果一个数据集中，图片都是同一尺寸，但有的是竖着的，有的是横着的，可使用下面的代码其他们统一
import os

from PIL import Image
cascade_mask_rcnn_water_train_imgs_dir=r'F:\Dataset\train\color_img'
cascade_mask_rcnn_water_train_marks_dir=r'F:\Jun\Dataset\train\binary_img'
cascade_mask_rcnn_water_val_imgs_dir=r'F:\Jun\Dataset\val\color_img_all'
cascade_mask_rcnn_water_val_marks_dir=r'F:\Jun\Dataset\val\binary_img_all'
def rotate_image(img_path,save_path):
    imgs_path = os.listdir(img_path)
    for r in imgs_path:
        img=os.path.join(img_path,r)
        img=Image.open(img)
        w, h = img.size
        if w>h:
            img = img.transpose(Image.ROTATE_90)
        img.save(os.path.join(save_path,r))