Merge pull request #1 from allenai/easier-sampling

Large scale refactor to make things more modular and generally easier to modify the generation process.

procthor/constants.py

@@ -15,14 +15,6 @@
 SCHEMA = "0.0.1"
 """The schema version of the json file to create the house."""
 
-PRIORITY_ASSET_TYPES = {
-    "Bedroom": ["Bed", "Dresser"],
-    "LivingRoom": ["Television", "DiningTable", "Sofa"],
-    "Kitchen": ["CounterTop", "Fridge"],
-    "Bathroom": ["Toilet", "Sink"],
-}
-"""These objects should be placed first inside of the rooms."""
-
 MARGIN = {
     "middle": 0.35,
     "edge": {"front": 0.5, "back": 0, "sides": 0},

procthor/databases/__init__.py

@@ -1,23 +1,47 @@
 import json
 import os
+from collections import defaultdict
 from functools import lru_cache
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Union, Tuple
 
 import numpy as np
 import pandas as pd
 from ai2thor.controller import Controller
+from attr import define
+
 from procthor.constants import USE_ITHOR_SPLITS
 from procthor.utils.types import Split
 
 
+@define(eq=False)
+class ProcTHORDatabase:
+    SOLID_WALL_COLORS: List[Dict[str, float]]
+    MATERIAL_DATABASE: Dict[str, List[str]]
+    SKYBOXES: Dict[str, Dict[str, str]]
+    OBJECTS_IN_RECEPTACLES: Dict[str, Dict[str, Dict[str, Union[float, int]]]]
+    ASSET_DATABASE: Dict[str, List[Dict[str, Any]]]
+    ASSET_ID_DATABASE: Dict[str, Any]
+    PLACEMENT_ANNOTATIONS: pd.DataFrame
+    AI2THOR_OBJECT_METADATA: Dict[str, List[List[Dict[str, Any]]]]
+    ASSET_GROUPS: Dict[str, Any]
+    ASSETS_DF: pd.DataFrame
+    WALL_HOLES: Dict[str, Dict[str, Dict[str, float]]]
+    FLOOR_ASSET_DICT: Dict[Tuple[str, str], Tuple[Dict[str, Any], pd.DataFrame]]
+
+    PRIORITY_ASSET_TYPES: Dict[
+        str, List[str]
+    ]  # TODO: Move to sampling parameters in some way?
+    """These objects should be placed first inside of the rooms."""
+
+
 def _load_json_from_database(json_file: str) -> Union[list, dict]:
     dirname = os.path.dirname(__file__)
     filepath = os.path.join(dirname, json_file)
     with open(filepath, "r") as f:
         return json.load(f)
 
 
-def _get_solid_wall_colors() -> dict:
+def _get_solid_wall_colors() -> List[Dict[str, float]]:
     return _load_json_from_database("solid-wall-colors.json")
 
 
@@ -113,7 +137,7 @@ def _get_wall_holes() -> Dict[str, Any]:
 
 @lru_cache(maxsize=None)
 def get_spawnable_asset_group_info(
-    split: Split, controller: Controller
+    split: Split, controller: Controller, pt_db: ProcTHORDatabase
 ) -> pd.DataFrame:
     from procthor.generation.asset_groups import AssetGroupGenerator
 
@@ -127,6 +151,7 @@ def get_spawnable_asset_group_info(
             split=split,
             data=asset_group_data,
             controller=controller,
+            pt_db=pt_db,
         )
 
         dims = asset_group_generator.dimensions
@@ -167,14 +192,63 @@ def get_spawnable_asset_group_info(
     return pd.DataFrame(data)
 
 
-solid_wall_colors = _get_solid_wall_colors()
-material_database = _get_material_database()
-skyboxes = _get_skyboxes()
-objects_in_receptacles = _get_object_in_receptacles()
-asset_database = _get_asset_database()
-asset_id_database = _get_asset_id_database()
-placement_annotations = _get_placement_annotations()
-ai2thor_object_metadata = _get_ai2thor_object_metadata()
-asset_groups = _get_asset_groups()
-assets_df = _get_assets_df()
-wall_holes = _get_wall_holes()
+def _get_floor_assets(
+    room_type: str, split: str, pt_db: ProcTHORDatabase
+) -> Tuple[Any, pd.DataFrame]:
+    floor_types = pt_db.PLACEMENT_ANNOTATIONS[
+        pt_db.PLACEMENT_ANNOTATIONS["onFloor"]
+        & (pt_db.PLACEMENT_ANNOTATIONS[f"in{room_type}s"] > 0)
+    ]
+    assets = pd.DataFrame(
+        [
+            {
+                "assetId": asset["assetId"],
+                "assetType": asset["objectType"],
+                "split": asset["split"],
+                "xSize": asset["boundingBox"]["x"],
+                "ySize": asset["boundingBox"]["y"],
+                "zSize": asset["boundingBox"]["z"],
+            }
+            for asset_type in floor_types.index
+            for asset in pt_db.ASSET_DATABASE[asset_type]
+        ]
+    )
+    assets = pd.merge(assets, floor_types, on="assetType", how="left")
+    assets = assets[assets["split"].isin([split, None])]
+    assets.set_index("assetId", inplace=True)
+    return floor_types, assets
+
+
+def _get_default_floor_assets_from_key(key: Tuple[str, str]):
+    return _get_floor_assets(*key, pt_db=DEFAULT_PROCTHOR_DATABASE)
+
+
+class keydefaultdict(defaultdict):
+    def __missing__(self, key):
+        if self.default_factory is None:
+            raise KeyError(key)
+        else:
+            ret = self[key] = self.default_factory(key)
+            return ret
+
+
+DEFAULT_PROCTHOR_DATABASE = ProcTHORDatabase(
+    SOLID_WALL_COLORS=_get_solid_wall_colors(),
+    MATERIAL_DATABASE=_get_material_database(),
+    SKYBOXES=_get_skyboxes(),
+    OBJECTS_IN_RECEPTACLES=_get_object_in_receptacles(),
+    ASSET_DATABASE=_get_asset_database(),
+    ASSET_ID_DATABASE=_get_asset_id_database(),
+    PLACEMENT_ANNOTATIONS=_get_placement_annotations(),
+    AI2THOR_OBJECT_METADATA=_get_ai2thor_object_metadata(),
+    ASSET_GROUPS=_get_asset_groups(),
+    ASSETS_DF=_get_assets_df(),
+    WALL_HOLES=_get_wall_holes(),
+    FLOOR_ASSET_DICT=keydefaultdict(_get_default_floor_assets_from_key),
+    PRIORITY_ASSET_TYPES={
+        "Bedroom": ["Bed", "Dresser"],
+        "LivingRoom": ["Television", "DiningTable", "Sofa"],
+        "Kitchen": ["CounterTop", "Fridge"],
+        "Bathroom": ["Toilet", "Sink"],
+    },
+)