1. 获取分割结果
2. 把分割结果画图片上
# --*-- coding:utf-8 -*-
import cv2
import numpy as np
import xml.etree.ElementTree as ET
def xml_reader(filename):
tree = ET.parse(filename)
objects = []
for obj in tree.findall('image'):
obj_struct = {}
obj_attrib = obj.attrib
obj_struct['name'] = obj_attrib['name']
polygons = []
for polygon in obj.findall('polygon'):
polygon_attrib = polygon.attrib
polygon_struct = {}
polygon_struct['label'] = polygon_attrib['label']
polygon_struct['points'] = polygon_attrib['points']
polygons.append(polygon_struct)
obj_struct['polygons'] = polygons
objects.append(obj_struct)
return objects
def vis_mask(img, box, col, alpha=0.7, show_border=True, border_thick=3):
"""
Visualizes a single binary mask.
col = np.array([0.2, 0.7, 0.4])*225
box = np.array([[x1, y1], [x2, y2], [x3, y3], [x4, y4]], dtype = np.int32)
"""
img = img.astype(np.float32)
mask = np.zeros((img.shape[0], img.shape[1]), dtype = np.uint8)
cv2.fillPoly(mask,[box], 1)
idx = np.nonzero(mask)
img[idx[0], idx[1], :] *= 1.0 - alpha
img[idx[0], idx[1], :] += alpha * col
if show_border:
a = cv2.findContours(mask.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
ctr = np.array(a[1][0]).reshape(-1,1,2).astype(np.int32)
cv2.drawContours(img, [ctr], -1, (255,255,0), border_thick)
return img.astype(np.uint8)
if __name__ == "__main__":
objects = xml_reader('0.xml')
for idx, obj in enumerate(objects):
img_path = obj['name']
img= cv2.imread(img_path)
for polygon in obj['polygons']:
label = polygon['label']
if label == "person":
col = np.array([0.1, 0.1, 0.8])*225
elif label == "box":
col = np.array([0.1, 0.8, 0.1])*225
else:
col = np.array([0.8, 0.1, 0.1])*225
points = polygon['points']
points = [i.split(",") for i in points.split(";")]
points = np.array(points, dtype= np.float)
box = np.array(points, dtype= np.int32)
img = vis_mask(img, box, col)
cv2.imwrite("%s.jpg" % idx, img)
看效果:
end!