def __init__(self,
data_dir,
data_root_dir,
scene_indices=(0, 4000),
p_auxiliary=0.7,
seed=None,
ablation=None):
"""
Parameters
----------
data_root_dir (String): root dir where all data lives
data_dir (String): directory where this dataset lives (relative to data_root_dir)
scene_indices (tuple[int, int]): list of indices of scenes (in data_dir) that are considered part of this set
p_auxiliary (int): probability that a auxiliary category is chosen
Note that since (existing centroid) is sparse, it is actually treated as non-auxiliary when sampling
seed (int or None, optional): random seed, to replicate stuff, if set
ablation (string or None, optional): see data.RenderedComposite.get_composite, and paper
"""
self.category_map = ObjectCategories()
self.seed = seed
self.data_dir = data_dir
self.data_root_dir = data_root_dir
self.scene_indices = scene_indices
self.p_auxiliary = p_auxiliary
self.ablation = ablation
def __init__(self,
details=False,
model_details=False,
save_freq=True,
save_dest=""):
"""
Parameters
----------
details (bool, optional): If true, then frequency information will be shown on screen
model_details (bool, optional): since there are so many model ids, this additional bool
controls if those should be printed
save_freq (bool, optional): if true, then the category frequency information will be saved
save_dest (string, optional): directory to which frequency information is saved
"""
self.room_count = 0
self.object_count = 0
self.room_types_count = {}
self.fine_categories_count = {}
self.coarse_categories_count = {}
self.final_categories_count = {}
self.models_count = {}
self.object_category = ObjectCategories()
self.floor_node_only = False
self.details = details
self.model_details = model_details
self.save_freq = save_freq
data_dir = utils.get_data_root_dir()
self.save_dest = f"{data_dir}/{save_dest}"
def __init__(self, categorization_type='final', sim_mode='direct'):
self._object_data = ObjectData()
self._object_categories = ObjectCategories()
self._objects = {}
self._room = None
self._categorization_type = categorization_type
self._sim = Simulator(mode=sim_mode)
def learn(self, data_folder="bedroom_final", data_root_dir=None):
if not data_root_dir:
data_root_dir = utils.get_data_root_dir()
data_dir = f"{data_root_dir}/{data_folder}"
self.data_dir = data_dir
self.category_map = ObjectCategories()
files = os.listdir(data_dir)
files = [f for f in files if ".pkl" in f and not "domain" in f and not "_" in f]
with open(f"{data_dir}/final_categories_frequency", "r") as f:
lines = f.readlines()
cats = [line.split()[0] for line in lines]
self.category_count = [int(line.split()[1]) for line in lines if line.split()[0] not in ["window", "door"]]
self.categories = [cat for cat in cats if cat not in set(['window', 'door'])]
self.cat_to_index = {self.categories[i]:i for i in range(len(self.categories))}
self.num_categories = len(self.categories)
self.categories.append("floor")
N = self.num_categories
self.support_count = [[0 for i in range(N+1)] for j in range(N)]
for index in range(len(files)):
print(index)
with open(f"{data_dir}/{index}.pkl", "rb") as f:
(_, _, nodes), _ = pickle.load(f)
object_nodes = []
id_to_cat = {}
for node in nodes:
modelId = node["modelId"]
category = self.category_map.get_final_category(modelId)
if not category in ["door", "window"]:
object_nodes.append(node)
id_to_cat[node["id"]] = self.cat_to_index[category]
node["category"] = self.cat_to_index[category]
for node in object_nodes:
parent = node["parent"]
category = node["category"]
if parent == "Floor" or parent is None:
self.support_count[category][-1] += 1
else:
self.support_count[category][id_to_cat[parent]] += 1
#quit()
self.possible_supports={}
for i in range(self.num_categories):
print(f"Support for {self.categories[i]}:")
supports = [(c, self.support_count[i][c]/self.category_count[i]) for c in range(N+1)]
supports = sorted(supports, key = lambda x:-x[1])
supports = [s for s in supports if s[1] > 0.01]
for s in supports:
print(f" {self.categories[s[0]]}:{s[1]:4f}")
self.possible_supports[i] = [s[0] for s in supports]
print(self.possible_supports)
self.N = N
def __init__(self,
scene_indices=(0, 4000),
data_folder="bedroom_fin_256",
data_root_dir=None,
seed=None,
do_rotation_augmentation=False,
cat_only=False,
use_same_category_batches=False,
importance_order=False):
super(LatentDataset, self).__init__()
self.category_map = ObjectCategories()
self.seed = seed
self.data_folder = data_folder
self.data_root_dir = data_root_dir
self.scene_indices = scene_indices
self.do_rotation_augmentation = do_rotation_augmentation
self.cat_only = cat_only
self.cat_name2index = None
self.cat_index2scenes = None
if self.data_root_dir is None:
self.data_root_dir = utils.get_data_root_dir()
with open(
f"{self.data_root_dir}/{self.data_folder}/final_categories_frequency",
"r") as f:
lines = f.readlines()
names = [line.split()[0] for line in lines]
names = [
name for name in names
if ((name != 'door') and (name != 'window'))
]
self.catnames = names
self.cat_name2index = {names[i]: i for i in range(0, len(names))}
self.n_categories = len(names)
self.cat2freq = {}
for line in lines:
cat, freq = line.split(' ')
self.cat2freq[cat] = int(freq)
maxfreq = max(self.cat2freq.values())
self.cat2freq_normalized = {
cat: freq / maxfreq
for cat, freq in self.cat2freq.items()
}
self.build_cat2scene()
self.build_cats_in_scene_indices()
self.compute_cat_sizes()
# See 'prepare_same_category_batches' below for info
self.use_same_category_batches = use_same_category_batches
if use_same_category_batches:
self.same_category_batch_indices = []
else:
self.same_category_batch_indices = None
self.importance_order = importance_order
def __init__(self,
arrangement_priors,
num_angle_divisions=8,
num_pairwise_priors=-1,
sim_mode='direct'):
self._objects = ObjectCollection(sim_mode=sim_mode)
self.room_id = None
self._priors = arrangement_priors
self._num_angle_divisions = num_angle_divisions
self._num_pairwise_priors = num_pairwise_priors
self.category_map = ObjectCategories()
def __init__(self,
scene_indices=(0, 4000),
data_folder="bedroom_fin_256",
data_root_dir=None,
seed=None,
do_rotation_augmentation=False,
cat_only=False,
use_same_category_batches=False,
importance_order=False,
epoch_size=None):
super(LatentDataset, self).__init__()
self.category_map = ObjectCategories()
self.seed = seed
self.data_folder = data_folder
self.data_root_dir = data_root_dir
self.scene_indices = scene_indices
self.do_rotation_augmentation = do_rotation_augmentation
self.cat_only = cat_only
self.cat_name2index = None
self.cat_index2scenes = None
if self.data_root_dir is None:
self.data_root_dir = utils.get_data_root_dir()
self.catnames = self.category_map.all_non_arch_categories(
self.data_root_dir, data_folder)
self.cat_name2index = {
self.catnames[i]: i
for i in range(0, len(self.catnames))
}
self.n_categories = len(self.catnames)
self.build_cat2scene()
self.build_cats_in_scene_indices()
self.cat_importances = self.category_map.all_non_arch_category_importances(
self.data_root_dir, data_folder)
# See 'prepare_same_category_batches' below for info
self.use_same_category_batches = use_same_category_batches
if use_same_category_batches:
self.same_category_batch_indices = []
assert (epoch_size is not None)
self.epoch_size = epoch_size
else:
self.same_category_batch_indices = None
self.importance_order = importance_order
def __init__(self,
scene_indices=(0, 4000),
data_folder="bedroom",
data_root_dir=None,
seed=None):
super(LocDataset, self).__init__()
self.category_map = ObjectCategories()
self.seed = seed
self.data_folder = data_folder
self.data_root_dir = data_root_dir
self.scene_indices = scene_indices
data_root_dir = utils.get_data_root_dir()
with open(f"{data_root_dir}/{data_folder}/final_categories_frequency",
"r") as f:
lines = f.readlines()
self.n_categories = len(lines) - 2 # -2 for 'window' and 'door'
def render_graph(self, room, root, targets):
target_identifiers = list(map(lambda x: x[0], targets))
projection = self.pgen.get_projection(room)
visualization = np.zeros((self.size,self.size))
nodes = []
for i, node in enumerate(room.nodes):
modelId = node.modelId #Camelcase due to original json
t = np.asarray(node.transform).reshape(4,4)
o = Obj(modelId)
# NEW
objspace_width = np.linalg.norm(o.front_left - o.front_right)
objspace_depth = np.linalg.norm(o.front_left - o.back_left)
#END NEW
t = projection.to_2d(t)
o.transform(t)
# NEW
worldspace_width = np.linalg.norm(o.front_left - o.front_right)
worldspace_depth = np.linalg.norm(o.front_left - o.back_left)
#END NEW
t = projection.to_2d()
bbox_min = np.dot(np.asarray([node.xmin, node.zmin, node.ymin, 1]), t)
bbox_max = np.dot(np.asarray([node.xmax, node.zmax, node.ymax, 1]), t)
xmin = math.floor(bbox_min[0])
ymin = math.floor(bbox_min[2])
xsize = math.ceil(bbox_max[0]) - xmin + 1
ysize = math.ceil(bbox_max[2]) - ymin + 1
description = {}
description["modelId"] = modelId
description["transform"] = node.transform
description["bbox_min"] = bbox_min
description["bbox_max"] = bbox_max
#Since it is possible that the bounding box information of a room
#Was calculated without some doors/windows
#We need to handle these cases
if ymin < 0:
ymin = 0
if xmin < 0:
xmin = 0
# render object
rendered = self.render_object(o, xmin, ymin, xsize, ysize, self.size)
description["height_map"] = torch.from_numpy(rendered).float()
tmp = np.zeros((self.size, self.size))
tmp[xmin:xmin+rendered.shape[0],ymin:ymin+rendered.shape[1]] = rendered
# render bbox
for idx, line in enumerate(o.bbox_lines()):
direction = line[1] - line[0]
distance = np.linalg.norm(direction)
norm_direction = direction / distance
for t in range(math.floor(distance)):
point = line[0] + t * norm_direction
x = min(math.floor(point[0]), self.size - 1)
y = min(math.floor(point[2]), self.size - 1)
tmp[x][y] = 1
# temporarily darken image to see more easily
category = ObjectCategories().get_coarse_category(modelId)
identifier = f"{category}_{i}"
if identifier in target_identifiers:
tmp *= 0.8
elif identifier != root:
tmp *= 0.1
else:
for ray in o.front_rays():
for t in range(self.size):
point = ray.origin + t * ray.direction
if point[0] < 0 or point[0] > self.size or point[2] < 0 or point[2] > self.size:
break
color = 1 if int(t * objspace_depth / worldspace_depth) % 2 == 0 else -100
tmp[math.floor(point[0])][math.floor(point[2])] = color
for ray in o.back_rays():
for t in range(self.size):
point = ray.origin + t * ray.direction
if point[0] < 0 or point[0] > self.size or point[2] < 0 or point[2] > self.size:
break
color = 1 if int(t * objspace_depth / worldspace_depth) % 2 else -100
tmp[math.floor(point[0])][math.floor(point[2])] = color
for ray in o.left_rays():
for t in range(self.size):
point = ray.origin + t * ray.direction
if point[0] < 0 or point[0] > self.size or point[2] < 0 or point[2] > self.size:
break
color = 1 if int(t * objspace_width / worldspace_width) % 2 else -100
tmp[math.floor(point[0])][math.floor(point[2])] = color
for ray in o.right_rays():
for t in range(self.size):
point = ray.origin + t * ray.direction
if point[0] < 0 or point[0] > self.size or point[2] < 0 or point[2] > self.size:
break
color = 1 if int(t * objspace_width / worldspace_width) % 2 else -100
tmp[math.floor(point[0])][math.floor(point[2])] = color
visualization += tmp
nodes.append(description)
#Render the floor
o = Obj(room.modelId+"f", room.house_id, is_room=True)
t = projection.to_2d()
o.transform(t)
floor = self.render_object(o, 0, 0, self.size, self.size, self.size)
visualization += floor
floor = torch.from_numpy(floor).float()
#Render the walls
o = Obj(room.modelId+"w", room.house_id, is_room=True)
t = projection.to_2d()
o.transform(t)
wall = self.render_object(o, 0, 0, self.size, self.size, self.size)
visualization += wall
wall = torch.from_numpy(wall).float()
return (visualization, (floor, wall, nodes))
class RenderedScene():
"""
Loading a rendered room
Attributes
----------
category_map (ObjectCategories): object category mapping
that should be the same across all instances of the class
categories (list[string]): all categories present in this room type.
Loaded once when the first room is loaded to reduce disk access.
cat_to_index (dict[string, int]): maps a category to corresponding index
current_data_dir (string): keep track of the current data directory, if
it changes, then categories and cat_to_index should be recomputed
"""
category_map = ObjectCategories()
categories = None
cat_to_index = None
current_data_dir = None
def __init__(self, index, data_dir, data_root_dir=None, \
shuffle=True, load_objects=True, seed=None, rotation=0):
"""
Load a rendered scene from file
Parameters
----------
index (int): room number
data_dir (string): location of the pre-rendered rooms
data_root_dir (string or None, optional): if specified,
use this as the root directory
shuffle (bool, optional): If true, randomly order the objects
in the room. Otherwise use the default order as written
in the original dataset
load_objects (bool, optional): If false, only load the doors
and windows. Otherwise load all objects in the room
seed (int or None, optional): if set, use a fixed random seed
so we can replicate a particular experiment
"""
if seed:
random.seed(seed)
if not data_root_dir:
data_root_dir = utils.get_data_root_dir()
if RenderedScene.categories is None or RenderedScene.current_data_dir != data_dir:
with open(f"{data_root_dir}/{data_dir}/final_categories_frequency",
"r") as f:
lines = f.readlines()
cats = [line.split()[0] for line in lines]
RenderedScene.categories = [
cat for cat in cats if cat not in set(['window', 'door'])
]
RenderedScene.cat_to_index = {
RenderedScene.categories[i]: i
for i in range(len(RenderedScene.categories))
}
RenderedScene.current_data_dir = data_dir
#print(index, rotation)
if rotation != 0:
fname = f"{index}_{rotation}"
else:
fname = index
with open(f"{data_root_dir}/{data_dir}/{fname}.pkl", "rb") as f:
(self.floor, self.wall, nodes), self.room = pickle.load(f)
self.index = index
self.rotation = rotation
self.object_nodes = []
self.door_window_nodes = []
for node in nodes:
category = RenderedScene.category_map.get_final_category(
node["modelId"])
if category in ["door", "window"]:
node["category"] = category
self.door_window_nodes.append(node)
elif load_objects:
node["category"] = RenderedScene.cat_to_index[category]
self.object_nodes.append(node)
if shuffle:
random.shuffle(self.object_nodes)
self.size = self.floor.shape[0]
def create_composite(self):
"""
Create a initial composite that only contains the floor,
wall, doors and windows. See RenderedComposite for how
to add more objects
"""
r = RenderedComposite(RenderedScene.categories, self.floor, self.wall,
self.door_window_nodes)
return r
def learn(self, data_folder, data_root_dir=None):
if data_root_dir is None:
data_root_dir = utils.get_data_root_dir()
data_dir = f"{data_root_dir}/{data_folder}"
self.data_dir = data_dir
self.category_map = ObjectCategories()
files = os.listdir(data_dir)
files = [f for f in files if ".pkl" in f and not "domain" in f]
with open(f"{data_dir}/final_categories_frequency", "r") as f:
lines = f.readlines()
cats = [line.split()[0] for line in lines]
self.categories = [
cat for cat in cats if cat not in set(['window', 'door'])
]
self.cat_to_index = {
self.categories[i]: i
for i in range(len(self.categories))
}
with open(f"{data_dir}/model_frequency", "r") as f:
lines = f.readlines()
models = [line.split()[0] for line in lines]
self.model_freq = [int(l[:-1].split()[1]) for l in lines]
self.models = [
model for model in models if self.category_map.get_final_category(
model) not in set(['window', 'door'])
]
self.model_to_index = {
self.models[i]: i
for i in range(len(self.models))
}
N = len(self.models)
self.num_categories = len(self.categories)
self.model_index_to_cat = [
self.cat_to_index[self.category_map.get_final_category(
self.models[i])] for i in range(N)
]
self.count = [[0 for i in range(N)] for j in range(N)]
for index in range(len(files)):
#for index in range(100):
with open(f"{data_dir}/{index}.pkl", "rb") as f:
(_, _, nodes), _ = pickle.load(f)
object_nodes = []
for node in nodes:
modelId = node["modelId"]
category = self.category_map.get_final_category(modelId)
if not category in ["door", "window"]:
object_nodes.append(node)
for i in range(len(object_nodes)):
for j in range(i + 1, len(object_nodes)):
a = self.model_to_index[object_nodes[i]["modelId"]]
b = self.model_to_index[object_nodes[j]["modelId"]]
self.count[a][b] += 1
self.count[b][a] += 1
print(index, end="\r")
self.N = N
def learn(self, data_folder="bedroom_final", data_root_dir=None):
if not data_root_dir:
data_root_dir = utils.get_data_root_dir()
data_dir = f"{data_root_dir}/{data_folder}"
self.data_dir = data_dir
self.category_map = ObjectCategories()
files = os.listdir(data_dir)
files = [
f for f in files
if ".pkl" in f and not "domain" in f and not "_" in f
]
self.categories = self.category_map.all_non_arch_categories(
data_root_dir, data_folder)
self.cat_to_index = {
self.categories[i]: i
for i in range(len(self.categories))
}
with open(f"{data_dir}/model_frequency", "r") as f:
lines = f.readlines()
models = [line.split()[0] for line in lines]
self.model_freq = [int(l[:-1].split()[1]) for l in lines]
self.models = [
model for model in models if not self.category_map.is_arch(
self.category_map.get_final_category(model))
]
self.model_to_index = {
self.models[i]: i
for i in range(len(self.models))
}
N = len(self.models)
self.num_categories = len(self.categories)
self.model_index_to_cat = [
self.cat_to_index[self.category_map.get_final_category(
self.models[i])] for i in range(N)
]
self.count = [[0 for i in range(N)] for j in range(N)]
for index in range(len(files)):
#for index in range(100):
with open(f"{data_dir}/{index}.pkl", "rb") as f:
(_, _, nodes), _ = pickle.load(f)
object_nodes = []
for node in nodes:
modelId = node["modelId"]
category = self.category_map.get_final_category(modelId)
if not self.category_map.is_arch(category):
object_nodes.append(node)
for i in range(len(object_nodes)):
for j in range(i + 1, len(object_nodes)):
a = self.model_to_index[object_nodes[i]["modelId"]]
b = self.model_to_index[object_nodes[j]["modelId"]]
self.count[a][b] += 1
self.count[b][a] += 1
print(index)
self.N = N
use_jitter = False
jitter_stdev = 0.01
which_to_load = 500
parser = argparse.ArgumentParser(description='orient')
parser.add_argument('--save-dir', type=str, required=True)
parser.add_argument('--data-folder', type=str, default="")
args = parser.parse_args()
outdir = f'./output/{args.save_dir}'
utils.ensuredir(outdir)
data_folder = args.data_folder
data_root_dir = utils.get_data_root_dir()
categories = ObjectCategories().all_non_arch_categories(
data_root_dir, data_folder)
num_categories = len(categories)
num_input_channels = num_categories + 8
nc = num_categories
# Dataset size is based on the number of available scenes - the valid set size
dataset_size = len([f for f in os.listdir(f'{data_root_dir}/{data_folder}') \
if f.endswith('.jpg')]) - valid_set_size
dataset_size = int(dataset_size / batch_size) * batch_size
logfile = open(f"{outdir}/log.txt", 'w')
def LOG(msg):
print(msg)
logfile.write(msg + '\n')
class GlobalCategoryFilter():
category_map = ObjectCategories()
def get_filter():
door_window = GlobalCategoryFilter.category_map.all_arch_categories()
second_tier = [
"table_lamp", "chandelier", "guitar", "amplifier", "keyboard",
"drumset", "microphone", "accordion", "toy", "xbox", "playstation",
"fishbowl", "chessboard", "iron", "helmet", "telephone",
"stationary_container", "ceiling_fan", "bottle", "fruit_bowl",
"glass", "knife_rack", "plates", "books", "book", "television",
"wood_board", "switch", "pillow", "laptop", "clock", "helmet",
"bottle", "trinket", "glass", "range_hood", "candle", "soap_dish"
]
wall_objects = ["wall_lamp", "mirror", "curtain", "blind"]
unimportant = [
"toy", "fish_tank", "tricycle", "vacuum_cleaner", "weight_scale",
"heater", "picture_frame", "beer", "shoes", "weight_scale",
"decoration", "ladder", "tripod", "air_conditioner", "cart",
"fireplace_tools", "vase"
]
inhabitants = ["person", "cat", "bird", "dog", "pet"]
special_filter = [
"rug",
]
filtered = second_tier + unimportant + inhabitants + special_filter + wall_objects
unwanted_complex_structure = ["partition", "column", "arch", "stairs"]
set_items = [
"chair_set", "stereo_set", "table_and_chair",
"slot_machine_and_chair", "kitchen_set", "double_desk",
"double_desk_with_chairs", "dressing_table_with_stool",
"kitchen_island_with_range_hood_and_table"
]
outdoor = [
"lawn_mower", "car", "motorcycle", "bicycle", "garage_door",
"outdoor_seating", "fence"
]
rejected = unwanted_complex_structure + set_items + outdoor
return filtered, rejected, door_window
def get_filter_latent():
door_window = GlobalCategoryFilter.category_map.all_arch_categories()
second_tier_include = [
"table_lamp",
"television",
"picture_frame",
"books",
"book",
"laptop",
"floor_lamp",
"vase",
"plant",
"console",
"stereo_set",
"toy",
"fish_tank",
"cup",
"glass",
"fruit_bowl",
"bottle",
"fishbowl",
"pillow",
]
second_tier = [
"chandelier", "guitar", "amplifier", "keyboard", "drumset",
"microphone", "accordion", "chessboard", "iron", "helmet",
"stationary_container", "ceiling_fan", "knife_rack", "plates",
"wood_board", "switch", "clock", "helmet", "trinket", "range_hood",
"candle", "soap_dish"
]
wall_objects = [
"wall_lamp", "mirror", "curtain", "wall_shelf", "blinds", "blind"
]
unimportant = [
"tricycle",
"fish_tank",
"vacuum_cleaner",
"weight_scale",
"heater",
"picture_frame",
"beer",
"shoes",
"weight_scale",
"decoration",
"ladder",
"tripod",
"air_conditioner",
"cart",
"fireplace_tools",
"ironing_board",
]
inhabitants = ["person", "cat", "bird", "dog", "pet"]
special_filter = [
"rug",
]
filtered = second_tier + unimportant + inhabitants + special_filter + wall_objects
unwanted_complex_structure = ["partition", "column", "arch", "stairs"]
set_items = [
"chair_set", "stereo_set", "table_and_chair",
"slot_machine_and_chair", "kitchen_set", "double_desk",
"double_desk_with_chairs", "desk_with_shelves",
"dressing_table_with_stool", "armchair_with_ottoman",
"kitchen_island_with_range_hood_and_table"
]
outdoor = [
"lawn_mower", "car", "motorcycle", "bicycle", "garage_door",
"outdoor_seating", "fence"
]
rejected = unwanted_complex_structure + set_items + outdoor
return filtered, rejected, door_window, second_tier_include
class RenderedScene():
"""
Loading a rendered room
Attributes
----------
category_map (ObjectCategories): object category mapping
that should be the same across all instances of the class
categories (list[string]): all categories present in this room type.
Loaded once when the first room is loaded to reduce disk access.
cat_to_index (dict[string, int]): maps a category to corresponding index
current_data_dir (string): keep track of the current data directory, if
it changes, then categories and cat_to_index should be recomputed
"""
category_map = ObjectCategories()
categories = None
cat_to_index = None
current_data_dir = None
def __init__(self, index, data_dir, data_root_dir=None, \
shuffle=True, load_objects=True, seed=None, rotation=0):
"""
Load a rendered scene from file
Parameters
----------
index (int): room number
data_dir (string): location of the pre-rendered rooms
data_root_dir (string or None, optional): if specified,
use this as the root directory
shuffle (bool, optional): If true, randomly order the objects
in the room. Otherwise use the default order as written
in the original dataset
load_objects (bool, optional): If false, only load the doors
and windows. Otherwise load all objects in the room
seed (int or None, optional): if set, use a fixed random seed
so we can replicate a particular experiment
"""
if seed:
random.seed(seed)
if not data_root_dir:
data_root_dir = utils.get_data_root_dir()
if RenderedScene.categories is None or RenderedScene.current_data_dir != data_dir:
RenderedScene.categories = RenderedScene.category_map.all_non_arch_categories(
data_root_dir, data_dir)
RenderedScene.cat_to_index = {
RenderedScene.categories[i]: i
for i in range(len(RenderedScene.categories))
}
RenderedScene.current_data_dir = data_dir
#print(RenderedScene.cat_to_index)
#print(index, rotation)
if rotation != 0:
fname = f"{index}_{rotation}"
else:
fname = index
with open(f"{data_root_dir}/{data_dir}/{fname}.pkl", "rb") as f:
(self.floor, wall, nodes), self.room = pickle.load(f)
self.size = self.floor.shape[0]
self.wall = wall["height_map"]
self.wall_segments = wall["segments"]
# Compute normals for wall segment objects
mask = (self.floor + self.wall) != 0
for wall_seg in self.wall_segments:
seg = [np.array([x[0], x[2]])
for x in wall_seg["points"]] # reduce to 2d
v_diff = seg[1] - seg[0]
v_norm = np.array([-v_diff[1], v_diff[0]])
v_norm = v_norm / np.linalg.norm(v_norm)
#### Decide whether to flip the normal vector
midp = (seg[0] + seg[1]) / 2
test_point = np.ceil(midp +
5 * v_norm * np.array([-1, 1])).astype(int)
# If (1) a point slightly along the normal falls outside the image
# (2) the pixel in the room mask image at this point is not filled
# then this must be pointing out of the room and thus should be flipped.
out = (test_point >= self.size).any() or (test_point < 0).any()
if out or not mask[int(test_point[0]), int(test_point[1])]:
v_norm = -v_norm
wall_seg["normal"] = v_norm
wall_seg["points"] = seg
self.index = index
self.rotation = rotation
self.object_nodes = []
self.door_window_nodes = []
for node in nodes:
category = RenderedScene.category_map.get_final_category(
node["modelId"])
if RenderedScene.category_map.is_arch(category):
node["category"] = category
self.door_window_nodes.append(node)
elif load_objects:
node["category"] = RenderedScene.cat_to_index[category]
self.object_nodes.append(node)
if shuffle:
random.shuffle(self.object_nodes)
def create_composite(self):
"""
Create a initial composite that only contains the floor,
wall, doors and windows. See RenderedComposite for how
to add more objects
"""
r = RenderedComposite(RenderedScene.categories, self.floor, self.wall,
self.door_window_nodes)
return r