[自用代码]根据SSDD official的xml文件生成VOC格式mask

感谢官方提出了inshore offshore的sub set, 希望能够规范后续SAR ship detection的论文.
这里结合链接1 链接2 实现对于ssdd official中船只分割VOC格式的mask生成. 具体代码如下:

# -*- coding: utf-8 -*-
import numpy as np
import cv2
import os
import xml.etree.ElementTree as ET

def create_dir_not_exist(path):
    if not os.path.exists(path):
        os.mkdir(path)

 
def xml2mask(image_path, xml, id2class_dict):
    class2id_dict = id2class_dict
    tree = ET.parse(xml)
    root = tree.getroot()
 
    img_size = cv2.imread(image_path).shape
    height = img_size[0]
    weight = img_size[1]
 
    mask = np.zeros([height, weight], dtype=np.uint8)
 
    for instance in root.iter('object'):
        classname = instance.findall('name')[0].text
        cnt_points = []
        for pt in instance.iter('segm'):
            for xy in pt.findall('point'):
                ptx,pty = float(str(xy.text).split(',')[0]),float(str(xy.text).split(',')[1])
                cnt_points.append([ptx, pty])

        pts = np.asarray([cnt_points], dtype=np.int32)
        cv2.fillPoly(img=mask, pts=pts, color=class2id_dict[str(classname)])
 
    return mask

if __name__=='__main__':


    dir = "Annotations"
    img_dir = 'JPEGImages'
    id2class_dict = {'ship':(1,0,0)}

    out_dir = 'SegmentationClass'
    create_dir_not_exist(out_dir)

    files = os.listdir(dir)
    i = 0
    for file in files:
        xml_path = os.path.join(dir,file)
        img_name = str(file.split('.')[0]) + '.jpg'
        img_path = os.path.join(img_dir,img_name)
        if os.path.exists(img_path) and os.path.exists(xml_path):
            # print(xml_path)
            mask = xml2mask(img_path, xml_path, id2class_dict)
            cv2.imwrite(out_dir+ '\\' + os.path.splitext(file)[0] + ".png", mask)
            i+=1
            print('已完成{0}幅图像!'.format(i))
    print("全部完成!")

参考文献:
labelme xml 标注文件输出掩码图片
python实现xml标注文件生成mask
Zhang, T.; Zhang, X.; Li, J.; Xu, X.; Wang, B.; Zhan, X.; Xu, Y.; Ke, X.; Zeng, T.; Su, H.; Ahmad, I.; Pan, D.; Liu, C.; Zhou, Y.; Shi, J.; Wei, S. SAR Ship Detection Dataset (SSDD): Official Release and Comprehensive Data Analysis. Remote Sens. 2021, 13, 3690. https://doi.org/10.3390/rs13183690