converct_VOC_to_mask

converct_VOC_to_mask.py

# read VOC style segmentation labels (.xml) and convert them to gray masked image(.jpg)
# reference: http://www.programmersought.com/article/9208990502/


import os
import cv2
import xml.dom.minidom
 
image_path = "/home/seung/Dropbox/Project/IKEA-Assembly-Instuctions-Dataset/supervisely/chair/img/"
annotation_path = "/home/seung/Annotations_step/"
save_path = "/home/seung/mask/" 

file_names = os.listdir(image_path)
color_list = [(80, 80, 80), (130, 130, 130), (180, 180, 180)]

for filename in file_names:
    filename, extension = os.path.splitext(filename)
    img_path = image_path+filename+'.jpg'
    xml_path = annotation_path+filename+'.xml'
    try:
        open(xml_path)
    except FileNotFoundError as e:
        continue

    img = cv2.imread(img_path)
    if img is None:
        pass

    mask_img = cv2.rectangle(img, (int(0), int(0)), (int(img.shape[1]), int(img.shape[0])), (0, 0, 0), -1)

    dom = xml.dom.minidom.parse(xml_path)
    root = dom.documentElement
    objects = dom.getElementsByTagName("object")
    color = (0, 0, 0)
    for i, object_ in enumerate(objects):
        bndbox = root.getElementsByTagName('bndbox')[i]
        name = object_.getElementsByTagName('name')[0]
        name = str(name.childNodes[0].data)
        if name == 'step':
            color = color_list[0]
        else:
            print("wrong name!")
            exit()

        xmin = bndbox.getElementsByTagName('xmin')[0]
        ymin = bndbox.getElementsByTagName('ymin')[0]
        xmax = bndbox.getElementsByTagName('xmax')[0]
        ymax = bndbox.getElementsByTagName('ymax')[0]
        xmin_data = xmin.childNodes[0].data
        ymin_data = ymin.childNodes[0].data
        xmax_data = xmax.childNodes[0].data
        ymax_data = ymax.childNodes[0].data
        mask_img = cv2.rectangle(img, (int(xmin_data), int(ymin_data)), (int(xmax_data), int(ymax_data)), color, -1)

    flag = 0
    flag = cv2.imwrite(save_path + "{}.png".format(filename), mask_img)
    if (flag):
        print(filename, "done")