Dagm数据集转换为VOC格式

最新推荐文章于 2024-05-30 11:43:13 发布

Finddog

最新推荐文章于 2024-05-30 11:43:13 发布

阅读量145

点赞数 1

文章标签： python 图像处理

本文链接：https://blog.csdn.net/Finddog/article/details/132153820

版权

Dagm数据集转换为VOC格式

由于dagm数据集的标签过于简陋，是椭圆形状的标签，最近实验室同学重新标注了dagm数据集，用来跑其他代码时，都是采用的VOC格式，因此写了以下脚本，将DAGM格式的数据集转换为VOC格式，VOC格式的数据集各个文件夹内容不懂的可以参照 VOC数据集格式解读，使用时直接指定DAGM数据集的根目录和需要存放转换后数据集的文件夹，另外创建一个palette.json,并将内容输入就可以了。

注：class序号从1开始

import os
import cv2
import numpy as np
import os
import shutil
from tqdm import tqdm
from PIL import Image
import json
# voc 格式介绍请参照 https://www.cnblogs.com/gshang/p/12964606.html
def makDirs(path):
    if not(os.path.exists(path)):
        os.makedirs(path)

def creat_xml(xmldir,file,classname,xmin,ymin,xmax,ymax):
    xml_file = open(os.path.join(xmldir,file+".xml"), 'w')
    xml_file.write('<annotation>\n')
    xml_file.write('\t<folder>'+"JPEGImages"+'</folder>\n')
    xml_file.write('\t<filename>' + file+".jpg" + '</filename>\n')
    # xml_file.write('\t<path>' + imagePath + '</path>\n')
    xml_file.write('\t<source>\n')
    xml_file.write('\t\t<database>' + 'Unknown' + '</database>\n')
    # xml_file.write('\t\t<annotation>' + 'simple' + '</annotation>\n')
    # xml_file.write('\t\t<image>flickr</image>\n')
    # xml_file.write('\t\t<flickrid>325991873</flickrid>\n')
    xml_file.write('\t</source>\n')
    # xml_file.write('\t<owner>\n')
    # xml_file.write('\t\t<flickrid>archin</flickrid>\n')
    # xml_file.write('\t\t<name>?</name>\n')
    # xml_file.write('\t</owner>\n')
    xml_file.write('\t<size>\n')
    xml_file.write('\t\t<width>' + "512" + '</width>\n')
    xml_file.write('\t\t<height>' + "512" + '</height>\n')
    xml_file.write('\t\t<depth>3</depth>\n')
    xml_file.write('\t</size>\n')
    xml_file.write('\t<segmented>0</segmented>\n')
    # write the region of text on xml file
    for ann in [1]:
        xml_file.write('\t<object>\n')
        xml_file.write('\t\t<name>' + classname + '</name>\n')
        xml_file.write('\t\t<pose>Unspecified</pose>\n')
        xml_file.write('\t\t<truncated>0</truncated>\n')
        xml_file.write('\t\t<difficult>0</difficult>\n')
        xml_file.write('\t\t<bndbox>\n')
        xml_file.write('\t\t\t<xmin>' + str(int(xmin)) + '</xmin>\n')
        xml_file.write('\t\t\t<ymin>' + str(int(ymin)) + '</ymin>\n')
        xml_file.write('\t\t\t<xmax>' + str(int(xmax)) + '</xmax>\n')
        xml_file.write('\t\t\t<ymax>' + str(int(ymax)) + '</ymax>\n')
        xml_file.write('\t\t</bndbox>\n')
        xml_file.write('\t</object>\n')
    xml_file.write('</annotation>')
    xml_file.close()

