from imageio import imread
from keras.preprocessing import image
from matplotlib import pyplot as plt
import xml.dom.minidom
import numpy as np
#Train.txt格式类似如下
#data/123.jpg
lines = []
with open("./dataset/Train.txt",‘r’,encoding=“utf-8”) as f:
lines = [line[:-1] for line in f.readlines()]
getObject=[]
def paraXml(path):
DomTree= xml.dom.minidom.parse(path)
collection =DomTree.documentElement
fileName= collection.getElementsByTagName(‘filename’)[0].childNodes[0].data
size = collection.getElementsByTagName(‘size’)[0]
width= size.getElementsByTagName(‘width’)[0].childNodes[0].data
height= size.getElementsByTagName(‘height’)[0].childNodes[0].data
objects = collection.getElementsByTagName(‘object’)
for obj in objects:
# 获取类别
name = obj.getElementsByTagName(“name”)[0].childNodes[0].data
# 获取坐标
bndbox = obj.getElementsByTagName(“bndbox”)[0]
xmin = bndbox.getElementsByTagName(“xmin”)[0].childNodes[0].data
ymin = bndbox.getElementsByTagName(“ymin”)[0].childNodes[0].data
xmax = bndbox.getElementsByTagName(“xmax”)[0].childNodes[0].data
ymax = bndbox.getElementsByTagName(“ymax”)[0].childNodes[0].data
object=[]
object.append(name)
object.append(xmin)
object.append(ymin)
object.append(xmax)
object.append(ymax)
getObject.append(object)
# 构造输出对象
orig_images = [] # Store the images here.
input_images = [] # Store resized versions of the images here.
img_height = 300
img_width = 300
orig_images.append(imread(lines[0]))
orig_height = orig_images[0].shape[0]
orig_width = orig_images[0].shape[1]
img = image.load_img(lines[0], target_size=(img_height, img_width))
img = image.img_to_array(img)
input_images.append(img)
input_images = np.array(input_images)
input_images = input_images.astype(int)
colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
classes = [‘background’,‘headstock’, ‘plate’]
plt.figure(figsize=(20,12))
plt.imshow(orig_images[0])
path = lines[0].split(".jpg")[0]+".xml"
paraXml(path)
current_axis = plt.gca()
for box in getObject:
# Transform the predicted bounding boxes for the 300x300 image to the original image dimensions.
xmin = int(box[1])
ymin = int(box[2])
xmax = int(box[3])
ymax = int(box[4])
if box[0] == “headstock”:
label = ‘{}’.format(classes[1])
color = colors[1]
current_axis.add_patch(plt.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, color=color, fill=False, linewidth=2))
current_axis.text(xmin, ymin, label, size=‘x-large’, color=‘white’, bbox={‘facecolor’:color, ‘alpha’:1.0})
elif box[0] == “plate”:
label = ‘{}’.format(classes[2])
color = colors[2]
current_axis.add_patch(
plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color=color, fill=False, linewidth=2))
current_axis.text(xmin, ymin, label, size=‘x-large’, color=‘white’, bbox={‘facecolor’: color, ‘alpha’: 1.0})
plt.show()
plt.figure(figsize=(20,12))
plt.imshow(input_images[0])
path = lines[0].split(".jpg")[0]+".xml"
paraXml(path)
current_axis = plt.gca()
for box in getObject:
# Transform the predicted bounding boxes for the 300x300 image to the original image dimensions.
xmin = int(box[1])*img_width/orig_width
ymin = int(box[2])*img_height/orig_height
xmax = int(box[3])*img_width/orig_width
ymax = int(box[4])*img_height/orig_height
if box[0] == “headstock”:
label = ‘{}’.format(classes[1])
color = colors[1]
current_axis.add_patch(plt.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, color=color, fill=False, linewidth=2))
current_axis.text(xmin, ymin, label, size=‘x-large’, color=‘white’, bbox={‘facecolor’:color, ‘alpha’:1.0})
elif box[0] == “plate”:
label = ‘{}’.format(classes[2])
color = colors[2]
current_axis.add_patch(
plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color=color, fill=False, linewidth=2))
current_axis.text(xmin, ymin, label, size=‘x-large’, color=‘white’, bbox={‘facecolor’: color, ‘alpha’: 1.0})
plt.show()