Fix project label not being set

MLSTRUCT · Aug 19, 2024 · 1434e0c · 1434e0c
1 parent eb05f74
commit 1434e0c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/MLStructFP/__init__.py b/MLStructFP/__init__.py
@@ -9,7 +9,7 @@
 __description__ = 'Machine learning structural floor plan dataset'
 __keywords__ = ['ml', 'ai', 'floor plan', 'architectural', 'dataset', 'cnn']
 __email__ = 'pablo@ppizarror.com'
-__version__ = '0.6.6'
+__version__ = '0.6.7'
 
 # URL
 __url__ = 'https://github.com/MLSTRUCT/MLSTRUCT-FP'

diff --git a/MLStructFP/db/_db_loader.py b/MLStructFP/db/_db_loader.py
@@ -60,7 +60,10 @@ def __init__(self, db: str, floor_only: bool = False) -> None:
                 item_types[ic[0]] = (cat, ic[1])
             project_label: Dict[int, str] = {}
             for pid in (meta['project_label'] if 'project_label' in meta else {}):
-                project_label[pid] = meta['project_label'][pid]
+                try:
+                    project_label[int(pid)] = meta['project_label'][pid]
+                except ValueError:
+                    pass
             room_categories: Dict[int, Tuple[str, str]] = {}
             for cat in (meta['room_categories'] if 'room_categories' in meta else {}):
                 rc = meta['room_categories'][cat]
@@ -69,12 +72,14 @@ def __init__(self, db: str, floor_only: bool = False) -> None:
             # Load floors
             for f_id in data.get('floor', {}):
                 f_data: dict = data['floor'][f_id]
-                f_cat = int(f_data['category'] if 'category' in f_data else 0)
+                f_cat: int = int(f_data['category'] if 'category' in f_data else 0)
+                project_id: int = f_data['project'] if 'project' in f_data else -1
                 self.__floor[int(f_id)] = Floor(
                     floor_id=int(f_id),
                     image_path=os.path.join(self.__path, f_data['image']),
                     image_scale=f_data['scale'],
-                    project_id=f_data['project'] if 'project' in f_data else -1,
+                    project_id=project_id,
+                    project_label=project_label[project_id] if project_id in project_label else '',
                     category=f_cat,
                     category_name=floor_categories.get(f_cat, ''),
                     elevation=f_data['elevation'] if 'elevation' in f_data else False