语义分割数据集抽取特定类

以VOC2012数据集为基准，制作成此类格式：# pascal_voc_seg#+ VOCdevkit#+ VOC2012#+ JPEGImages#+ SegmentationClassVOC抽取人person类import xml.etree.ElementTree as ETimport pickleimport osimport sysimport argparseimport shuti

三寸光阴___

455人浏览 · 2021-06-03 09:03:40

三寸光阴___ · 2021-06-03 09:03:40 发布

以VOC2012数据集为基准，制作成此类格式：

# pascal_voc_seg
#       + VOCdevkit
#         + VOC2012
#           + JPEGImages
#           + SegmentationClass

VOC抽取人person类

import xml.etree.ElementTree as ET
import pickle
import os
import sys
import argparse
import shutil
import numpy as np
import cv2

classes = ['person']  

def copy(read_dir, save_dir, xml_name, img_name):
    src_xml_dir = os.path.join(read_dir,"VOCdevkit","VOC2012","Annotations")
    src_image_dir = os.path.join(read_dir,"VOCdevkit","VOC2012","JPEGImages")
    src_SegmentationClassRaw_dir = os.path.join(read_dir,"VOCdevkit","VOC2012","SegmentationClassRaw")

    save_xml_dir = os.path.join(save_dir,"VOCdevkit","VOC2012","Annotations")
    save_image_dir = os.path.join(save_dir,"VOCdevkit","VOC2012","JPEGImages")
    save_SegmentationClassRaw_dir = os.path.join(save_dir,"VOCdevkit","VOC2012","SegmentationClassRaw")

    in_file = open(os.path.join(src_xml_dir,xml_name))
    tree=ET.parse(in_file)
    root = tree.getroot()

    flag = False
    segmented = root.find("segmented").text
    if int(segmented) == 1:
        for obj in root.iter('object'):
            cls = obj.find('name').text
            if cls in classes:
                flag = True
    if flag == True:
        shutil.copy(os.path.join(src_image_dir,img_name),os.path.join(save_image_dir,img_name))
        shutil.copy(os.path.join(src_SegmentationClassRaw_dir,img_name[:-4]+".png"),os.path.join(save_SegmentationClassRaw_dir,img_name[:-4]+".png"))

def correct(save_dir):
    save_SegmentationClassRaw_dir = os.path.join(save_dir,"VOCdevkit","VOC2012","SegmentationClassRaw")
    for image_name in os.listdir(save_SegmentationClassRaw_dir):
        print(image_name,end="\r")
        full_image_path = os.path.join(save_SegmentationClassRaw_dir,image_name)
        img = cv2.imread(full_image_path,-1)
        if len(img.shape) !=2:
            assert False
        img[np.multiply(img!=15, img!=255)]= 0
        # img[img!=15] = 0
        img[img==15] = 1
        cv2.imwrite(full_image_path,img)

