Regression in pretraining

legacy polygon status flag would cause crash in pretraining without explicit validation set. Fixes #626.
mittagessen · Jul 23, 2024 · 4730c5e · 4730c5e
1 parent e82ee4d
commit 4730c5e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/kraken/lib/pretrain/model.py b/kraken/lib/pretrain/model.py
@@ -203,11 +203,11 @@ def __init__(self,
             self.train_set, self.val_set = random_split(train_set, (train_len, val_len))
 
         if format_type == 'binary':
-            legacy_train_status = train_set.legacy_polygons_status
-            if val_set and val_set.legacy_polygons_status != legacy_train_status:
+            legacy_train_status = self.train_set.dataset.legacy_polygons_status
+            if self.val_set.dataset.legacy_polygons_status != legacy_train_status:
                 logger.warning('Train and validation set have different legacy '
                                f'polygon status: {legacy_train_status} and '
-                               f'{val_set.legacy_polygons_status}. Train set '
+                               f'{self.val_set.dataset.legacy_polygons_status}. Train set '
                                'status prevails.')
             if legacy_train_status == "mixed":
                 logger.warning('Mixed legacy polygon status in training dataset. Consider recompilation.')

diff --git a/kraken/lib/train.py b/kraken/lib/train.py
@@ -394,11 +394,11 @@ def __init__(self,
                              'command. Please add valid XML, line, or binary data.')
 
         if format_type == 'binary':
-            legacy_train_status = train_set.legacy_polygons_status
-            if val_set and val_set.legacy_polygons_status != legacy_train_status:
+            legacy_train_status = self.train_set.dataset.legacy_polygons_status
+            if self.val_set.dataset.legacy_polygons_status != legacy_train_status:
                 logger.warning('Train and validation set have different legacy '
                                f'polygon status: {legacy_train_status} and '
-                               f'{val_set.legacy_polygons_status}. Train set '
+                               f'{self.val_set.dataset.legacy_polygons_status}. Train set '
                                'status prevails.')
             if legacy_train_status == "mixed":
                 logger.warning('Mixed legacy polygon status in training dataset. Consider recompilation.')