From 92f1a1195dab016d2dffdb035a51a48e802ddb8b Mon Sep 17 00:00:00 2001 From: Etienne David Date: Wed, 22 Jan 2020 18:51:24 +0100 Subject: [PATCH] Few mistakes in from_xml (#29) * handle images without object * Update dataset.py * Update dataset.py correct few mistakes during iterations --- imantics/dataset.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/imantics/dataset.py b/imantics/dataset.py index 106d545..fdd359f 100644 --- a/imantics/dataset.py +++ b/imantics/dataset.py @@ -28,14 +28,16 @@ def from_xml(cls, xml_folder, name="XML Dataset"): for ext in extensions: xml_list += list(xml_folder.glob(f"*.{ext}")) categories = [] - if "object" in xml["annotation"].keys(): - if type(xml["annotation"]["object"]) is not list: - cat = xml["annotation"]["object"]["name"]["$"] - categories.append(cat) - else: - for ann in xml["annotation"]["object"]: - cat = ann["name"]["$"] + for idx, imgp in enumerate(xml_list): + xml = bf.data(fromstring(open(imgp.with_suffix(".xml"),"r").read())) + if "object" in xml["annotation"].keys(): + if type(xml["annotation"]["object"]) is not list: + cat = xml["annotation"]["object"]["name"]["$"] categories.append(cat) + else: + for ann in xml["annotation"]["object"]: + cat = ann["name"]["$"] + categories.append(cat) categories = list(set(categories)) @@ -47,13 +49,13 @@ def from_xml(cls, xml_folder, name="XML Dataset"): image.dataset = name - + xml = bf.data(fromstring(open(imgp.with_suffix(".xml"),"r").read())) if "object" in xml["annotation"].keys(): # Handle single object case if type(xml["annotation"]["object"]) is not list: xml["annotation"]["object"] = [xml["annotation"]["object"]] - + for ann in xml["annotation"]["object"]: i = ann["bndbox"] cat = ann["name"]["$"] @@ -63,10 +65,9 @@ def from_xml(cls, xml_folder, name="XML Dataset"): fin_ann = Annotation(id=id_counter, image=image, bbox=bbox,category=xml_categories[cat]) id_counter += 1 - + image.add(fin_ann) - - dataset.add(image) + dataset.add(image) return dataset