def correct_txt(src_full_txt_name,save_dir):
    src_file = open(src_full_txt_name,'r')
    src_name_list = src_file.read().splitlines()
    save_txt_dir = os.path.join(save_dir,"VOCdevkit","VOC2012","ImageSets","Segmentation")
    save_SegmentationClassRaw_dir = os.path.join(save_dir,"VOCdevkit","VOC2012","SegmentationClassRaw")
    
    save_txt = os.path.join(save_txt_dir,src_full_txt_name.split("/")[-1])
    with open(save_txt,'w') as save_file:
        for image_name in os.listdir(save_SegmentationClassRaw_dir):
            if os.path.splitext(image_name)[0] in src_name_list:
                save_file.write(os.path.splitext(image_name)[0])
                save_file.write("\n")
                # pass


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-read_dir', type=str, help='image dir')
    parser.add_argument('-save_dir', type=str, help='image dir')
    parser.add_argument('-image_extension', type=str, default=".jpg", help='image extension')
    opt = parser.parse_args()
    
    read_dir = opt.read_dir
    save_dir = opt.save_dir
    image_extension = opt.image_extension

    src_xml_dir = os.path.join(read_dir,"VOCdevkit","VOC2012","Annotations")
    scr_txt_dir = os.path.join(read_dir,"VOCdevkit","VOC2012","ImageSets","Segmentation")

    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        os.makedirs(os.path.join(save_dir,"VOCdevkit","VOC2012"))
        os.makedirs(os.path.join(save_dir,"VOCdevkit","VOC2012","Annotations"))
        os.makedirs(os.path.join(save_dir,"VOCdevkit","VOC2012","ImageSets"))
        os.makedirs(os.path.join(save_dir,"VOCdevkit","VOC2012","ImageSets","Segmentation"))
        os.makedirs(os.path.join(save_dir,"VOCdevkit","VOC2012","JPEGImages"))
        os.makedirs(os.path.join(save_dir,"VOCdevkit","VOC2012","SegmentationClass"))
        os.makedirs(os.path.join(save_dir,"VOCdevkit","VOC2012","SegmentationObject"))
        os.makedirs(os.path.join(save_dir,"VOCdevkit","VOC2012","SegmentationClassRaw"))

    print("start select image")
    for xml_name in os.listdir(src_xml_dir):
        img_name = xml_name.replace(".xml",image_extension)    
        copy(read_dir,save_dir,xml_name,img_name)
        print(img_name,end="\r")

    print("start correct image")
    correct(save_dir)

    print("start correct txt")
    for txt_name in os.listdir(scr_txt_dir):
        src_full_txt_name = os.path.join(scr_txt_dir,txt_name)
        correct_txt(src_full_txt_name,save_dir)

    print("complete!")
    # python get_voc_personv2.py -read_dir /home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/pascal_voc_seg -save_dir /home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/pascal_voc_seg_person

cityscapes抽取天空（sky）类

import xml.etree.ElementTree as ET
import pickle
import os
import sys
import argparse
import shutil
import numpy as np
import cv2
import imgviz


def copy_and_correct(read_dir_src, read_dir_label, save_dir_jpg, save_dir_label, label_id): # 23:sky
    
    for image_label_name in os.listdir(read_dir_label):
        if image_label_name.endswith("labelIds.png"):
            img = cv2.imread(os.path.join(read_dir_label,image_label_name),-1)
            if len(img.shape)==3:
                img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
            if np.any(img==23) == True:
                img[img!=23]= 0
                img[img==23] = 2
                #print(image_label_name)
                cv2.imwrite(os.path.join(save_dir_label,image_label_name[:-20]+".png"),img)
                read_dir_png_name = image_label_name[:-20] + "_leftImg8bit.png" # _gtFine_labelIds.png
                if os.path.isfile(os.path.join(read_dir_src, read_dir_png_name)):
                    print(os.path.join(save_dir_jpg,read_dir_png_name[:-16]+".jpg"))
                    img_png = cv2.imread(os.path.join(read_dir_src,read_dir_png_name))
                    cv2.imwrite(os.path.join(save_dir_jpg,read_dir_png_name[:-16]+".jpg"),img_png)
            # else:
            #     print("---",image_label_name)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-read_dir_src', type=str, help='image dir') # /mldb/dataset/cityscapes/leftImg8bit/test/berlin/
    parser.add_argument('-read_dir_label', type=str, help='image dir') # /mldb/dataset/cityscapes/gtFine/test/berlin/
    parser.add_argument('-save_dir_jpg', type=str, default="/home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/custom_voc_seg/VOCdevkit/VOC2012/JPEGImages", help='image dir')
    parser.add_argument('-save_dir_label', type=str, default="/home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/custom_voc_seg/VOCdevkit/VOC2012/SegmentationClassRaw", help='image dir')
    parser.add_argument('-label_id', type=str, help='image dir')
    parser.add_argument('-image_extension', type=str, default=".jpg", help='image extension')

    opt = parser.parse_args()
    
    read_dir_src = opt.read_dir_src
    read_dir_label = opt.read_dir_label
    save_dir_jpg = opt.save_dir_jpg
    save_dir_label = opt.save_dir_label
    label_id = opt.label_id
    image_extension = opt.image_extension

    if not os.path.exists(save_dir_jpg):
        os.makedirs(save_dir_jpg)
    if not os.path.exists(save_dir_label):
        os.makedirs(save_dir_label)

    print("start copy and correct image")

    #  python get_voc_class_cityscapes.py -read_dir_src /mldb/dataset/cityscapes/leftImg8bit -read_dir_label /mldb/dataset/cityscapes/gtFine
    for dataset in os.listdir(read_dir_src):
        for city in os.listdir(os.path.join(read_dir_src, dataset)):
            read_dir_src_ = os.path.join(read_dir_src, dataset, city)
            read_dir_label_ = os.path.join(read_dir_label, dataset, city)
            print("-----:",read_dir_label_)
            copy_and_correct(read_dir_src_, read_dir_label_, save_dir_jpg, save_dir_label, label_id)
            
    #copy_and_correct(read_dir_src, read_dir_label, save_dir_jpg, save_dir_label, label_id)

    #  python get_voc_class_cityscapes.py -read_dir_src /mldb/dataset/cityscapes/leftImg8bit/test/berlin -read_dir_label /mldb/dataset/cityscapes/gtFine/test/berlin

    print("complete!")
    # python get_voc_class_cityscapes.py -read_dir_src /mldb/dataset/cityscapes/leftImg8bit -read_dir_label /mldb/dataset/cityscapes/gtFine

