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)
