"""Load the SUNCG dataset, beware that the SUNCG dataset can not be downloaded anymore."""
import csv
import json
import math
import os
from typing import List, Optional, Dict, Any, Union, Tuple
import bpy
import numpy as np
from mathutils import Matrix, Vector
from blenderproc.python.utility.LabelIdMapping import LabelIdMapping
from blenderproc.python.types.MaterialUtility import Material
from blenderproc.python.utility.MathUtility import change_coordinate_frame_of_point
from blenderproc.python.types.EntityUtility import create_empty, Entity
from blenderproc.python.types.MeshObjectUtility import create_primitive, MeshObject
from blenderproc.python.utility.Utility import Utility, resolve_path, resolve_resource
from blenderproc.python.loader.ObjectLoader import load_obj
[docs]
def load_suncg(house_path: str, label_mapping: LabelIdMapping,
suncg_dir: Optional[str] = None) -> List[Union[Entity, MeshObject]]:
""" Loads a house.json file into blender.
- Loads all objects files specified in the house.json file.
- Orders them hierarchically (level -> room -> object)
- Writes metadata into the custom properties of each object
:param house_path: The path to the house.json file which should be loaded.
:param label_mapping: A dict which maps the names of the objects to ids.
:param suncg_dir: The path to the suncg root directory which should be used for loading objects,
rooms, textures etc.
:return: The list of loaded mesh objects.
"""
# If not suncg root directory has been given, determine it via the given house directory.
if suncg_dir is None:
suncg_dir = os.path.join(os.path.dirname(house_path), "..", "..")
_SuncgLoader.suncg_dir = suncg_dir
_SuncgLoader.collection_of_loaded_objs = {}
# there are only two types of materials, textures and diffuse
_SuncgLoader.collection_of_loaded_mats = {"texture": {}, "diffuse": {}}
with open(resolve_path(house_path), "r", encoding="utf-8") as f:
config = json.load(f)
model_category_mapping_path = resolve_resource(os.path.join('suncg', 'Better_labeling_for_NYU.csv'))
result = _SuncgLoader.read_model_category_mapping(model_category_mapping_path)
object_label_map, object_fine_grained_label_map, object_coarse_grained_label_map = result
house_id = config["id"]
loaded_objects = []
for level in config["levels"]:
# Build empty level object which acts as a parent for all rooms on the level
level_obj = create_empty("Level#" + level["id"])
level_obj.set_cp("suncg_type", "Level")
if "bbox" in level:
level_obj.set_cp("bbox", _SuncgLoader.correct_bbox_frame(level["bbox"]))
else:
print("Warning: The level with id " + level[
"id"] + " is missing the bounding box attribute in the given house.json file!")
loaded_objects.append(level_obj)
room_per_object: Dict[int, Entity] = {}
for node in level["nodes"]:
# Skip invalid nodes (This is the same behavior as in the SUNCG Toolbox)
if "valid" in node and node["valid"] == 0:
continue
# Metadata is directly stored in the objects custom data
metadata = {
"type": node["type"],
"is_suncg": True
}
if "modelId" in node:
metadata["modelId"] = node["modelId"]
if node["modelId"] in object_fine_grained_label_map:
metadata["fine_grained_class"] = object_fine_grained_label_map[node["modelId"]]
metadata["coarse_grained_class"] = object_coarse_grained_label_map[node["modelId"]]
metadata["category_id"] = label_mapping.id_from_label(object_label_map[node["modelId"]])
if "bbox" in node:
metadata["bbox"] = _SuncgLoader.correct_bbox_frame(node["bbox"])
if "transform" in node:
transform = Matrix([node["transform"][i * 4:(i + 1) * 4] for i in range(4)])
# Transpose, as given transform matrix was col-wise, but blender expects row-wise
transform.transpose()
else:
transform = None
if "materials" in node:
material_adjustments = node["materials"]
else:
material_adjustments = []
# Lookup if the object belongs to a room
object_id = int(node["id"].split("_")[-1])
parent = room_per_object.get(object_id, level_obj)
if node["type"] == "Room":
loaded_objects += _SuncgLoader.load_room(node, metadata, material_adjustments, transform, house_id,
level_obj, room_per_object, label_mapping)
elif node["type"] == "Ground":
loaded_objects += _SuncgLoader.load_ground(node, metadata, material_adjustments, transform, house_id,
parent, label_mapping)
elif node["type"] == "Object":
loaded_objects += _SuncgLoader.load_object(node, metadata, material_adjustments, transform, parent)
elif node["type"] == "Box":
loaded_objects += _SuncgLoader.load_box(node, material_adjustments, transform, parent, label_mapping)
_SuncgLoader.rename_materials()
return loaded_objects
[docs]
class _SuncgLoader:
suncg_dir: Optional[str] = None
collection_of_loaded_objs: Dict[str, List[MeshObject]] = {}
collection_of_loaded_mats: Dict[str, Dict[str, Material]] = {}
[docs]
@staticmethod
def rename_materials():
"""
Rename all materials based on their texture if they have one
This makes the accessing later on easier
"""
# TODO: should only be done to suncg materials
for material in bpy.data.materials:
if material.use_nodes:
nodes = material.node_tree.nodes
textures = Utility.get_nodes_with_type(nodes, "ShaderNodeTexImage")
if len(textures) == 1:
material.name = textures[0].image.name
[docs]
@staticmethod
def load_room(node: Dict[str, Any], metadata: Dict[str, Union[str, int]],
material_adjustments: List[Dict[str, str]], transform: Matrix, house_id: str,
parent: Entity, room_per_object: Dict[int, Entity], label_mapping: LabelIdMapping) \
-> List[Union[Entity, MeshObject]]:
""" Load the room specified in the given node.
:param node: The node dict which contains information from house.json..
:param metadata: A dict of metadata which will be written into the object's custom data.
:param material_adjustments: Adjustments to the materials which were specified inside house.json.
:param transform: The transformation that should be applied to the loaded objects.
:param house_id: The id of the current house.
:param parent: The parent object to which the room should be linked
:param room_per_object: A dict for object -> room lookup (Will be written into)
:return: The list of loaded mesh objects.
"""
# Build empty room object which acts as a parent for all objects inside
room_obj = create_empty("Room#" + node["id"])
room_obj.set_cp("suncg_type", "Room")
room_obj.set_cp("bbox", _SuncgLoader.correct_bbox_frame(node["bbox"]))
room_obj.set_cp("roomTypes", node["roomTypes"])
room_obj.set_parent(parent)
loaded_objects = [room_obj]
# Store indices of all contained objects in
if "nodeIndices" in node:
for child_id in node["nodeIndices"]:
room_per_object[child_id] = room_obj
if "hideFloor" not in node or node["hideFloor"] != 1:
metadata["type"] = "Floor"
metadata["category_id"] = label_mapping.id_from_label("floor")
metadata["fine_grained_class"] = "floor"
loaded_objects += _SuncgLoader.load_obj(
os.path.join(_SuncgLoader.suncg_dir, "room", house_id, node["modelId"] + "f.obj"), metadata,
material_adjustments, transform, room_obj)
if "hideCeiling" not in node or node["hideCeiling"] != 1:
metadata["type"] = "Ceiling"
metadata["category_id"] = label_mapping.id_from_label("ceiling")
metadata["fine_grained_class"] = "ceiling"
loaded_objects += _SuncgLoader.load_obj(
os.path.join(_SuncgLoader.suncg_dir, "room", house_id, node["modelId"] + "c.obj"), metadata,
material_adjustments, transform, room_obj)
if "hideWalls" not in node or node["hideWalls"] != 1:
metadata["type"] = "Wall"
metadata["category_id"] = label_mapping.id_from_label("wall")
metadata["fine_grained_class"] = "wall"
loaded_objects += _SuncgLoader.load_obj(
os.path.join(_SuncgLoader.suncg_dir, "room", house_id, node["modelId"] + "w.obj"), metadata,
material_adjustments, transform, room_obj)
return loaded_objects
[docs]
@staticmethod
def load_ground(node: Dict[str, Any], metadata: Dict[str, Union[str, int]],
material_adjustments: List[Dict[str, str]],
transform: Matrix, house_id: str, parent: Entity, label_mapping: LabelIdMapping) -> List[
MeshObject]:
""" Load the ground specified in the given node.
:param node: The node dict which contains information from house.json..
:param metadata: A dict of metadata which will be written into the object's custom data.
:param material_adjustments: Adjustments to the materials which were specified inside house.json.
:param transform: The transformation that should be applied to the loaded objects.
:param house_id: The id of the current house.
:param parent: The parent object to which the ground should be linked
:return: The list of loaded mesh objects.
"""
metadata["type"] = "Ground"
metadata["category_id"] = label_mapping.id_from_label("floor")
metadata["fine_grained_class"] = "ground"
return _SuncgLoader.load_obj(os.path.join(_SuncgLoader.suncg_dir, "room", house_id, node["modelId"] + "f.obj"),
metadata, material_adjustments, transform, parent)
[docs]
@staticmethod
def load_object(node: Dict[str, Any], metadata: Dict[str, Union[str, int]],
material_adjustments: List[Dict[str, str]],
transform: Matrix, parent: Entity) -> List[MeshObject]:
""" Load the object specified in the given node.
:param node: The node dict which contains information from house.json..
:param metadata: A dict of metadata which will be written into the object's custom data.
:param material_adjustments: Adjustments to the materials which were specified inside house.json.
:param transform: The transformation that should be applied to the loaded objects.
:param parent: The parent object to which the ground should be linked
:return: The list of loaded mesh objects.
"""
if "state" not in node or node["state"] == 0:
file_path = os.path.join(_SuncgLoader.suncg_dir, "object", node["modelId"], node["modelId"] + ".obj")
return _SuncgLoader.load_obj(file_path, metadata, material_adjustments, transform, parent)
file_path = os.path.join(_SuncgLoader.suncg_dir, "object", node["modelId"], node["modelId"] + "_0.obj")
return _SuncgLoader.load_obj(file_path, metadata, material_adjustments, transform, parent)
[docs]
@staticmethod
def correct_bbox_frame(bbox: Dict[str, Union[np.ndarray, list, Vector]]) -> Dict[str, np.ndarray]:
""" Corrects the coordinate frame of the given bbox.
:param bbox: The bbox.
:return: The corrected bbox.
"""
return {
"min": change_coordinate_frame_of_point(bbox["min"], ["X", "-Z", "Y"]),
"max": change_coordinate_frame_of_point(bbox["max"], ["X", "-Z", "Y"])
}
[docs]
@staticmethod
def load_box(node: Dict[str, Any], material_adjustments: List[Dict[str, str]], transform: Matrix, parent: Entity,
label_mapping: LabelIdMapping) -> List[MeshObject]:
""" Creates a cube inside blender which follows the specifications of the given node.
:param node: The node dict which contains information from house.json..
:param material_adjustments: Adjustments to the materials which were specified inside house.json.
:param transform: The transformation that should be applied to the loaded objects.
:param parent: The parent object to which the ground should be linked
:return: The list of loaded mesh objects.
"""
box = create_primitive("CUBE")
box.set_name("Box#" + node["id"])
# Scale the cube to the required dimensions
local2world_mat = Matrix.Scale(node["dimensions"][0] / 2, 4, (1.0, 0.0, 0.0)) \
@ Matrix.Scale(node["dimensions"][1] / 2, 4, (0.0, 1.0, 0.0)) \
@ Matrix.Scale(node["dimensions"][2] / 2, 4, (0.0, 0.0, 1.0))
box.set_local2world_mat(local2world_mat)
bpy.ops.object.editmode_toggle()
bpy.ops.uv.cube_project()
bpy.ops.object.editmode_toggle()
# Create an empty material which is filled in the next step
box.new_material("material_0")
_SuncgLoader.transform_and_colorize_object(box, material_adjustments, transform, parent)
# set class to void
box.set_cp("category_id", label_mapping.id_from_label("void"))
# Rotate cube to match objects loaded from .obj, has to be done after transformations have been applied
box.set_local2world_mat(Matrix.Rotation(math.radians(90), 4, "X") @ Matrix(box.get_local2world_mat()))
return [box]
[docs]
@staticmethod
def load_obj(path: str, metadata: Dict[str, Union[str, int]], material_adjustments: List[Dict[str, str]],
transform: Optional[Matrix] = None, parent: Optional[Entity] = None) -> List[MeshObject]:
""" Load the wavefront object file from the given path and adjust according to the given arguments.
:param path: The path to the .obj file.
:param metadata: A dict of metadata which will be written into the object's custom data.
:param material_adjustments: Adjustments to the materials which were specified inside house.json.
:param transform: The transformation that should be applied to the loaded objects.
:param parent: The parent object to which the object should be linked
:return: The list of loaded mesh objects.
"""
if not os.path.exists(path):
print(f"Warning: {path} is missing!")
return []
object_already_loaded = path in _SuncgLoader.collection_of_loaded_objs
loaded_objects = load_obj(filepath=path, cached_objects=_SuncgLoader.collection_of_loaded_objs)
if object_already_loaded:
print(f"Duplicate object: {path}")
for obj in loaded_objects:
# the original object matrix from the .obj loader -> is not an identity matrix
obj.set_local2world_mat(Matrix([[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]))
# remove all custom properties
obj.clear_all_cps()
# Go through all imported objects
for obj in loaded_objects:
for key in metadata.keys():
used_key = key
if key == "type":
used_key = "suncg_type"
obj.set_cp(used_key, metadata[key])
_SuncgLoader.transform_and_colorize_object(obj, material_adjustments, transform, parent)
return loaded_objects
[docs]
@staticmethod
def get_type_and_value_from_mat(mat: Material) -> Tuple[str, str]:
"""
Returns the type of the material -> either diffuse or with texture (there are only two in SUNCG)
:param mat: the material where the type and value should be determined
:return: mat_type, value: mat_type is either "diffuse" or "texture", the value contains either name of the \
image or the color mapped to an RGB string of the values
"""
image_node = mat.get_nodes_with_type('TexImage')
if len(image_node) == 1:
# there is an image node -> type texture
mat_type = "texture"
image_node = image_node[0]
if image_node.image is None:
raise RuntimeError(f"The image does not have a texture for material: {mat.get_name()}")
value = image_node.image.name
if "." in value:
value = value[:value.find(".")]
else:
mat_type = "diffuse"
principled_node = mat.get_the_one_node_with_type("BsdfPrincipled")
used_keys = list(principled_node.inputs["Base Color"].default_value)
alpha = principled_node.inputs['Alpha'].default_value
used_keys.append(alpha)
value = "_".join([str(int(255. * ele)) for ele in used_keys])
return mat_type, value
[docs]
@staticmethod
def recreate_material_nodes(mat: Material, force_texture: bool):
""" Remove all nodes and recreate a diffuse node, optionally with texture.
This will replace all material nodes with only a diffuse and a texturing node (to speedup rendering).
:param mat: The blender material
:param force_texture: True, if a texture node should always be created even if the material
has no texture at the moment
"""
image_node = mat.get_nodes_with_type('TexImage')
# if there is no image no create one
if force_texture and len(image_node) == 0:
# The principled BSDF node contains all imported material properties
principled_node = mat.get_the_one_node_with_type("BsdfPrincipled")
uv_node = mat.new_node('ShaderNodeTexCoord')
# create an image node and link it
image_node = mat.new_node('ShaderNodeTexImage')
mat.link(uv_node.outputs['UV'], image_node.inputs['Vector'])
mat.link(image_node.outputs['Color'], principled_node.inputs['Base Color'])
[docs]
@staticmethod
def adjust_material_nodes(mat: Material, adjustments: Dict[str, str]):
""" Adjust the material node of the given material according to the given adjustments.
Textures or diffuse colors will be changed according to the given material_adjustments.
:param mat: The blender material.
:param adjustments: A dict containing a new "diffuse" color or a new "texture" path
"""
if "diffuse" in adjustments:
principle_node = mat.get_the_one_node_with_type("BsdfPrincipled")
principle_node.inputs['Base Color'].default_value = Utility.hex_to_rgba(adjustments["diffuse"])
if "texture" in adjustments:
image_path = os.path.join(_SuncgLoader.suncg_dir, "texture", adjustments["texture"])
image_path = resolve_path(image_path)
if os.path.exists(image_path + ".png"):
image_path += ".png"
else:
image_path += ".jpg"
image_node = mat.get_the_one_node_with_type("ShaderNodeTexImage")
if os.path.exists(image_path):
image_node.image = bpy.data.images.load(image_path, check_existing=True)
else:
print(f"Warning: Cannot load texture, path does not exist: {image_path}, remove image node again")
mat.remove_node(image_node)
[docs]
@staticmethod
def read_model_category_mapping(path: str):
""" Reads in the model category mapping csv.
:param path: The path to the csv file.
"""
object_label_map = {}
object_fine_grained_label_map = {}
object_coarse_grained_label_map = {}
with open(resolve_path(path), 'r', encoding="utf-8") as csvfile:
reader = csv.DictReader(csvfile)
for row in reader:
object_label_map[row["model_id"]] = row["nyuv2_40class"]
object_fine_grained_label_map[row["model_id"]] = row["fine_grained_class"]
object_coarse_grained_label_map[row["model_id"]] = row["coarse_grained_class"]
return object_label_map, object_fine_grained_label_map, object_coarse_grained_label_map