Merge regression in RO training

Fixes #549
mittagessen · Dec 4, 2023 · 99dc1c7 · 99dc1c7
1 parent 5b441ad
commit 99dc1c7
Showing 1 changed file with 10 additions and 22 deletions.
diff --git a/kraken/lib/dataset/ro.py b/kraken/lib/dataset/ro.py
@@ -80,17 +80,6 @@ def __init__(self, files: Sequence[Union[PathLike, str]] = None,
         self.data = []
 
         if mode in ['alto', 'page', 'xml']:
-            for file in files:
-                try:
-                    doc = XMLPage(file, filetype=mode)
-                    for tag in doc.tags:
-                        if tag not in self.class_mapping:
-                            self.class_mapping[tag] = self.num_classes
-                            self.num_classes += 1
-                except KrakenInputException as e:
-                    files.pop(file)
-                    logger.warning(e)
-                    continue
             for file in files:
                 try:
                     doc = XMLPage(file, filetype=mode)
@@ -104,6 +93,11 @@ def __init__(self, files: Sequence[Union[PathLike, str]] = None,
                         order = doc.get_sorted_regions(ro_id)
                     else:
                         raise ValueError(f'Invalid RO type {level}')
+                    for el in order:
+                        for tag in el.tags.values():
+                            if tag not in self.class_mapping:
+                                self.class_mapping[tag] = self.num_classes
+                                self.num_classes += 1
                     # traverse RO and substitute features.
                     w, h = doc.image_size
                     sorted_lines = []
@@ -186,17 +180,6 @@ def __init__(self, files: Sequence[Union[PathLike, str]] = None,
         self.data = []
 
         if mode in ['alto', 'page', 'xml']:
-            for file in files:
-                try:
-                    doc = XMLPage(file, filetype=mode)
-                    for tag in doc.tags:
-                        if tag not in self.class_mapping:
-                            self.class_mapping[tag] = self.num_classes
-                            self.num_classes += 1
-                except KrakenInputException as e:
-                    files.pop(file)
-                    logger.warning(e)
-                    continue
             for file in files:
                 try:
                     doc = XMLPage(file, filetype=mode)
@@ -210,6 +193,11 @@ def __init__(self, files: Sequence[Union[PathLike, str]] = None,
                         order = doc.get_sorted_regions(ro_id)
                     else:
                         raise ValueError(f'Invalid RO type {level}')
+                    for el in order:
+                        for tag in el.tags.values():
+                            if tag not in self.class_mapping:
+                                self.class_mapping[tag] = self.num_classes
+                                self.num_classes += 1
                     # traverse RO and substitute features.
                     w, h = doc.image_size
                     sorted_lines = []