图像分割-xml/json标注文件处理

1、xml文件

import cv2
import os
import argparse
from PIL import Image
import numpy as np
import xml.etree.ElementTree as ET
import matplotlib.pyplot as plt
import exifread
from PIL.ExifTags import TAGS

def apply_exif_orientation(fname):
    ret = {}
    try:
        image = Image.open(fname)
        if hasattr(image, '_getexif' ):
            exifinfo = image._getexif()
            if exifinfo != None:
                for tag, value in exifinfo.items():
                    decoded = TAGS.get(tag, tag)
                    ret[decoded] = value
    except IOError:
        print ('IOERROR ' + fname)

    orientation = ret.get('Orientation', None)
    print(ret['Orientation'], orientation)
    if orientation == 1:
         # do nothing
         return image
    elif orientation == 2:
        # left-to-right mirror
        return PIL.ImageOps.mirror(image)
    elif orientation == 3:
        # rotate 180
        return image.transpose(PIL.Image.ROTATE_180)
    elif orientation == 4:
        # top-to-bottom mirror
        return PIL.ImageOps.flip(image)
    elif orientation == 5:
        # top-to-left mirror
        return PIL.ImageOps.mirror(image.transpose(PIL.Image.ROTATE_270))
    elif orientation == 6:
        # rotate 270
        return image.transpose(PIL.Image.ROTATE_270)
    elif orientation == 7:
        # top-to-right mirror
        return PIL.ImageOps.mirror(image.transpose(PIL.Image.ROTATE_90))
    elif orientation == 8:
        # rotate 90
        return image.transpose(PIL.Image.ROTATE_90)
    else:
        return image

def fill_ploygan(image_path, xml_path, output_path):
    #img = apply_exif_orientation(image_path)
    img = cv2.imread(image_path)
    img_ = np.zeros(img.shape[:2], np.uint8)

    tree = ET.parse(xml_path)
    root = tree.getroot()

    triangle = {}
    for member in root.findall('object'):
        name = member.find('name').text

        x_list = []
        y_list = []
        for x in member.findall('./polygon/pt/x'):
            x_list.append(int(float(x.text)))
        for y in member.findall('./polygon/pt/y'):
            y_list.append(int(float(y.text)))

        triangle[name] = np.array(list(zip(x_list, y_list)))
    class_keys = triangle.keys()
    color_keys = {
        "label1": 1,
        "label2": 2,
    }

    for key in color_keys.keys():
        if key in class_keys:
            cv2.fillConvexPoly(img_, triangle[key], color_keys[key])
    cv2.imwrite(output_path.replace('jpg', 'png'), img_)
    #plt.imshow(img_)
    #plt.show()

def parsers():
    parser = argparse.ArgumentParser(description='xml2mask, a tool to get mask for image segamentation.')
    parser.add_argument('--imagedir', help='dir of image files')
    parser.add_argument('--xmldir', help='dir of xml files')
    parser.add_argument('--maskdir', help='dir of mask files')
    args = parser.parse_args()

    if not args.imagedir:
        print('You must supply a imagedir\n')
        parser.print_help()
        exit(0)

    if not args.xmldir:
        print('You must supply a xmldir\n')
        parser.print_help()
        exit(0)

    if not args.maskdir:
        print('You must supply a maskdir\n')
        parser.print_help()
        exit(0)
    return args

if __name__ == "__main__":
    args = parsers()
    for image_file in os.listdir(args.imagedir):
        image = os.path.join(args.imagedir, image_file)
        xml_file = image_file.replace('.jpg', '.xml')
        xml = os.path.join(args.xmldir, xml_file)
        mask_file = os.path.join(args.maskdir, image_file)
        fill_ploygan(image_path=image, xml_path=xml, output_path=mask_file)

2、json文件

import cv2
import os
import argparse
from PIL import Image
import numpy as np
import json
import matplotlib.pyplot as plt

def fill_ploygan(image_path, json_path, output_path):

    img = cv2.imread(image_path)
    mask = np.zeros(img.shape[:2], np.uint8)
    color_keys = {
        "round_food": 1,
        "square_food": 1
    }
    if not os.path.exists(json_path):
        pass
    else:
        with open(json_path, 'r') as load_f:
            load_dict = json.load(load_f)
            shapes = load_dict['shapes']
            for member in shapes:
                label = member['label']
                points = member['points']
                rectangular = np.array(points).astype(int)
                cv2.fillConvexPoly(mask, rectangular, color_keys[label])
    cv2.imwrite(output_path.replace('jpg', 'png'), mask)


def parsers():
    parser = argparse.ArgumentParser(description='json2mask, a tool to get mask for image segamentation.')
    parser.add_argument('--imagedir', help='dir of image files')
    parser.add_argument('--jsondir', help='dir of xml files')
    parser.add_argument('--maskdir', help='dir of mask files')
    args = parser.parse_args()

    if not args.imagedir:
        print('You must supply a imagedir\n')
        parser.print_help()
        exit(0)

    if not args.jsondir:
        print('You must supply a jsondir\n')
        parser.print_help()
        exit(0)

    if not args.maskdir:
        print('You must supply a maskdir\n')
        parser.print_help()
        exit(0)
    return args

if __name__ == "__main__":
    args = parsers()
    for image_file in os.listdir(args.imagedir):
        image_path = os.path.join(args.imagedir, image_file)
        json_file = image_file.replace('.jpg', '.json')
        json_path = os.path.join(args.jsondir, json_file)
        mask_path = os.path.join(args.maskdir, image_file)
        fill_ploygan(image_path=image_path, json_path=json_path, output_path=mask_path)