def main(dagm_dir,voc_dir):
    # parase dagm
    CLASSES = ('Class1', 'Class2', 'Class3', 'Class4', 'Class5',
               'Class6', 'Class7', 'Class8', 'Class9', 'Class10')
    class2index=dict(zip(CLASSES,range(1,11)))
    dagmdataset=[]
    for classname in CLASSES:
        for datatype in ('Train','Test'):
            imgdir = os.path.join(dagm_dir, classname, datatype)
            label_imgdir = os.path.join(
                dagm_dir, classname, datatype, 'Label')
            imgs = os.listdir(imgdir)
            imgs = [x for x in imgs if '.PNG' in x]
            imgs_with_label = os.listdir(label_imgdir)
            imgs_with_label = [
                x[:-10]+'.PNG' for x in imgs_with_label if '.PNG' in x]
            imgs_no_label = [i for i in imgs if i not in imgs_with_label]
            classdata=dict(
                imgroot=dagm_dir,
                classname=classname,
                datatype=datatype,
                label_data=imgs_with_label,
                unlabel_data=imgs_no_label
            )
            dagmdataset.append(classdata)


    #
    train_list=[]
    test_list=[]
    for classdata in tqdm(dagmdataset):
        class_train_list_label=[]
        class_test_list_label=[]
        for imgname in tqdm(classdata["label_data"]):
            new_name=classdata["classname"][5:].rjust(2,'0')+"_"+imgname.split(".")[0]+".jpg"
            if classdata["datatype"]=="Train":
                class_train_list_label.append(new_name)
                train_list.append(new_name)
            else:
                class_test_list_label.append(new_name)
                test_list.append(new_name)
            imgpath=os.path.join(classdata["imgroot"],classdata["classname"],classdata["datatype"],imgname)
            dst_imgpath=os.path.join(voc_dir,"JPEGImages",new_name)
            src_img=cv2.imread(imgpath)
            cv2.imwrite(dst_imgpath,src_img)
            # shutil.copy(imgpath, dst_imgpath)

            segpath=os.path.join(classdata["imgroot"],classdata["classname"],classdata["datatype"],"Label",imgname[:4]+"_label.PNG")
            dst_segpath=os.path.join(voc_dir,"SegmentationClass",new_name.split(".")[0]+".png")
            mask=cv2.imread(segpath, 0)
            contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
            mask[np.where(mask==255)]=class2index[classdata["classname"]]
            cv2.drawContours(mask, contours, -1, 255, 1)
            cv2.imwrite(dst_segpath,mask)
            points=contours[0][:,0,:]
            xmin,xmax=points[:,0].min(),points[:,0].max()
            ymin,ymax=points[:,1].min(),points[:,1].max()
            creat_xml(os.path.join(voc_dir,"Annotations"),new_name.split(".")[0],classdata["classname"],xmin,ymin,xmax,ymax)
                        
        class_train_list_nolabel=[]
        class_test_list_nolabel=[]
        for imgname in tqdm(classdata["unlabel_data"]):
            
            new_name=classdata["classname"][5:].rjust(2,'0')+"_"+imgname.split(".")[0]+".jpg"
            if classdata["datatype"]=="Train":
                class_train_list_nolabel.append(new_name)
                # train_list.append(new_name)
            else:
                class_test_list_nolabel.append(new_name)
                # test_list.append(new_name)

            imgpath=os.path.join(classdata["imgroot"],classdata["classname"],classdata["datatype"],imgname)
            dst_imgpath=os.path.join(voc_dir,"JPEGImages",new_name)
            src_img=cv2.imread(imgpath)
            cv2.imwrite(dst_imgpath,src_img)
            # shutil.copy(imgpath, dst_imgpath)

            
            # dst_segpath=os.path.join(voc_dir,"SegmentationClass",new_name)
            # mask = np.zeros((512,512))
            # cv2.imwrite(dst_segpath,mask)
        if classdata["datatype"]=="Train":
            with open(os.path.join(voc_dir,"ImageSets","Main",classdata["classname"]+"_train.txt"),"w") as f:
                f.write("\n".join([x.split(".")[0]+" 1" for x in class_train_list_label]))
                f.write("\n")
                f.write("\n".join([x.split(".")[0]+" 0" for x in class_train_list_nolabel]))
        else:
            with open(os.path.join(voc_dir,"ImageSets","Main",classdata["classname"]+"_test.txt"),"w") as f:
                f.write("\n".join([x.split(".")[0]+" 1" for x in class_test_list_label]))
                f.write("\n")
                f.write("\n".join([x.split(".")[0]+" 0" for x in class_test_list_nolabel]))

    with open(os.path.join(voc_dir,"ImageSets","Segmentation","train.txt"),"w") as f:
        f.write("\n".join([x.split(".")[0] for x in train_list]))
    with open(os.path.join(voc_dir,"ImageSets","Segmentation","test.txt"),"w") as f:
        f.write("\n".join([x.split(".")[0] for x in test_list]))


