def __getitem__(self, index):
if self.seed:
random.seed(self.seed)
i = int(index + self.scene_indices[0] / self.num_per_epoch)
scene = RenderedScene(i, self.data_dir, self.data_root_dir)
composite = scene.create_composite()
num_categories = len(scene.categories)
existing_categories = torch.zeros(num_categories)
# Flip a coin for whether we're going remove objects or treat this as a complete scene
is_complete = random.random() < self.complete_prob
if not is_complete:
# If we decide to remove objects, then remove a random number of them
num_objects = random.randint(0, len(scene.object_nodes) - 1)
else:
num_objects = len(scene.object_nodes)
for i in range(num_objects):
node = scene.object_nodes[i]
composite.add_node(node)
existing_categories[node["category"]] += 1
inputs = composite.get_composite(num_extra_channels=0,
ablation=self.ablation)
# Output is a boolean for "should we continue adding objects?"
output = not is_complete
return inputs, output, existing_categories
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 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 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)
input_nodes = object_nodes
output_node = None
return scene, input_nodes, output_node
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 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 __getitem__(self, index):
if self.seed:
random.seed(self.seed)
i = index + self.scene_indices[0]
scene = RenderedScene(i, self.data_folder, self.data_root_dir)
composite = scene.create_composite()
object_nodes = scene.object_nodes
random.shuffle(object_nodes)
if 'parent' in object_nodes[0]:
#print([a["category"] for a in object_nodes])
# 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))
#print([a["category"] for a in object_nodes])
#print("________________")
num_objects = random.randint(0, len(object_nodes))
#num_objects = len(object_nodes)
num_categories = len(scene.categories)
centroids = []
parent_ids = ["Floor", "Wall"]
for i in range(num_objects):
node = object_nodes[i]
if node["parent"] == "Wall":
print("Massive messup!")
composite.add_node(node)
xsize, ysize = node["height_map"].shape
xmin, _, ymin, _ = node["bbox_min"]
xmax, _, ymax, _ = node["bbox_max"]
parent_ids.append(node["id"])
inputs = composite.get_composite(num_extra_channels=0)
size = inputs.shape[1]
for i in range(num_objects, len(object_nodes)):
node = object_nodes[i]
if node["parent"] == "Wall":
print("Massive messup!")
if node["parent"] in parent_ids:
xsize, ysize = node["height_map"].shape
xmin, _, ymin, _ = node["bbox_min"]
xmax, _, ymax, _ = node["bbox_max"]
centroids.append(
((xmin + xmax) / 2, (ymin + ymax) / 2, node["category"]))
output = torch.zeros((64, 64)).long()
for (x, y, label) in centroids:
output[math.floor(x / 4), math.floor(y / 4)] = label + 1
return inputs, output
def __getitem__(self, index):
if self.seed:
random.seed(self.seed)
i = int(index + self.scene_indices[0] / self.num_per_epoch)
scene = RenderedScene(i, self.data_dir, self.data_root_dir)
composite = scene.create_composite()
num_categories = len(scene.categories)
# Flip a coin for whether we're going remove objects or treat this as a complete scene
num_objects = len(scene.object_nodes)
object_nodes = scene.object_nodes
# 理解：p_existing是输入的p
# 一个场景的num_categories数量固定，每种标签至少有一个物体，但可能不止一个
# 因此，列数是num_categories，而行数暂时先多填了5个
# 疑问：1.每个场景的num_categories不一样，形成的one-hot vector长度不一致，是否
# 需要改成固定长度的num_categories
p_existing = torch.zeros(num_objects, num_categories)
for i in range(num_objects):
existing_categories = torch.zeros(num_categories)
node = scene.object_nodes[i]
composite.add_node(node)
existing_categories[node["category"]] = 1
p_existing[i] = existing_categories
coordinates_existing = torch.zeros(num_objects, 4)
wall = scene.wall
wall_mask = wall.clone()
index_nonzero = torch.nonzero(wall_mask)
xmin_scene, ymin_scene = index_nonzero[0][0], index_nonzero[0][1]
xmax_scene, ymax_scene = index_nonzero[index_nonzero.shape[0] - 1][0], \
index_nonzero[index_nonzero.shape[0] - 1][1]
for i in range(num_objects):
#existing_coordinates = torch.zeros(4)
node = object_nodes[i]
xmin, _, ymin, _ = node["bbox_min"]
xmax, _, ymax, _ = node["bbox_max"]
# TO DO
# 1 scale coordinates(need to pre-define the height and width of map)
# 获取房间俯视图的xmin_scene,xmax_scene,ymin_scene,ymax_scene
# 将坐标归一化到0-1之间(房间的边缘是0和1)
xmin = (xmin - xmin_scene) / (xmax_scene - xmin_scene).double()
xmax = (xmax - xmin_scene) / (xmax_scene - xmin_scene).double()
ymin = (ymin - ymin_scene) / (ymax_scene - ymin_scene).double()
ymax = (ymax - ymin_scene) / (ymax_scene - ymin_scene).double()
existing_coordinates = torch.Tensor((xmin, ymin, xmax, ymax))
coordinates_existing[i] = existing_coordinates
existing_object = torch.cat((p_existing, coordinates_existing), 1)
non_existing = torch.zeros(num_categories + 5 - num_objects,
num_categories + 4)
output = torch.cat((existing_object, non_existing), 0)
#print("output shape=",output.shape)
return output
def get_scene(self, index):
return RenderedScene(index,
self.data_dir,
self.data_root_dir,
shuffle=False)
def __getitem__(self, index):
if self.seed:
random.seed(self.seed)
i = index + self.scene_indices[0]
scene = RenderedScene(i, self.data_dir, self.data_root_dir)
composite = scene.create_composite() #get empty composite
#Select a subset of objects randomly. Number of objects is uniformly
#distributed between [0, total_number_of_objects]
#Doors and windows do not count here
object_nodes = scene.object_nodes
num_objects = random.randint(0, len(object_nodes))
num_categories = len(scene.categories)
OUTSIDE = num_categories + 2
EXISTING = num_categories + 1
NOTHING = num_categories
centroids = []
existing_categories = torch.zeros(num_categories)
future_categories = torch.zeros(num_categories)
#Process existing objects
for i in range(num_objects):
#Add object to composite
node = object_nodes[i]
composite.add_node(node)
#Add existing centroids
xmin, _, ymin, _ = node["bbox_min"]
xmax, _, ymax, _ = node["bbox_max"]
centroids.append(((xmin + xmax) / 2, (ymin + ymax) / 2, EXISTING))
existing_categories[node["category"]] += 1
inputs = composite.get_composite(ablation=self.ablation)
size = inputs.shape[1]
#Process removed objects
for i in range(num_objects, len(object_nodes)):
node = object_nodes[i]
xmin, _, ymin, _ = node["bbox_min"]
xmax, _, ymax, _ = node["bbox_max"]
centroids.append(
((xmin + xmax) / 2, (ymin + ymax) / 2, node["category"]))
future_categories[node["category"]] += 1
resample = True
if random.uniform(0,
1) > self.p_auxiliary: #Sample an object at random
x, y, output_centroid = random.choice(centroids)
x = int(x)
y = int(y)
else: #Or sample an auxiliary category
while resample:
x, y = random.randint(0, 511), random.randint(
0, 511
) #Should probably remove this hardcoded size at somepoint
good = True
for (xc, yc, _) in centroids:
#We don't want to sample an empty space that's too close to a centroid
#That is, if it can fall within the ground truth attention mask of an object
if x - 4 < xc < x + 5 and y - 4 < yc < y + 5:
good = False
if good:
if not inputs[0][x][y]:
output_centroid = OUTSIDE #Outside of room
if random.uniform(
0, 1
) > 0.8: #Just some hardcoded stuff simple outside room is learned very easily
resample = False
else:
output_centroid = NOTHING #Inside of room
resample = False
#Attention mask
xmin = max(x - 4, 0)
xmax = min(x + 5, size)
ymin = max(y - 4, 0)
ymax = min(y + 5, size)
inputs[-1, xmin:xmax, ymin:ymax] = 1 #Create attention mask
#Compute weight for L_Global
#If some objects in this cateogory are removed, weight is zero
#Otherwise linearly scaled based on completeness of the room
#See paper for details
penalty = torch.zeros(num_categories)
penalty[future_categories == 0] = num_objects / len(object_nodes)
return inputs, output_centroid, existing_categories, penalty
def __getitem__(self, index):
if self.seed:
random.seed(self.seed)
i = index + self.scene_indices[0]
scene = RenderedScene(i, self.data_dir, self.data_root_dir)
composite = scene.create_composite() #get empty composite
object_nodes = scene.object_nodes
#Since we need to at least rotate one object, this differs from location dataset slightly
num_objects = random.randint(0, len(object_nodes) - 1)
num_categories = len(scene.categories)
for i in range(num_objects):
node = object_nodes[i]
composite.add_node(node)
#Select the node we want to rotate
node = object_nodes[num_objects]
modelId = node["modelId"]
#Just some made up distribution of different cases
#Focusing on 180 degree, then 90, then others
ran = random.uniform(0, 1)
if ran < 0.2:
r = math.pi
target = 0
elif ran < 0.4:
r = math.pi / 2 * random.randint(1, 3)
target = 0
elif ran < 0.6:
r = math.pi / 8 * random.randint(1, 15)
target = 0
else:
r = 0
target = 1
o = Obj(modelId)
#Get the transformation matrix from object space to scene space
t = RotationDataset.pgen.get_projection(scene.room).to_2d(
np.asarray(node["transform"]).reshape(4, 4))
#Since centered already in object space, rotating the object in object space is the easier option
sin, cos = math.sin(r), math.cos(r)
t_rot = np.asarray([[cos, 0, -sin, 0], \
[0, 1, 0, 0], \
[sin, 0, cos, 0], \
[0, 0, 0, 1]])
o.transform(np.dot(t_rot, t))
#Render the rotated view of the object
rotated = torch.from_numpy(
TopDownView.render_object_full_size(o, composite.size))
#Calculate the relevant info needed to composite it to the input
sin, cos = composite.get_transformation(node["transform"])
original_r = math.atan2(sin, cos)
sin = math.sin(original_r + r)
cos = math.cos(original_r + r)
composite.add_height_map(rotated, node["category"], sin, cos)
inputs = composite.get_composite(ablation=self.ablation)
#Add attention channel, which is just the outline of the targeted object
rotated[rotated > 0] = 1
inputs[-1] = rotated
return inputs, target