mirror of
https://github.com/m5stack/UNITV_TRAIN.git
synced 2026-05-20 11:05:22 -07:00
optimize xml parser
This commit is contained in:
Neucrack
+49
-51
@@ -648,63 +648,61 @@ class Detector(Train_Base):
|
||||
# decode xml
|
||||
input_shape_checked = False
|
||||
for xml_path in xmls:
|
||||
with open(xml_path) as f:
|
||||
xml = f.read()
|
||||
ok, result = decode_pascal_voc_xml(xml)
|
||||
if not ok:
|
||||
result = f"decode xml {xml_path} fail, reason: {result}"
|
||||
self.on_warning_message(result)
|
||||
ok, result = decode_pascal_voc_xml(xml_path)
|
||||
if not ok:
|
||||
result = f"decode xml {xml_path} fail, reason: {result}"
|
||||
self.on_warning_message(result)
|
||||
continue
|
||||
# shape
|
||||
img_shape = (result['height'], result['width'], result['depth'])
|
||||
# check first image shape, and switch to proper supported input_shape
|
||||
if not input_shape_checked:
|
||||
if not self._check_update_input_shape(img_shape) and not self.allow_reshape:
|
||||
return False, "not supported input size, supported: {}".format(self.support_shapes), [], None, None, None
|
||||
input_shape_checked = True
|
||||
if img_shape != self.input_shape:
|
||||
msg = f"decode xml {xml_path} ok, but shape {img_shape} not the same as expected: {self.input_shape}"
|
||||
if not self.allow_reshape:
|
||||
self.on_warning_message(msg)
|
||||
continue
|
||||
# shape
|
||||
img_shape = (result['height'], result['width'], result['depth'])
|
||||
# check first image shape, and switch to proper supported input_shape
|
||||
if not input_shape_checked:
|
||||
if not self._check_update_input_shape(img_shape) and not self.allow_reshape:
|
||||
return False, "not supported input size, supported: {}".format(self.support_shapes), [], None, None, None
|
||||
input_shape_checked = True
|
||||
if img_shape != self.input_shape:
|
||||
msg = f"decode xml {xml_path} ok, but shape {img_shape} not the same as expected: {self.input_shape}"
|
||||
if not self.allow_reshape:
|
||||
self.on_warning_message(msg)
|
||||
continue
|
||||
else:
|
||||
msg += ", will automatically reshape"
|
||||
self.on_warning_message(msg)
|
||||
# load image
|
||||
dir_name = os.path.split(os.path.split(result['path'])[0])[-1] # class1 / images
|
||||
# images/class1/tututututut.jpg
|
||||
img_path = os.path.join(img_dir, dir_name, result['filename'])
|
||||
else:
|
||||
msg += ", will automatically reshape"
|
||||
self.on_warning_message(msg)
|
||||
# load image
|
||||
dir_name = os.path.split(os.path.split(result['path'])[0])[-1] # class1 / images
|
||||
# images/class1/tututututut.jpg
|
||||
img_path = os.path.join(img_dir, dir_name, result['filename'])
|
||||
if os.path.exists(img_path):
|
||||
img = np.array(Image.open(img_path), dtype='uint8')
|
||||
else:
|
||||
# images/tututututut.jpg
|
||||
img_path = os.path.join(img_dir, result['filename'])
|
||||
if os.path.exists(img_path):
|
||||
img = np.array(Image.open(img_path), dtype='uint8')
|
||||
else:
|
||||
# images/tututututut.jpg
|
||||
img_path = os.path.join(img_dir, result['filename'])
|
||||
if os.path.exists(img_path):
|
||||
img = np.array(Image.open(img_path), dtype='uint8')
|
||||
else:
|
||||
result = f"decode xml {xml_path}, can not find iamge: {result['path']}"
|
||||
self.on_warning_message(result)
|
||||
continue
|
||||
# load bndboxes
|
||||
y = []
|
||||
for bbox in result['bboxes']:
|
||||
if not bbox[4] in labels:
|
||||
result = f"decode xml {xml_path}, can not find iamge: {result['path']}"
|
||||
self.on_warning_message(result)
|
||||
continue
|
||||
label_idx = labels.index(bbox[4])
|
||||
bbox[4] = label_idx # replace label text with label index
|
||||
classes_data_counts[label_idx] += 1
|
||||
# range to [0, 1]
|
||||
y.append( bbox[:5])
|
||||
if len(y) < 1:
|
||||
result = f"decode xml {xml_path}, no object, skip"
|
||||
result = f"decode xml {xml_path}, can not find iamge: {result['path']}"
|
||||
self.on_warning_message(result)
|
||||
continue
|
||||
if img_shape != self.input_shape:
|
||||
img, y = self._reshape_image(img, self.input_shape, y)
|
||||
datasets_x.append(img)
|
||||
datasets_y.append(y)
|
||||
# load bndboxes
|
||||
y = []
|
||||
for bbox in result['bboxes']:
|
||||
if not bbox[4] in labels:
|
||||
result = f"decode xml {xml_path}, can not find iamge: {result['path']}"
|
||||
self.on_warning_message(result)
|
||||
continue
|
||||
label_idx = labels.index(bbox[4])
|
||||
bbox[4] = label_idx # replace label text with label index
|
||||
classes_data_counts[label_idx] += 1
|
||||
# range to [0, 1]
|
||||
y.append( bbox[:5])
|
||||
if len(y) < 1:
|
||||
result = f"decode xml {xml_path}, no object, skip"
|
||||
self.on_warning_message(result)
|
||||
continue
|
||||
if img_shape != self.input_shape:
|
||||
img, y = self._reshape_image(img, self.input_shape, y)
|
||||
datasets_x.append(img)
|
||||
datasets_y.append(y)
|
||||
return True, "ok", labels, classes_data_counts, datasets_x, datasets_y
|
||||
|
||||
def _decode_pbtxt_file(self, file_path):
|
||||
|
||||
@@ -10,8 +10,10 @@
|
||||
|
||||
import re
|
||||
|
||||
def decode_pascal_voc_xml(xml):
|
||||
|
||||
def decode_pascal_voc_xml(xml_path, ordered = True):
|
||||
'''
|
||||
@ordered parse labelimg ordered xml by RE, or will use xml parser
|
||||
@reuturn bool, info
|
||||
res = {
|
||||
"filename": ,
|
||||
@@ -22,27 +24,69 @@ def decode_pascal_voc_xml(xml):
|
||||
"bboxes": [(xmin, ymin, xmax, ymax, label, difficult)]
|
||||
}
|
||||
'''
|
||||
if ordered:
|
||||
with open(xml_path) as f:
|
||||
xml = f.read()
|
||||
try:
|
||||
rule = "<filename>(.*)</filename>.*<path>(.*)</path>.*<size>.*<width>(.*)</width>.*<height>(.*)</height>.*<depth>(.*)</depth>.*</size>"
|
||||
match = re.findall(rule, xml, re.MULTILINE|re.DOTALL)
|
||||
if len(match) < 1:
|
||||
return False, "decode error"
|
||||
res = {
|
||||
"filename": match[0][0].replace("\\", "/"),
|
||||
"path": match[0][1].replace("\\", "/"),
|
||||
"width": int(match[0][2]),
|
||||
"height": int(match[0][3]),
|
||||
"depth": int(match[0][4]),
|
||||
"bboxes": []
|
||||
}
|
||||
rule = "<object>.*?<name>(.*?)</name>.*?<difficult>(.*?)</difficult>.*?<bndbox>.*?<xmin>(.*?)</xmin>.*?<ymin>(.*?)</ymin>.*?<xmax>(.*?)</xmax>.*?<ymax>(.*?)</ymax>.*?</bndbox>.*?</object>"
|
||||
match = re.findall(rule, xml, re.MULTILINE|re.DOTALL)
|
||||
if len(match) < 1:
|
||||
return False, "no object in this iamge"
|
||||
for bbox in match:
|
||||
bbox = [int(bbox[2]), int(bbox[3]), int(bbox[4]), int(bbox[5]), bbox[0], int(bbox[1])]
|
||||
res["bboxes"].append(bbox)
|
||||
except Exception as e:
|
||||
return False, "decode error： {}".format(e)
|
||||
return True, res
|
||||
try:
|
||||
rule = "<filename>(.*)</filename>.*<path>(.*)</path>.*<size>.*<width>(.*)</width>.*<height>(.*)</height>.*<depth>(.*)</depth>.*</size>"
|
||||
match = re.findall(rule, xml, re.MULTILINE|re.DOTALL)
|
||||
if len(match) < 1:
|
||||
return False, "decode error"
|
||||
from xml.etree.ElementTree import parse
|
||||
tree = parse(xml_path)
|
||||
root = tree.getroot()
|
||||
filename = root.find("filename").text
|
||||
path = root.find("path")
|
||||
width = -1
|
||||
height = -1
|
||||
depth = -1
|
||||
for elem in tree.iter():
|
||||
if "width" in elem.tag:
|
||||
width = int(elem.text)
|
||||
elif "height" in elem.tag:
|
||||
height = int(elem.text)
|
||||
elif "depth" in elem.tag:
|
||||
depth = int(elem.text)
|
||||
obj_tags = root.findall("object")
|
||||
res = {
|
||||
"filename": match[0][0],
|
||||
"path": match[0][1],
|
||||
"width": int(match[0][2]),
|
||||
"height": int(match[0][3]),
|
||||
"depth": int(match[0][4]),
|
||||
"filename": filename,
|
||||
"path": filename if path is None else path.text,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"depth": depth,
|
||||
"bboxes": []
|
||||
}
|
||||
rule = "<object>.*?<name>(.*?)</name>.*?<difficult>(.*?)</difficult>.*?<bndbox>.*?<xmin>(.*?)</xmin>.*?<ymin>(.*?)</ymin>.*?<xmax>(.*?)</xmax>.*?<ymax>(.*?)</ymax>.*?</bndbox>.*?</object>"
|
||||
match = re.findall(rule, xml, re.MULTILINE|re.DOTALL)
|
||||
if len(match) < 1:
|
||||
return False, "no object in this iamge"
|
||||
for bbox in match:
|
||||
bbox = [int(bbox[2]), int(bbox[3]), int(bbox[4]), int(bbox[5]), bbox[0], int(bbox[1])]
|
||||
res["filename"] = res["filename"].replace("\\", "/"),
|
||||
res["path"] = res["path"].replace("\\", "/"),
|
||||
for t in obj_tags:
|
||||
name = t.find("name").text
|
||||
box_tag = t.find("bndbox")
|
||||
difficult = int(t.find("difficult").text)
|
||||
x1 = int(float(box_tag.find("xmin").text))
|
||||
y1 = int(float(box_tag.find("ymin").text))
|
||||
x2 = int(float(box_tag.find("xmax").text))
|
||||
y2 = int(float(box_tag.find("ymax").text))
|
||||
bbox = [x1, y1, x2, y2, name, difficult ]
|
||||
res["bboxes"].append(bbox)
|
||||
except Exception:
|
||||
return False, "decode error"
|
||||
return True, res
|
||||
|
||||
return True, res
|
||||
except Exception as e:
|
||||
return False, "decode error： {}".format(e)
|
||||
|
||||
Reference in New Issue
Block a user