def corlor_segpng(segdir,palette_path):
    with open(palette_path, "rb") as f:
        pallette_dict = json.load(f)
        pallette = []
        for v in pallette_dict.values():
            pallette += v

    files = os.listdir(segdir)
    for file in tqdm(files):
        if ".png" not in file :
            continue
        filepath=os.path.join(segdir,file)
        mask=Image.open(filepath)
        mask = Image.fromarray(np.asarray(mask))
        mask.putpalette(pallette)
        mask.save(filepath)

if __name__ == '__main__':
    dagm_dir="/home/fj/data/DAGM_relabel_me"
    voc_dir ="/home/fj/data/DAGM_VOC_M"

    makDirs(os.path.join(voc_dir,"Annotations"))
    makDirs(os.path.join(voc_dir,"ImageSets","Main"))
    makDirs(os.path.join(voc_dir,"ImageSets","Segmentation"))
    makDirs(os.path.join(voc_dir,"JPEGImages"))
    makDirs(os.path.join(voc_dir,"SegmentationClass"))
    makDirs(os.path.join(voc_dir,"SegmentationObject"))
    main(dagm_dir,voc_dir)
    palette_path = "./palette.json" #调色板
    corlor_segpng(os.path.join(voc_dir,"SegmentationClass"),palette_path)

palette.json 文件内容

