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
def get_filter(source, filter_type, *args):
#Just brute force enumerate possible implemented filters
#See filters/
if filter_type == "good_house":
from filters.good_house import good_house_criteria
house_f = good_house_criteria
dataset_f = DatasetFilter(house_filters=[house_f])
elif filter_type == "room_type":
from filters.room_type import room_type_criteria
room_f = room_type_criteria(*args)
dataset_f = DatasetFilter(room_filters=[room_f])
elif filter_type == "bedroom":
from filters.bedroom import bedroom_filter
dataset_f = bedroom_filter(*args, source)
elif filter_type == "office":
from filters.office import office_filter
dataset_f = office_filter(*args, source)
elif filter_type == "livingroom":
from filters.livingroom import livingroom_filter
dataset_f = livingroom_filter(*args, source)
elif filter_type == "floor_node":
from filters.floor_node import floor_node_filter
dataset_f = floor_node_filter(*args)
elif filter_type == "renderable":
from filters.renderable import renderable_room_filter
dataset_f = renderable_room_filter(*args)
elif filter_type == "collision":
from filters.collision import collision_filter
from priors.observations import ObjectCollection
oc = ObjectCollection()
dataset_f = collision_filter(oc)
else:
raise NotImplementedError
return dataset_f
def __init__(self, location_epoch, rotation_epoch, continue_epoch, \
data_dir, model_dir, data_root_dir=None, model_root_dir=None, size=512):
"""
Parameters
----------
location_epoch, rotation_epoch, continue_epoch (int):
the epoch number of the respective trained models to be loaded
data_dir (string): location of the dataset relative to data_root_dir
model_dir (string): location of the trained models relative to model_root_dir
data_root_dir (string or None, optional): if not set, use the default data location,
see utils.get_data_root_dir
model_root_dir (string or None, optional): if not set, use the directory of this script
size (int): size of the input image
"""
Node.warning = False
self.data_dir_relative = data_dir #For use in RenderedScene
if not data_root_dir:
self.data_root_dir = utils.get_data_root_dir()
self.data_dir = f"{self.data_root_dir}/{data_dir}"
if not model_root_dir:
model_root_dir = os.path.dirname(os.path.abspath(__file__))
self.model_dir = f"{model_root_dir}/{model_dir}"
#Loads category and model information
self.categories, self.cat_to_index = self._load_category_map()
self.num_categories = len(self.categories)
self.possible_models = self._load_possible_models()
self.model_set_list = self._load_model_set_list()
#Loads trained models and build up NNs
self.model_location, self.fc_location = self._load_location_model(
location_epoch)
self.model_rotation = self._load_rotation_model(rotation_epoch)
self.model_continue, self.fc_continue = self._load_continue_model(
continue_epoch)
self.softmax = nn.Softmax(dim=1)
self.softmax.cuda()
#Misc Handling
self.pgen = ProjectionGenerator()
self.model_sampler = ModelPrior()
self.model_sampler.load(self.data_dir)
self.object_collection = ObjectCollection()
class ArrangementGreedySampler:
"""
Iterative optimization of object arrangements using greedy sampling of ArrangementPriors
"""
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()
@property
def objects(self):
return self._objects
def init(self, house, only_architecture=True, room_id=None):
if not room_id:
room_id = house.rooms[0].id
self._objects.init_from_room(house,
room_id,
only_architecture=only_architecture)
self.room_id = room_id
def log_prob(self, filter_ref_obj=None, ignore_categories=list()):
observations = self._objects.get_relative_observations(
self.room_id,
filter_ref_obj=filter_ref_obj,
ignore_categories=ignore_categories)
observations_by_key = {}
# top_k_prior_categories = None
# if filter_ref_obj and num_pairwise_priors specified, filter observations to only those in top k priors
# if filter_ref_obj and self._num_pairwise_priors > 0:
# category = self._objects.category(filter_ref_obj.modelId, scheme='final')
# priors = list(filter(lambda p: p.ref_obj_category == category, self._priors.pairwise_priors))
# k = min(self._num_pairwise_priors, len(priors))
# priors = list(sorted(priors, key=lambda p: self._priors.pairwise_occurrence_log_prob(p)))[-k:]
# top_k_prior_categories = set(map(lambda p: p.obj_category, priors))
for o in observations.values():
# only pairwise observations in which filter_ref_obj is the reference
if filter_ref_obj and o.ref_id != filter_ref_obj.id:
continue
key = self._objects.get_observation_key(o)
os_key = observations_by_key.get(key, [])
os_key.append(o)
observations_by_key[key] = os_key
return self._priors.log_prob(observations_by_key)
def get_candidate_transform(self, node, max_iterations=100):
num_checks = 0
zmin = self._objects.room.zmin
while True:
num_checks += 1
p = self._objects.room.obb.sample()
ray_from = [p[0], zmin - .5, p[2]]
ray_to = [p[0], zmin + .5, p[2]]
intersection = ags._objects.simulator.ray_test(ray_from, ray_to)
if intersection.id == self.room_id + 'f' or num_checks > max_iterations:
break
xform = Transform()
xform.set_translation([p[0], zmin + .1, p[2]])
angle = random() * 2 * math.pi
angular_resolution = 2 * math.pi / self._num_angle_divisions
angle = round(angle / angular_resolution) * angular_resolution
xform.set_rotation(radians=angle)
return xform
def sample_placement(self,
node,
n_samples,
houses_log=None,
max_attempts_per_sample=10,
ignore_categories=list(),
collision_threshold=0):
"""
Sample placement for given node
"""
self._objects.add_object(node)
max_lp = -np.inf
max_xform = None
max_house = None
num_noncolliding_samples = 0
for i in range(max_attempts_per_sample * n_samples):
xform = self.get_candidate_transform(node)
self._objects.update(node, xform=xform, update_sim=True)
collisions = self._objects.get_collisions(obj_id_a=node.id)
# print(f'i={i}, samples_so_far={num_noncolliding_samples}, n_samples={n_samples},'
# f'max_attempts_per_sample={max_attempts_per_sample}')
if collision_threshold > 0:
if min(collisions.values(), key=lambda c: c.distance
).distance < -collision_threshold:
continue
elif len(collisions) > 0:
continue
lp = self.log_prob(filter_ref_obj=node,
ignore_categories=ignore_categories)
print(f'lp={lp}')
if lp > max_lp:
max_xform = xform
max_lp = lp
if houses_log is not None:
max_house = self._objects.as_house()
num_noncolliding_samples += 1
if num_noncolliding_samples == n_samples:
break
if houses_log is not None:
houses_log.append(max_house)
self._objects.update(node, xform=max_xform, update_sim=True)
def placeable_objects_sorted_by_size(self, house):
objects = []
fixed_objects = []
for n in house.levels[0].nodes:
if n.type != 'Object':
continue
category = self._objects.category(n.modelId, scheme='final')
if self.category_map.is_arch(category):
fixed_objects.append(n)
continue
objects.append(n)
for o in objects:
# to_delete = []
# for k in o.__dict__:
# if k not in ['id', 'modelId', 'transform', 'type', 'valid', 'bbox']:
# to_delete.append(k)
# for k in to_delete:
# delattr(o, k)
dims = self._objects._object_data.get_aligned_dims(o.modelId)
o.volume = dims[0] * dims[1] * dims[2]
objects = list(sorted(objects, key=lambda x: x.volume, reverse=True))
return objects, fixed_objects
def __init__(self, data_dir, data_root_dir=None, size=256):
"""
Parameters
----------
location_epoch, rotation_epoch, continue_epoch (int):
the epoch number of the respective trained models to be loaded
data_dir (string): location of the dataset relative to data_root_dir
data_root_dir (string or None, optional): if not set, use the default data location,
see utils.get_data_root_dir
size (int): size of the input image
"""
Node.warning = False
self.data_dir_relative = data_dir #For use in RenderedScene
if not data_root_dir:
self.data_root_dir = utils.get_data_root_dir()
self.data_dir = f"{self.data_root_dir}/{data_dir}"
#Loads category and model information
self.categories, self.cat_to_index = self._load_category_map()
#Hardcode second tier here because we don't have planit yet
self.second_tiers = [
"table_lamp",
"television",
"picture_frame",
"books",
"book",
"laptop",
"vase",
"plant",
"console",
"stereo_set",
"toy",
"fish_tank",
"cup",
"glass",
"fruit_bowl",
"bottle",
"fishbowl",
"pillow",
]
self.num_categories = len(self.categories)
#Misc Handling
self.pgen = ProjectionGenerator()
self.possible_models = self._load_possible_models()
self.model_set_list = self._load_model_set_list()
#Loads trained models and build up NNs
self.model_cat = self._load_category_model()
self.model_location, self.fc_location = self._load_location_model()
self.model_orient = self._load_orient_model()
self.model_dims = self._load_dims_model()
self.softmax = nn.Softmax(dim=1)
self.softmax.cuda()
self.model_sampler = ModelPrior()
self.model_sampler.load(self.data_dir)
sp = SupportPrior()
sp.load(self.data_dir)
self.possible_supports = sp.possible_supports
self.obj_data = ObjectData()
self.object_collection = ObjectCollection()