consummating xml_style.py (#2457)

* consummating xml_style.py * Update xml_style.py update xml_style.py * Update xml_style.py fix the format
open-mmlab · Apr 26, 2020 · c40cafb · c40cafb
1 parent cd4e716
commit c40cafb
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/mmdet/datasets/xml_style.py b/mmdet/datasets/xml_style.py
@@ -3,6 +3,7 @@
 
 import mmcv
 import numpy as np
+from PIL import Image
 
 from .custom import CustomDataset
 from .registry import DATASETS
@@ -26,8 +27,16 @@ def load_annotations(self, ann_file):
             tree = ET.parse(xml_path)
             root = tree.getroot()
             size = root.find('size')
-            width = int(size.find('width').text)
-            height = int(size.find('height').text)
+            width = 0
+            height = 0
+            if size is not None:
+                width = int(size.find('width').text)
+                height = int(size.find('height').text)
+            else:
+                img_path = osp.join(self.img_prefix, 'JPEGImages',
+                                    '{}.jpg'.format(img_id))
+                img = Image.open(img_path)
+                width, height = img.size
             img_infos.append(
                 dict(id=img_id, filename=filename, width=width, height=height))
         return img_infos