{"0": [0, 0, 0], "1": [128, 0, 0], "2": [0, 128, 0], "3": [128, 128, 0], "4": [0, 0, 128], "5": [128, 0, 128], "6": [0, 128, 128], "7": [128, 128, 128], "8": [64, 0, 0], "9": [192, 0, 0], "10": [64, 128, 0], "11": [192, 128, 0], "12": [64, 0, 128], "13": [192, 0, 128], "14": [64, 128, 128], "15": [192, 128, 128], "16": [0, 64, 0], "17": [128, 64, 0], "18": [0, 192, 0], "19": [128, 192, 0], "20": [0, 64, 128], "21": [128, 64, 128], "22": [0, 192, 128], "23": [128, 192, 128], "24": [64, 64, 0], "25": [192, 64, 0], "26": [64, 192, 0], "27": [192, 192, 0], "28": [64, 64, 128], "29": [192, 64, 128], "30": [64, 192, 128], "31": [192, 192, 128], "32": [0, 0, 64], "33": [128, 0, 64], "34": [0, 128, 64], "35": [128, 128, 64], "36": [0, 0, 192], "37": [128, 0, 192], "38": [0, 128, 192], "39": [128, 128, 192], "40": [64, 0, 64], "41": [192, 0, 64], "42": [64, 128, 64], "43": [192, 128, 64], "44": [64, 0, 192], "45": [192, 0, 192], "46": [64, 128, 192], "47": [192, 128, 192], "48": [0, 64, 64], "49": [128, 64, 64], "50": [0, 192, 64], "51": [128, 192, 64], "52": [0, 64, 192], "53": [128, 64, 192], "54": [0, 192, 192], "55": [128, 192, 192], "56": [64, 64, 64], "57": [192, 64, 64], "58": [64, 192, 64], "59": [192, 192, 64], "60": [64, 64, 192], "61": [192, 64, 192], "62": [64, 192, 192], "63": [192, 192, 192], "64": [32, 0, 0], "65": [160, 0, 0], "66": [32, 128, 0], "67": [160, 128, 0], "68": [32, 0, 128], "69": [160, 0, 128], "70": [32, 128, 128], "71": [160, 128, 128], "72": [96, 0, 0], "73": [224, 0, 0], "74": [96, 128, 0], "75": [224, 128, 0], "76": [96, 0, 128], "77": [224, 0, 128], "78": [96, 128, 128], "79": [224, 128, 128], "80": [32, 64, 0], "81": [160, 64, 0], "82": [32, 192, 0], "83": [160, 192, 0], "84": [32, 64, 128], "85": [160, 64, 128], "86": [32, 192, 128], "87": [160, 192, 128], "88": [96, 64, 0], "89": [224, 64, 0], "90": [96, 192, 0], "91": [224, 192, 0], "92": [96, 64, 128], "93": [224, 64, 128], "94": [96, 192, 128], "95": [224, 192, 128], "96": [32, 0, 64], "97": [160, 0, 64], "98": [32, 128, 64], "99": [160, 128, 64], "100": [32, 0, 192], "101": [160, 0, 192], "102": [32, 128, 192], "103": [160, 128, 192], "104": [96, 0, 64], "105": [224, 0, 64], "106": [96, 128, 64], "107": [224, 128, 64], "108": [96, 0, 192], "109": [224, 0, 192], "110": [96, 128, 192], "111": [224, 128, 192], "112": [32, 64, 64], "113": [160, 64, 64], "114": [32, 192, 64], "115": [160, 192, 64], "116": [32, 64, 192], "117": [160, 64, 192], "118": [32, 192, 192], "119": [160, 192, 192], "120": [96, 64, 64], "121": [224, 64, 64], "122": [96, 192, 64], "123": [224, 192, 64], "124": [96, 64, 192], "125": [224, 64, 192], "126": [96, 192, 192], "127": [224, 192, 192], "128": [0, 32, 0], "129": [128, 32, 0], "130": [0, 160, 0], "131": [128, 160, 0], "132": [0, 32, 128], "133": [128, 32, 128], "134": [0, 160, 128], "135": [128, 160, 128], "136": [64, 32, 0], "137": [192, 32, 0], "138": [64, 160, 0], "139": [192, 160, 0], "140": [64, 32, 128], "141": [192, 32, 128], "142": [64, 160, 128], "143": [192, 160, 128], "144": [0, 96, 0], "145": [128, 96, 0], "146": [0, 224, 0], "147": [128, 224, 0], "148": [0, 96, 128], "149": [128, 96, 128], "150": [0, 224, 128], "151": [128, 224, 128], "152": [64, 96, 0], "153": [192, 96, 0], "154": [64, 224, 0], "155": [192, 224, 0], "156": [64, 96, 128], "157": [192, 96, 128], "158": [64, 224, 128], "159": [192, 224, 128], "160": [0, 32, 64], "161": [128, 32, 64], "162": [0, 160, 64], "163": [128, 160, 64], "164": [0, 32, 192], "165": [128, 32, 192], "166": [0, 160, 192], "167": [128, 160, 192], "168": [64, 32, 64], "169": [192, 32, 64], "170": [64, 160, 64], "171": [192, 160, 64], "172": [64, 32, 192], "173": [192, 32, 192], "174": [64, 160, 192], "175": [192, 160, 192], "176": [0, 96, 64], "177": [128, 96, 64], "178": [0, 224, 64], "179": [128, 224, 64], "180": [0, 96, 192], "181": [128, 96, 192], "182": [0, 224, 192], "183": [128, 224, 192], "184": [64, 96, 64], "185": [192, 96, 64], "186": [64, 224, 64], "187": [192, 224, 64], "188": [64, 96, 192], "189": [192, 96, 192], "190": [64, 224, 192], "191": [192, 224, 192], "192": [32, 32, 0], "193": [160, 32, 0], "194": [32, 160, 0], "195": [160, 160, 0], "196": [32, 32, 128], "197": [160, 32, 128], "198": [32, 160, 128], "199": [160, 160, 128], "200": [96, 32, 0], "201": [224, 32, 0], "202": [96, 160, 0], "203": [224, 160, 0], "204": [96, 32, 128], "205": [224, 32, 128], "206": [96, 160, 128], "207": [224, 160, 128], "208": [32, 96, 0], "209": [160, 96, 0], "210": [32, 224, 0], "211": [160, 224, 0], "212": [32, 96, 128], "213": [160, 96, 128], "214": [32, 224, 128], "215": [160, 224, 128], "216": [96, 96, 0], "217": [224, 96, 0], "218": [96, 224, 0], "219": [224, 224, 0], "220": [96, 96, 128], "221": [224, 96, 128], "222": [96, 224, 128], "223": [224, 224, 128], "224": [32, 32, 64], "225": [160, 32, 64], "226": [32, 160, 64], "227": [160, 160, 64], "228": [32, 32, 192], "229": [160, 32, 192], "230": [32, 160, 192], "231": [160, 160, 192], "232": [96, 32, 64], "233": [224, 32, 64], "234": [96, 160, 64], "235": [224, 160, 64], "236": [96, 32, 192], "237": [224, 32, 192], "238": [96, 160, 192], "239": [224, 160, 192], "240": [32, 96, 64], "241": [160, 96, 64], "242": [32, 224, 64], "243": [160, 224, 64], "244": [32, 96, 192], "245": [160, 96, 192], "246": [32, 224, 192], "247": [160, 224, 192], "248": [96, 96, 64], "249": [224, 96, 64], "250": [96, 224, 64], "251": [224, 224, 64], "252": [96, 96, 192], "253": [224, 96, 192], "254": [96, 224, 192], "255": [224, 224, 192]}