def __getitem__(self, idx):
file = self.files[idx]
name = None
if self.is_chunks:
name = os.path.splitext(os.path.basename(file))[0]
inputs, targets, dims, world2grid, target_known, target_hierarchy = data_util.load_train_file(file)
else:
input_file = file[0]
target_file = file[1]
name = os.path.splitext(os.path.basename(input_file))[0]
inputs, dims, world2grid = data_util.load_scene(input_file)
targets, dims, world2grid = data_util.load_scene(target_file)
target_known = data_util.load_scene_known(os.path.splitext(target_file)[0] + '.knw')
targets = data_util.sparse_to_dense_np(targets[0], targets[1], dims[2], dims[1], dims[0], -float('inf'))
target_hierarchy = None
orig_dims = torch.LongTensor(targets.shape)
if not self.is_chunks:
# add padding
hierarchy_factor = pow(2, self.num_hierarchy_levels-1)
max_input_dim = np.array(sdf.shape)
if self.max_input_height > 0 and max_input_dim[self.UP_AXIS] > self.max_input_height:
max_input_dim[self.UP_AXIS] = self.max_input_height
mask_input = input[0][:,self.UP_AXIS] < self.max_input_height
input[0] = input[0][mask_input]
input[1] = input[1][mask_input]
max_input_dim = ((max_input_dim + (hierarchy_factor*4) - 1) // (hierarchy_factor*4)) * (hierarchy_factor*4)
# pad target to max_input_dim
padded = np.zeros((max_input_dim[0], max_input_dim[1], max_input_dim[2]), dtype=np.float32)
padded.fill(-float('inf'))
padded[:min(self.max_input_height, sdf.shape[0]), :sdf.shape[1], :sdf.shape[2]] = sdf[:self.max_input_height, :, :]
sdf = padded
if target_known is not None:
known_pad = np.ones((max_input_dim[0], max_input_dim[1], max_input_dim[2]), dtype=np.uint8) * 255
known_pad[:min(self.max_input_height,target_known.shape[0]), :target_known.shape[1], :target_known.shape[2]] = target_known[:self.max_input_height, :, :]
target_known = known_pad
else:
if self.num_hierarchy_levels < 4:
target_hierarchy = target_hierarchy[4-self.num_hierarchy_levels:]
mask = np.abs(inputs[1]) < self.truncation
input_locs = inputs[0][mask]
input_vals = inputs[1][mask]
inputs = [torch.from_numpy(input_locs).long(), torch.from_numpy(input_vals[:,np.newaxis]).float()]
targets = targets[np.newaxis,:]
targets = torch.from_numpy(targets)
if target_hierarchy is not None:
for h in range(len(target_hierarchy)):
target_hierarchy[h] = torch.from_numpy(target_hierarchy[h][np.newaxis,:])
world2grid = torch.from_numpy(world2grid)
target_known = target_known[np.newaxis,:]
target_known = torch.from_numpy(target_known)
sample = {'name': name, 'input': inputs, 'sdf': targets, 'world2grid': world2grid, 'known': target_known, 'hierarchy': target_hierarchy, 'orig_dims': orig_dims}
return sample
def test(scene_name, eval_file):
print('scene', scene_name)
scene_file = os.path.join(opt.data_path_3d, scene_name + '.sdf.ann')
scene_image_file = os.path.join(opt.data_path_3d, scene_name + '.image')
if not os.path.exists(scene_file) or not os.path.exists(scene_image_file):
print(scene_file, os.path.exists(scene_file))
print(scene_image_file, os.path.exists(scene_image_file))
raise
scene_occ, scene_label = data_util.load_scene(scene_file, num_classes,
opt.has_gt)
if scene_occ.shape[1] > column_height:
scene_occ = scene_occ[:, :column_height, :, :]
if opt.has_gt:
scene_label = scene_label[:column_height, :, :]
scene_occ_sz = scene_occ.shape[1:]
depth_images, color_images, poses, frame_ids, world_to_grids = data_util.load_scene_image_info_multi(
scene_image_file, scene_name, opt.data_path_2d, proj_image_dims,
input_image_dims, num_classes, color_mean, color_std)
input_occ = torch.cuda.FloatTensor(1, 2, grid_dims[2], grid_dims[1],
grid_dims[0])
depth_image = torch.cuda.FloatTensor(num_images, proj_image_dims[1],
proj_image_dims[0])
color_image = torch.cuda.FloatTensor(num_images, 3, input_image_dims[1],
input_image_dims[0])
world_to_grid = torch.cuda.FloatTensor(num_images, 4, 4)
pose = torch.cuda.FloatTensor(num_images, 4, 4)
output_probs = np.zeros(
[num_classes, scene_occ_sz[0], scene_occ_sz[1], scene_occ_sz[2]])
# make sure nonsingular
for k in range(num_images):
pose[k] = torch.eye(4)
world_to_grid[k] = torch.eye(4)
# go thru all columns
for y in range(grid_padY, scene_occ_sz[1] - grid_padY):
for x in range(grid_padX, scene_occ_sz[2] - grid_padX):
input_occ.fill_(0)
input_occ[0, :, :scene_occ_sz[0], :, :] = torch.from_numpy(
scene_occ[:, :, y - grid_padY:y + grid_padY + 1,
x - grid_padX:x + grid_padX + 1])
cur_frame_ids = frame_ids[:, y, x][np.greater_equal(
frame_ids[:, y, x], 0)]
if len(cur_frame_ids) < num_images or torch.sum(
input_occ[0, 0, :, grid_padY, grid_padX]) == 0:
continue
for k in range(num_images):
depth_image[k] = depth_images[cur_frame_ids[k]]
color_image[k] = color_images[cur_frame_ids[k]]
pose[k] = poses[cur_frame_ids[k]]
world_to_grid[k] = torch.from_numpy(world_to_grids[y, x])
proj_mapping = [
projection.compute_projection(d, c, t)
for d, c, t in zip(depth_image, pose, world_to_grid)
]
if None in proj_mapping: #invalid sample
#print('(invalid sample)')
continue
proj_mapping = list(zip(*proj_mapping))
proj_ind_3d = torch.stack(proj_mapping[0])
proj_ind_2d = torch.stack(proj_mapping[1])
imageft_fixed = model2d_fixed(torch.autograd.Variable(color_image))
imageft = model2d_trainable(imageft_fixed)
out = model(torch.autograd.Variable(input_occ), imageft,
torch.autograd.Variable(proj_ind_3d),
torch.autograd.Variable(proj_ind_2d), grid_dims)
output = out.data[0].permute(1, 0)
output_probs[:, :, y, x] = output[:, :scene_occ_sz[0]]
sys.stdout.write('\r[ %d | %d ]' %
(y + 1, scene_occ_sz[1] - grid_padY))
sys.stdout.flush()
sys.stdout.write('\n')
pred_label = np.argmax(output_probs, axis=0)
mask = np.equal(scene_occ[0], 1)
pred_label[np.logical_not(mask)] = 0
util.write_array_to_file(
pred_label.astype(np.uint8),
os.path.join(opt.output_path, scene_name + '.bin'))
eval_scene = None
if opt.has_gt:
eval_scene = evaluate_prediction(scene_occ, scene_label, pred_label)
return eval_scene