ADE20k抽取天空（sky）、草（grass）类

import xml.etree.ElementTree as ET
import pickle
import os
import sys
import argparse
import shutil
import numpy as np
import cv2
import imgviz

def copy_and_correct(read_dir_src, read_dir_label, save_dir_jpg, save_dir_label, label_id): # 3:sky ,10:grass
    
    for image_label_name in os.listdir(read_dir_label):
        img = cv2.imread(os.path.join(read_dir_label,image_label_name),-1)
        if np.any(img==3) == True or np.any(img==10) == True:
            img[np.multiply(img!=3, img!=10)] = 0
            img[img==3] = 2
            img[img==10] = 3
            
            cv2.imwrite(os.path.join(save_dir_label,image_label_name),img)
            jpg_name = image_label_name.replace("png","jpg")
            print(os.path.join(save_dir_jpg,jpg_name))
            
            shutil.copy(os.path.join(read_dir_src,jpg_name),os.path.join(save_dir_jpg,jpg_name))
        # else:
        #     print("---",image_label_name)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-read_dir_src', type=str, help='image dir') 
    parser.add_argument('-read_dir_label', type=str, help='image dir')
    parser.add_argument('-save_dir_jpg', type=str, default="/home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/custom_voc_seg/VOCdevkit/VOC2012/JPEGImages", help='image dir')
    parser.add_argument('-save_dir_label', type=str, default="/home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/custom_voc_seg/VOCdevkit/VOC2012/SegmentationClassRaw", help='image dir')
    parser.add_argument('-label_id', type=str, help='image dir')
    parser.add_argument('-image_extension', type=str, default=".jpg", help='image extension')

    opt = parser.parse_args()
    
    read_dir_src = opt.read_dir_src
    read_dir_label = opt.read_dir_label
    save_dir_jpg = opt.save_dir_jpg
    save_dir_label = opt.save_dir_label
    label_id = opt.label_id
    image_extension = opt.image_extension

    if not os.path.exists(save_dir_jpg):
        os.makedirs(save_dir_jpg)
    if not os.path.exists(save_dir_label):
        os.makedirs(save_dir_label)

    print("start copy and correct image")

    #  python get_voc_class_ade20k.py -read_dir_src /home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/ADE20K/ADEChallengeData2016/images -read_dir_label /home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/ADE20K/ADEChallengeData2016/annotations
    for dataset in os.listdir(read_dir_src):
        read_dir_src_ = os.path.join(read_dir_src, dataset)
        read_dir_label_ = os.path.join(read_dir_label, dataset)
        print("-----:",read_dir_label_)
        copy_and_correct(read_dir_src_, read_dir_label_, save_dir_jpg, save_dir_label, label_id)

    print("complete!")
    # python get_voc_class_ade20k.py -read_dir_src /home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/ADE20K/ADEChallengeData2016/images -read_dir_label /home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/ADE20K/ADEChallengeData2016/annotations

