import glob
import os
import random
import warnings
from collections import OrderedDict
from typing import Union, List
import bpy
from blenderproc.python.types.MeshObjectUtility import MeshObject
from blenderproc.python.loader.ObjectLoader import load_obj
[docs]def load_ikea(data_dir: str = 'resources/IKEA', obj_categories: Union[list, str] = None, obj_style: str = None) -> List[MeshObject]:
""" Loads ikea objects based on selected type and style.
If there are multiple options it picks one randomly or if style or type is None it picks one randomly.
:param data_dir: The directory with all the IKEA models.
:param obj_categories: The category to use for example: 'bookcase'. This can also be a list of elements. Available: ['bed', 'bookcase', 'chair', 'desk', 'sofa', 'table', 'wardrobe']
:param obj_style: The IKEA style to use for example: 'hemnes'. See data_dir for other options.
:return: The list of loaded mesh objects.
"""
obj_dict = IKEALoader._generate_object_dict(data_dir)
# If obj_categories is given, make sure it is a list
if obj_categories is not None and not isinstance(obj_categories, list):
obj_categories = [obj_categories]
if obj_categories is not None and obj_style is not None:
object_lst = []
for obj_category in obj_categories:
object_lst.extend([obj[0] for (key, obj) in obj_dict.items() \
if obj_style in key.lower() and obj_category in key])
if not object_lst:
selected_obj = random.choice(obj_dict.get(random.choice(list(obj_dict.keys()))))
warnings.warn("Could not find object of type: {}, and style: {}. Selecting random object...".format(
obj_categories, obj_style), category=Warning)
else:
# Multiple objects with same type and style are possible: select randomly from list.
selected_obj = random.choice(object_lst)
elif obj_categories is not None:
object_lst = []
for obj_category in obj_categories:
object_lst.extend(IKEALoader._get_object_by_type(obj_category, obj_dict))
selected_obj = random.choice(object_lst)
elif obj_style is not None:
object_lst = IKEALoader._get_object_by_style(obj_style, obj_dict)
selected_obj = random.choice(object_lst)
else:
random_key = random.choice(list(obj_dict.keys()))
# One key can have multiple object files as value: select randomly from list.
selected_obj = random.choice(obj_dict.get(random_key))
print("Selected object: ", os.path.basename(selected_obj))
loaded_obj = load_obj(selected_obj)
# extract the name from the path:
selected_dir_name = os.path.dirname(selected_obj)
selected_name = ""
if os.path.basename(selected_dir_name).startswith("IKEA_"):
selected_name = os.path.basename(selected_dir_name)
else:
selected_dir_name = os.path.dirname(selected_dir_name)
if os.path.basename(selected_dir_name).startswith("IKEA_"):
selected_name = os.path.basename(selected_dir_name)
if selected_name:
for obj in loaded_obj:
obj.set_name(selected_name)
# extract the file unit from the .obj file to convert every object to meters
file_unit = ""
with open(selected_obj, "r") as file:
first_lines = [next(file) for x in range(5)]
for line in first_lines:
if "File units" in line:
file_unit = line.strip().split(" ")[-1]
if file_unit not in ["inches", "meters", "centimeters", "millimeters"]:
raise Exception("The file unit type could not be found, check the selected "
"file: {}".format(selected_obj))
break
for obj in loaded_obj:
# convert all objects to meters
if file_unit == "inches":
scale = 0.0254
elif file_unit == "centimeters":
scale = 0.01
elif file_unit == "millimeters":
scale = 0.001
elif file_unit == "meters":
scale = 1.0
else:
raise Exception("The file unit type: {} is not defined".format(file_unit))
if scale != 1.0:
# move all object centers to the world origin and set the bounding box correctly
bpy.ops.object.select_all(action='DESELECT')
obj.select()
bpy.context.view_layer.objects.active = obj.blender_obj
# scale object down
bpy.ops.object.mode_set(mode='EDIT')
bpy.ops.transform.resize(value=(scale, scale, scale))
bpy.ops.object.mode_set(mode='OBJECT')
bpy.ops.object.select_all(action='DESELECT')
# removes the x axis rotation found in all ShapeNet objects, this is caused by importing .obj files
# the object has the same pose as before, just that the rotation_euler is now [0, 0, 0]
for obj in loaded_obj:
obj.persist_transformation_into_mesh(location=False, rotation=True, scale=False)
# move the origin of the object to the world origin and on top of the X-Y plane
# makes it easier to place them later on, this does not change the `.location`
for obj in loaded_obj:
obj.move_origin_to_bottom_mean_point()
bpy.ops.object.select_all(action='DESELECT')
return loaded_obj
[docs]class IKEALoader:
"""
This class loads objects from the IKEA dataset.
Objects can be selected randomly, based on object type, object style, or both.
"""
[docs] @staticmethod
def _generate_object_dict(data_dir: str) -> dict:
"""
Generates a dictionary of all available objects, i.e. all .obj files that have an associated .mtl file.
:param data_dir: The directory with all the IKEA models.
:return: dict: {IKEA_<type>_<style> : [<path_to_obj_file>, ...]}
"""
obj_dict = {}
counter = 0
obj_files = glob.glob(os.path.join(data_dir, "IKEA", "*", "*.obj"))
# this will be ensure that the call is deterministic
obj_files.sort()
for obj_file in obj_files:
category = [s for s in obj_file.split('/') if 'IKEA_' in s][0]
if IKEALoader._check_material_file(obj_file):
obj_dict.setdefault(category, []).append(obj_file)
counter += 1
print('Found {} object files in dataset belonging to {} categories'.format(counter, len(obj_dict)))
if len(obj_dict) == 0:
raise Exception("No obj file was found, check if the correct folder is provided!")
# to avoid randomness while accessing the dict
obj_dict = OrderedDict(obj_dict)
return obj_dict
[docs] @staticmethod
def _check_material_file(path: str) -> bool:
"""
Checks whether there is a texture file (.mtl) associated to the object available.
:param path: path to object
:return: texture file exists
"""
name = os.path.basename(path).split(".")[0]
obj_dir = os.path.dirname(path)
mtl_path = os.path.join(obj_dir, name + ".mtl")
return os.path.exists(mtl_path)
[docs] @staticmethod
def _get_object_by_type(obj_type: str, obj_dict: dict) -> list:
"""
Finds all available objects with a specific type.
:param obj_type: type of object e.g. 'table'
:return: list of available objects with specified type
"""
object_lst = [obj[0] for (key, obj) in obj_dict.items() if obj_type in key]
if not object_lst:
warnings.warn("There were no objects found matching the type: {}.".format(obj_type), category=Warning)
return object_lst
[docs] @staticmethod
def _get_object_by_style(obj_style, obj_dict):
"""
Finds all available objects with a specific style, i.e. IKEA product series.
:param obj_type: (str) type of object e.g. 'table'
:return: (list) list of available objects with specified style
"""
object_lst = [obj[0] for (key, obj) in obj_dict.items() if obj_style in key.lower()]
if not object_lst:
warnings.warn("There were no objects found matching the style: {}.".format(obj_style), category=Warning)
return object_lst