class ModelPrior():
def __init__(self):
pass
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
def save(self, dest=None):
if dest == None:
dest = f"{self.data_dir}/model_prior.pkl"
with open(dest, "wb") as f:
pickle.dump(self.__dict__, f, pickle.HIGHEST_PROTOCOL)
def load(self, data_dir):
source = f"{data_dir}/model_prior.pkl"
with open(source, "rb") as f:
self.__dict__ = pickle.load(f)
def sample(self, category, models):
N = self.N
indices = [
i for i in range(N) if self.model_index_to_cat[i] == category
]
p = [self.model_freq[indices[i]] for i in range(len(indices))]
p = np.asarray(p)
for model in models:
i = self.model_to_index[model]
p1 = [self.count[indices[j]][i] for j in range(len(indices))]
p1 = np.asarray(p1)
p1 = p1 / p1.sum()
p = p * p1
p = p / sum(p)
numbers = np.asarray([i for i in range(len(indices))])
return self.models[indices[np.random.choice(numbers, p=p)]]
def get_models(self, category, important, others):
N = self.N
indices = [
i for i in range(N) if self.model_index_to_cat[i] == category
]
to_remove = []
freq = [self.model_freq[indices[i]] for i in range(len(indices))]
total_freq = sum(freq)
for j in range(len(indices)):
if freq[j] / total_freq < 0.01:
if not indices[j] in to_remove:
to_remove.append(indices[j])
for model in important:
i = self.model_to_index[model]
freq = [self.count[indices[j]][i] for j in range(len(indices))]
total_freq = sum(freq)
if total_freq > 0:
for j in range(len(indices)):
if freq[j] / total_freq < 0.1:
if not indices[j] in to_remove:
to_remove.append(indices[j])
for model in others:
i = self.model_to_index[model]
freq = [self.count[indices[j]][i] for j in range(len(indices))]
total_freq = sum(freq)
if total_freq > 0:
for j in range(len(indices)):
if freq[j] / total_freq < 0.05:
if not indices[j] in to_remove:
to_remove.append(indices[j])
for item in to_remove:
if len(indices) > 1:
indices.remove(item)
return [self.models[index] for index in indices]
class LatentDataset(data.Dataset):
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
# Build a map from category index to the scene indices that contain an instance of that category
# This ignores scene_indices and does it for the whole data folder
def build_cat2scene(self):
self.cat_index2scenes = defaultdict(list)
data_root_dir = self.data_root_dir or utils.get_data_root_dir()
data_dir = f'{data_root_dir}/{self.data_folder}'
filename = f'{data_dir}/cat_index2scenes'
# Create new cached map file
if not os.path.exists(filename):
print(
'Building map of category to scenes containing an instance...')
pkls = [
path for path in os.listdir(data_dir) if path.endswith('.pkl')
]
pklnames = [os.path.splitext(path)[0] for path in pkls]
# Only get the .pkl files which are numbered scenes
indices = [
int(pklname) for pklname in pklnames if pklname.isdigit()
]
i = 0
for idx in indices:
i += 1
sys.stdout.write(f' {i}/{len(indices)}\r')
sys.stdout.flush()
scene = RenderedScene(idx, self.data_folder,
self.data_root_dir)
object_nodes = scene.object_nodes
for node in object_nodes:
self.cat_index2scenes[node['category']].append(idx)
pickle.dump(self.cat_index2scenes, open(filename, 'wb'))
print('')
# Load an existing cached map file from disk
else:
self.cat_index2scenes = pickle.load(open(filename, 'rb'))
def __len__(self):
if self.use_same_category_batches:
return self.epoch_size
else:
return self.scene_indices[1] - self.scene_indices[0]
# First, find the set of categories that occur within the scene indices
# We do this because it's possible that there might be some category that
# occurs in the dataset, but only in the test set...
def build_cats_in_scene_indices(self):
cats_seen = {}
for cat, scene_indices in self.cat_index2scenes.items():
scenes = [idx for idx in scene_indices if \
(idx >= self.scene_indices[0] and idx < self.scene_indices[1])]
if len(scenes) > 0:
cats_seen[cat] = True
cats_seen = list(cats_seen.keys())
self.cats_seen = cats_seen
# Use at the beginning of each epoch to support loading batches of all the same category
# NOTE: The data loader must have shuffle set to False for this to work
def prepare_same_category_batches(self, batch_size):
# Build a random list of category indices (grouped by batch_size)
# This requires than length of dataset is a multiple of batch_size
assert (len(self) % batch_size == 0)
num_batches = len(self) // batch_size
self.same_category_batch_indices = []
for i in range(num_batches):
# cat_index = random.randint(0, self.n_categories-1)
cat_index = random.choice(self.cats_seen)
for j in range(batch_size):
self.same_category_batch_indices.append(cat_index)
# 'importance' = a function of both size and observation frequency
def sort_object_nodes_by_importance(self,
object_nodes,
noise=None,
swap_prob=None):
# Build list of pairs of (index, importance)
index_imp_pairs = []
for i in range(0, len(object_nodes)):
node = object_nodes[i]
cat = node["category"]
imp = self.cat_importances[cat]
index_imp_pairs.append((i, imp))
# Optionally, add noise to these importance scores
# Noise is expressed as a multiple of the standard deviation of the importance scores
# A typical value might be really small, e.g. 0.05(?)
if noise is not None:
imps = [pair[1] for pair in index_imp_pairs]
istd = np.array(imps).std()
index_imp_pairs = [(index,
imp + noise * random.normalvariate(0, istd))
for index, imp in index_imp_pairs]
# Sort based on importance
index_imp_pairs.sort(key=lambda tup: tup[1], reverse=True)
sorted_nodes = [object_nodes[tup[0]] for tup in index_imp_pairs]
# Optionally, swap nodes with some probabilitiy
if swap_prob is not None:
indices = list(range(len(sorted_nodes)))
for i in range(len(indices)):
if random.random() < swap_prob:
indices_ = list(range(len(sorted_nodes)))
idx1 = random.choice(indices_)
indices_.remove(idx1)
idx2 = random.choice(indices_)
tmp = indices[idx1]
indices[idx1] = indices[idx2]
indices[idx2] = tmp
tmp = sorted_nodes[idx1]
sorted_nodes[idx1] = sorted_nodes[idx2]
sorted_nodes[idx2] = tmp
return sorted_nodes
def order_object_nodes(self, object_nodes):
if self.importance_order:
object_nodes = self.sort_object_nodes_by_importance(object_nodes)
else:
object_nodes = object_nodes[:]
random.shuffle(object_nodes)
# The following extra sorting passes only apply to datasets that have second-tier objects
# We can check for this by looking for the presence of certain object properties e.g. 'parent'
if 'parent' in object_nodes[0]:
# Make sure that all second-tier objects come *after* first tier ones
def is_second_tier(node):
return (node['parent'] != 'Wall') and \
(node['parent'] != 'Floor')
object_nodes.sort(key=lambda node: int(is_second_tier(node)))
# Make sure that all children come after their parents
def cmp_parent_child(node1, node2):
# Less than (negative): node1 is the parent of node2
if node2['parent'] == node1['id']:
return -1
# Greater than (postive): node2 is the parent of node1
elif node1['parent'] == node2['id']:
return 1
# Equal (zero): all other cases
else:
return 0
object_nodes.sort(key=cmp_to_key(cmp_parent_child))
# #### TEST: make sure office_chair comes *after* desk
# def chair_vs_desk(node1, node2):
# catname1 = self.catnames[node1['category']]
# catname2 = self.catnames[node2['category']]
# if (catname1 == 'desk') and (catname2 == 'office_chair'):
# return -1
# elif (catname1 == 'office_chair') and (catname2 == 'desk'):
# return 1
# else:
# return 0
# object_nodes.sort(key = cmp_to_key(chair_vs_desk))
return object_nodes
def get_scene(self, index, stop_prob=None):
i = index + self.scene_indices[0]
scene = RenderedScene(i, self.data_folder, self.data_root_dir)
object_nodes = self.order_object_nodes(scene.object_nodes)
# With some probability, sample a 'stop' (i.e. the end of the scene build sequence)
if stop_prob is not None and random.random() < stop_prob:
output_node = None
input_nodes = object_nodes
else:
# Pick a random index at which to split into (a) existing objects and
# (b) objects yet-to-be-added.
split_idx = random.randint(0, len(object_nodes) - 1)
# This object is the ouput node
output_node = object_nodes[split_idx]
# All object before this index are input nodes
input_nodes = object_nodes[0:split_idx]
return scene, input_nodes, output_node
def get_scene_specific_category(self, cat_index_or_name, empty_room=False):
if isinstance(cat_index_or_name, list):
cat_index_or_name = random.choice(cat_index_or_name)
if isinstance(cat_index_or_name, int):
cat_index = cat_index_or_name
else:
cat_name = cat_index_or_name
cat_index = self.cat_name2index[cat_name]
# Pull out a scene (within scene_indices) that has an instance of this category
scenes_for_cat = [idx for idx in self.cat_index2scenes[cat_index] if \
(idx >= self.scene_indices[0] and idx < self.scene_indices[1])]
scene_index = random.choice(scenes_for_cat)
scene = RenderedScene(scene_index, self.data_folder,
self.data_root_dir)
object_nodes = self.order_object_nodes(scene.object_nodes)
# Pick a random instance of the category
cat_indices = [
i for i in range(0, len(object_nodes))
if object_nodes[i]['category'] == cat_index
]
split_idx = random.choice(cat_indices)
# This object is the ouput node
output_node = object_nodes[split_idx]
if empty_room:
input_nodes = [] # No other objects in the scene
else:
input_nodes = object_nodes[
0:split_idx] # All object before this index are input nodes
return scene, input_nodes, output_node
def get_scene_same_category_batch(self, index):
cat_index = self.same_category_batch_indices[index]
return self.get_scene_specific_category(cat_index)
# Balance training data so that we train equally often on all target categories
def get_scene_uniform_category(self, stop_prob=None):
if stop_prob is not None and random.random() < stop_prob:
scene_index = random.randint(self.scene_indices[0],
self.scene_indices[1] - 1)
scene = RenderedScene(scene_index, self.data_folder,
self.data_root_dir)
output_node = None
input_nodes = self.order_object_nodes(scene.object_nodes)
return scene, input_nodes, output_node
else:
cat_index = random.choice(self.cats_seen)
return self.get_scene_specific_category(cat_index)
def __getitem__(self, index):
if self.seed:
random.seed(self.seed)
if self.use_same_category_batches:
scene, input_nodes, output_node = self.get_scene_same_category_batch(
index)
elif self.cat_only:
scene, input_nodes, output_node = self.get_scene(index,
stop_prob=0.1)
else:
scene, input_nodes, output_node = self.get_scene(index)
# Get the composite images
if not self.do_rotation_augmentation:
input_img = create_transformed_composite(scene, input_nodes, 0)
if not self.cat_only:
output_img = create_transformed_composite(
scene, [output_node], 0)
else:
# Data augmentation: Get the composite images under a random cardinal rotation
rot = random.choice([0, 90, 180, 270])
input_img = create_transformed_composite(scene, input_nodes, rot)
if not self.cat_only:
output_img = create_transformed_composite(
scene, [output_node], rot)
# Get the category of the object
# This is an integer index
if output_node is None:
cat = torch.LongTensor([self.n_categories])
else:
cat = torch.LongTensor([output_node["category"]])
# Also get the count of all categories currently in the scene
catcount = torch.zeros(self.n_categories)
for node in input_nodes:
catidx = node['category']
catcount[catidx] = catcount[catidx] + 1
# If the dataset is configured to only care about predicting the category, then we can go ahead
# and return now
if self.cat_only:
return input_img, cat, catcount
# Select just the object mask channel from the output image
output_img = output_img[2]
# Put a singleton dimension back in for the channel dimension
output_img = torch.unsqueeze(output_img, 0)
# Make sure that it has value 1 everywhere (hack: multiply by huge number and clamp)
output_img *= 1000
torch.clamp(output_img, 0, 1, out=output_img) # Clamp in place
# Get the location of the object
# Normalize the coordinates to [-1, 1], with (0,0) being the image center
loc = output_node['location']
x = loc[0]
y = loc[1]
w = output_img.size()[2]
x_ = ((x / w) - 0.5) * 2
y_ = ((y / w) - 0.5) * 2
loc = torch.Tensor([x_, y_])
# Get the orientation of the object
# Here, we assume that there is no scale, and that the only rotation is about the up vector
# (so we can just read the cos, sin values directly out of the transformation matrix)
xform = output_node["transform"]
cos = xform[0]
sin = xform[8]
orient = torch.Tensor([cos, sin])
# Get the object-space dimensions of the output object (in pixel space)
# (Normalize to [0, 1])
xsize, ysize = output_node['objspace_dims']
xsize = xsize / w
ysize = ysize / w
# dims = torch.Tensor([xsize, ysize])
dims = torch.Tensor([ysize, xsize
]) # Not sure why this flip is necessary atm...
return input_img, output_img, cat, loc, orient, dims, catcount
class SupportPrior():
def __init__(self):
pass
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.category_count = self.category_map.all_non_arch_category_counts(
data_root_dir, data_folder)
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 self.category_map.is_arch(category):
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 save(self, dest=None):
if dest == None:
dest = f"{self.data_dir}/support_prior.pkl"
with open(dest, "wb") as f:
pickle.dump(self.__dict__, f, pickle.HIGHEST_PROTOCOL)
def load(self, data_dir):
source = f"{data_dir}/support_prior.pkl"
with open(source, "rb") as f:
self.__dict__ = pickle.load(f)