LIP抽取可能带有皮肤的类（脸，左右胳膊，左右腿）

import xml.etree.ElementTree as ET
import pickle
import os
import sys
import argparse
import shutil
import numpy as np
import cv2
import imgviz

def copy_and_correct(read_dir_src, read_dir_label, save_dir_jpg, save_dir_label, label_id): # 13:face ,14: Left-arm ,15: Right-arm,16: Left-leg,17:Right-leg

    for image_label_name in os.listdir(read_dir_label):
        img = cv2.imread(os.path.join(read_dir_label,image_label_name),-1)
        if np.any(img==13) == True or np.any(img==14) == True or np.any(img==15) == True or np.any(img==16) == True or np.any(img==17) == True:
            img[img==13] = 1
            img[img==14] = 1
            img[img==15] = 1
            img[img==16] = 1
            img[img==17] = 1
            img[img!=1] = 0
            
            cv2.imwrite(os.path.join(save_dir_label,image_label_name),img)
            jpg_name = image_label_name.replace("png","jpg")
            print(os.path.join(save_dir_jpg,jpg_name))
            
            shutil.copy(os.path.join(read_dir_src,jpg_name),os.path.join(save_dir_jpg,jpg_name))
        # else:
        #     print("---",image_label_name)

    # viz = imgviz.label2rgb(
    #             label=lbl,
    #             img=imgviz.rgb2gray(img),
    #             font_size=15,
    #             label_names=class_names,
    #             loc="rb",
    #         )
    #         imgviz.io.imsave(out_viz_file, viz)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-read_dir_src', type=str, help='image dir') 
    parser.add_argument('-read_dir_label', type=str, help='image dir')
    parser.add_argument('-save_dir_jpg', type=str, default="/home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/custom_voc_seg/VOCdevkit/VOC2012/JPEGImages", help='image dir')
    parser.add_argument('-save_dir_label', type=str, default="/home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/custom_voc_seg/VOCdevkit/VOC2012/SegmentationClassRaw", help='image dir')
    parser.add_argument('-label_id', type=str, help='image dir')
    parser.add_argument('-image_extension', type=str, default=".jpg", help='image extension')

    opt = parser.parse_args()
    
    read_dir_src = opt.read_dir_src
    read_dir_label = opt.read_dir_label
    save_dir_jpg = opt.save_dir_jpg
    save_dir_label = opt.save_dir_label
    label_id = opt.label_id
    image_extension = opt.image_extension

    if not os.path.exists(save_dir_jpg):
        os.makedirs(save_dir_jpg)
    if not os.path.exists(save_dir_label):
        os.makedirs(save_dir_label)

    print("start copy and correct image")

    #  
    
    read_dir_src_ = read_dir_src
    read_dir_label_ = read_dir_label
    print("-----:",read_dir_label_)
    copy_and_correct(read_dir_src_, read_dir_label_, save_dir_jpg, save_dir_label, label_id)

    print("complete!")
    # python get_voc_class_lip.py -read_dir_src /data/xywang/dataset/LIP/src_image/ -read_dir_label /data/xywang/dataset/LIP/label_image/ -save_dir_jpg /home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/skin_voc_seg_lip/VOCdevkit/VOC2012/JPEGImages/ -save_dir_label /home/xywang/code/Semantic_segmentation/tf_DeepLabV3+/code/research/deeplab/datasets/skin_voc_seg_lip/VOCdevkit/VOC2012/SegmentationClassRaw/