def update_scannet_info_json(path, path_meta, test_only=False, verbose=2):
scenes = []
if not test_only:
scenes += sorted([
os.path.join('scans', scene)
for scene in os.listdir(os.path.join(path, 'scans'))
])
scenes += sorted([
os.path.join('scans_test', scene)
for scene in os.listdir(os.path.join(path, 'scans_test'))
])
for scene in scenes:
if verbose > 0:
print('update info json for %s' % scene)
info_file = os.path.join(path_meta, scene, 'info.json')
data = load_info_json(info_file)
folder, scene = scene.split('/')
data['path'] = path
data['file_name_mesh_gt'] = os.path.join(path, folder, scene,
scene + '_vh_clean_2.ply')
data['file_name_seg_indices'] = os.path.join(
path, folder, scene, scene + '_vh_clean_2.0.010000.segs.json')
data['file_name_seg_groups'] = os.path.join(
path, folder, scene, scene + '.aggregation.json')
frames = data['frames']
new_frames = []
for frame_id, frame in enumerate(frames):
frame['file_name_image'] = os.path.join(path, folder, scene,
'color',
'%d.jpg' % frame_id)
frame['file_name_depth'] = os.path.join(path, folder, scene,
'depth',
'%d.png' % frame_id)
if frame['file_name_instance'] != '':
frame['file_name_instance'] = os.path.join(
path, folder, scene, 'instance-filt', '%d.png' % frame_id)
new_frames.append(frame)
data['frames'] = new_frames
for voxel_size in [4, 8, 16]:
data['file_name_vol_%02d' % voxel_size] = os.path.join(
path_meta, folder, scene, 'tsdf_%02d.npz' % voxel_size)
json.dump(
data, open(os.path.join(path_meta, folder, scene, 'info.json'),
'w'))
def fuse_scene(path_meta,
scene,
voxel_size,
trunc_ratio=3,
max_depth=3,
vol_prcnt=.995,
vol_margin=1.5,
fuse_semseg=False,
device=0,
verbose=2):
""" Use TSDF fusion with GT depth maps to generate GT TSDFs
Args:
path_meta: path to save the TSDFs
(we recommend creating a parallel directory structure to save
derived data so that we don't modify the original dataset)
scene: name of scene to process
voxel_size: voxel size of TSDF
trunc_ratio: truncation distance in voxel units
max_depth: mask out large depth values since they are noisy
vol_prcnt: for computing the bounding volume of the TSDF... ignore outliers
vol_margin: padding for computing bounding volume of the TSDF
fuse_semseg: whether to accumulate semseg images for GT semseg
(prefered method is to not accumulate and insted transfer labels
from ground truth labeled mesh)
device: cpu/ which gpu
verbose: how much logging to print
Returns:
writes a TSDF (.npz) file into path_meta/scene
Notes: we use a conservative value of max_depth=3 to reduce noise in the
ground truth. However, this means some distant data is missing which can
create artifacts. Nevertheless, we found we acheived the best 2d metrics
with the less noisy ground truth.
"""
if verbose > 0:
print('fusing', scene, 'voxel size', voxel_size)
info_file = os.path.join(path_meta, scene, 'info.json')
# get gpu device for this worker
device = torch.device('cuda', device) # gpu for this process
# get the dataset
transform = transforms.Compose([
transforms.ResizeImage((640, 480)),
transforms.ToTensor(),
transforms.InstanceToSemseg('nyu40'),
transforms.IntrinsicsPoseToProjection(),
])
frame_types = ['depth', 'semseg'] if fuse_semseg else ['depth']
dataset = SceneDataset(info_file, transform, frame_types)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=None,
batch_sampler=None,
num_workers=4)
# find volume bounds and origin by backprojecting depth maps to point clouds
# use a subset of the frames to save time
if len(dataset) <= 200:
dataset1 = dataset
else:
inds = np.linspace(0, len(dataset) - 1, 200).astype(np.int)
dataset1 = torch.utils.data.Subset(dataset, inds)
dataloader1 = torch.utils.data.DataLoader(dataset1,
batch_size=None,
batch_sampler=None,
num_workers=4)
pts = []
for i, frame in enumerate(dataloader1):
projection = frame['projection'].to(device)
depth = frame['depth'].to(device)
depth[depth > max_depth] = 0
pts.append(depth_to_world(projection, depth).view(3, -1).T)
pts = torch.cat(pts)
pts = pts[torch.isfinite(pts[:, 0])].cpu().numpy()
# use top and bottom vol_prcnt of points plus vol_margin
origin = torch.as_tensor(
np.quantile(pts, 1 - vol_prcnt, axis=0) - vol_margin).float()
vol_max = torch.as_tensor(
np.quantile(pts, vol_prcnt, axis=0) + vol_margin).float()
vol_dim = ((vol_max - origin) / (float(voxel_size) / 100)).int().tolist()
# initialize tsdf
tsdf_fusion = TSDFFusion(vol_dim,
float(voxel_size) / 100,
origin,
trunc_ratio,
device,
label=fuse_semseg)
# integrate frames
for i, frame in enumerate(dataloader):
if verbose > 1 and i % 25 == 0:
print(scene, 'integrating voxel size', voxel_size, i, len(dataset))
projection = frame['projection'].to(device)
image = frame['image'].to(device)
depth = frame['depth'].to(device)
semseg = frame['semseg'].to(device) if fuse_semseg else None
# only use reliable depth
depth[depth > max_depth] = 0
tsdf_fusion.integrate(projection, depth, image, semseg)
# save mesh and tsdf
file_name_vol = os.path.join(path_meta, scene,
'tsdf_%02d.npz' % voxel_size)
file_name_mesh = os.path.join(path_meta, scene,
'mesh_%02d.ply' % voxel_size)
tsdf = tsdf_fusion.get_tsdf()
tsdf.save(file_name_vol)
mesh = tsdf.get_mesh()
mesh.export(file_name_mesh)
if fuse_semseg:
mesh = tsdf.get_mesh('instance')
mesh.export(file_name_mesh.replace('.ply', '_semseg.ply'))
# update info json
data = load_info_json(info_file)
data['file_name_vol_%02d' % voxel_size] = file_name_vol
json.dump(data, open(info_file, 'w'))
def label_scene(path_meta, scene, voxel_size, dist_thresh=.05, verbose=2):
""" Transfer instance labels from ground truth mesh to TSDF
For each voxel find the nearest vertex and transfer the label if
it is close enough to the voxel.
Args:
path_meta: path to save the TSDFs
(we recommend creating a parallel directory structure to save
derived data so that we don't modify the original dataset)
scene: name of scene to process
voxel_size: voxel size of TSDF to process
dist_thresh: beyond this distance labels are not transferd
verbose: how much logging to print
Returns:
Updates the TSDF (.npz) file with the instance volume
"""
# dist_thresh: beyond this distance to nearest gt mesh vertex,
# voxels are not labeled
if verbose > 0:
print('labeling', scene)
info_file = os.path.join(path_meta, scene, 'info.json')
data = load_info_json(info_file)
# each vertex in gt mesh indexs a seg group
segIndices = json.load(open(data['file_name_seg_indices'],
'r'))['segIndices']
# maps seg groups to instances
segGroups = json.load(open(data['file_name_seg_groups'], 'r'))['segGroups']
mapping = {
ind: group['id'] + 1
for group in segGroups for ind in group['segments']
}
# get per vertex instance ids (0 is unknown, [1,...] are objects)
n = len(segIndices)
instance_verts = torch.zeros(n, dtype=torch.long)
for i in range(n):
if segIndices[i] in mapping:
instance_verts[i] = mapping[segIndices[i]]
# load vertex locations
mesh = trimesh.load(data['file_name_mesh_gt'], process=False)
verts = mesh.vertices
# construct kdtree of vertices for fast nn lookup
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(verts)
kdtree = o3d.geometry.KDTreeFlann(pcd)
# load tsdf volume
tsdf = TSDF.load(data['file_name_vol_%02d' % voxel_size])
coords = coordinates(tsdf.tsdf_vol.size(), device=torch.device('cpu'))
coords = coords.type(torch.float) * tsdf.voxel_size + tsdf.origin.T
mask = tsdf.tsdf_vol.abs().view(-1) < 1
# transfer vertex instance ids to voxels near surface
instance_vol = torch.zeros(len(mask), dtype=torch.long)
for i in mask.nonzero():
_, inds, dist = kdtree.search_knn_vector_3d(coords[:, i], 1)
if dist[0] < dist_thresh:
instance_vol[i] = instance_verts[inds[0]]
tsdf.attribute_vols['instance'] = instance_vol.view(
list(tsdf.tsdf_vol.size()))
tsdf.save(data['file_name_vol_%02d' % voxel_size])
key = 'vol_%02d' % voxel_size
temp_data = {
key: tsdf,
'instances': data['instances'],
'dataset': data['dataset']
}
tsdf = transforms.InstanceToSemseg('nyu40')(temp_data)[key]
mesh = tsdf.get_mesh('semseg')
fname = data['file_name_vol_%02d' % voxel_size]
mesh.export(fname.replace('tsdf', 'mesh').replace('.npz', '_semseg.ply'))
def fuse_scene(path_meta, scene, info_file, dataset, voxel_size, device, origin, vol_max,
trunc_ratio=3, max_depth=3, fuse_semseg=False, verbose=2):
""" Use TSDF fusion with GT depth maps to generate GT TSDFs
Args:
path_meta: path to save the TSDFs
(we recommend creating a parallel directory structure to save
derived data so that we don't modify the original dataset)
scene: name of scene to process
voxel_size: voxel size of TSDF
trunc_ratio: truncation distance in voxel units
max_depth: mask out large depth values since they are noisy
vol_prcnt: for computing the bounding volume of the TSDF... ignore outliers
vol_margin: padding for computing bounding volume of the TSDF
fuse_semseg: whether to accumulate semseg images for GT semseg
(prefered method is to not accumulate and insted transfer labels
from ground truth labeled mesh)
device: cpu/ which gpu
verbose: how much logging to print
Returns:
writes a TSDF (.npz) file into path_meta/scene
Notes: we use a conservative value of max_depth=3 to reduce noise in the
ground truth. However, this means some distant data is missing which can
create artifacts. Nevertheless, we found we acheived the best 2d metrics
with the less noisy ground truth.
"""
if verbose>0:
print('fusing', scene, 'voxel size', voxel_size)
vol_dim = ((vol_max-origin)/(float(voxel_size)/100)).int().tolist()
# initialize tsdf
tsdf_fusion = TSDFFusion(vol_dim, float(voxel_size)/100, origin,
trunc_ratio, device, label=fuse_semseg)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=None,
batch_sampler=None, num_workers=4)
# integrate frames
for i, frame in enumerate(dataloader):
if verbose>1 and i%50==0:
print(scene, 'integrating voxel size', voxel_size, i, len(dataset))
projection = frame['projection'].to(device)
image = frame['image'].to(device)
depth = frame['depth'].to(device)
semseg = frame['semseg'].to(device) if fuse_semseg else None
# only use reliable depth
depth[depth>max_depth]=0
tsdf_fusion.integrate(projection, depth, image, semseg)
# save mesh and tsdf
file_name_vol = os.path.join(path_meta, scene, 'tsdf_%02d.npz'%voxel_size)
file_name_mesh = os.path.join(path_meta, scene, 'mesh_%02d.ply'%voxel_size)
tsdf = tsdf_fusion.get_tsdf()
tsdf.save(file_name_vol)
mesh = tsdf.get_mesh()
mesh.export(file_name_mesh)
if fuse_semseg:
mesh = tsdf.get_mesh('instance')
mesh.export(file_name_mesh.replace('.ply','_semseg.ply'))
# update info json
data = load_info_json(info_file)
data['file_name_vol_%02d'%voxel_size] = file_name_vol
json.dump(data, open(info_file, 'w'))
def process(info_file, pathout, stride, scale):
""" Run Colmap dense reconstruction with ground truth pose.
Copies and creates the necessary file structure required by Colmap.
Then runs Colmap.
Args:
info_file: path to info_json file for the scene
pathout: path to store intermediate and final results
stride: number of frames to skip (reduces runtime)
scale: how much to downsample images (reduces runtime and often
improves stereo matching results)
"""
info = load_info_json(info_file)
dataset = info['dataset']
scene = info['scene']
frames = info['frames'][::stride]
os.makedirs(os.path.join(pathout, dataset, scene, 'images'), exist_ok=True)
for i, frame in enumerate(frames):
if i % 25 == 0:
print(i, len(frames))
img = Image.open(frame['file_name_image'])
w = img.width // scale
h = img.height // scale
fname_out = os.path.split(frame['file_name_image'])[1]
fname_out = os.path.join(pathout, dataset, scene, 'images', fname_out)
img.resize((w, h), Image.BILINEAR).save(fname_out)
with open(os.path.join(pathout, dataset, scene, 'cameras.txt'), 'w') as fp:
fp.write('1 PINHOLE {w} {h} {fx} {fy} {cx} {cy}'.format(
w=w,
h=h,
fx=frames[0]['intrinsics'][0][0] / scale,
fy=frames[0]['intrinsics'][1][1] / scale,
cx=frames[0]['intrinsics'][0][2] / scale,
cy=frames[0]['intrinsics'][1][2] / scale,
))
with open(os.path.join(pathout, dataset, scene, 'points3D.txt'),
'w') as fp:
pass
cmd = 'colmap feature_extractor --database_path %s --image_path %s' % (
os.path.join(pathout, dataset, scene, 'database.db'),
os.path.join(pathout, dataset, scene, 'images'))
os.system(cmd)
cmd = 'colmap exhaustive_matcher --database_path %s' % (os.path.join(
pathout, dataset, scene, 'database.db'))
os.system(cmd)
conn = sqlite3.connect(os.path.join(pathout, dataset, scene,
'database.db'))
c = conn.cursor()
c.execute('SELECT image_id, name FROM images')
db_list = sorted(c.fetchall(), key=lambda x: x[1])
pose_dict = {
os.path.split(frame['file_name_image'])[1]: np.array(frame['pose'])
for frame in frames
}
with open(os.path.join(pathout, dataset, scene, 'images.txt'), 'w') as fp:
for ind, name in db_list:
pose = pose_dict[name]
pose = np.linalg.inv(pose)
q = R.from_matrix(pose[:3, :3]).as_quat()
t = pose[:3, 3]
fp.write(
'{i}, {qw}, {qx}, {qy}, {qz}, {tx}, {ty}, {tz}, 1, {name}\n\n'.
format(i=ind,
qw=q[3],
qx=q[0],
qy=q[1],
qz=q[2],
tx=t[0],
ty=t[1],
tz=t[2],
name=name))
cmd = ('colmap point_triangulator --database_path %s --image_path %s'
' --input_path %s --output_path %s') % (
os.path.join(pathout, dataset, scene, 'database.db'),
os.path.join(pathout, dataset, scene,
'images'), os.path.join(pathout, dataset, scene),
os.path.join(pathout, dataset, scene))
os.system(cmd)
cmd = 'colmap image_undistorter --image_path %s --input_path %s --output_path %s' % (
os.path.join(pathout, dataset, scene,
'images'), os.path.join(pathout, dataset, scene),
os.path.join(pathout, dataset, scene))
os.system(cmd)
cmd = 'colmap patch_match_stereo --workspace_path %s' % (os.path.join(
pathout, dataset, scene))
os.system(cmd)
cmd = 'colmap stereo_fusion --workspace_path %s --output_path %s' % (
os.path.join(pathout, dataset, scene),
os.path.join(pathout, dataset, scene, 'fused.ply'))
os.system(cmd)
cmd = ('colmap delaunay_mesher --input_path %s --output_path %s'
' --DelaunayMeshing.quality_regularization 5.'
' --DelaunayMeshing.max_proj_dist 10') % (os.path.join(
pathout, dataset,
scene), os.path.join(pathout, dataset, scene + '.ply'))
os.system(cmd)
def eval_scene(info_file, pathout):
""" Evaluates COLMAP inference compared to ground truth
Args:
info_file: path to info_json file for the scene
pathout: path where intermediate and final results are stored
"""
info = load_info_json(info_file)
dataset = info['dataset']
scene = info['scene']
frames = info['frames']
fnames = os.listdir(
os.path.join(pathout, dataset, scene, 'stereo', 'depth_maps'))
frames = [
frame for frame in frames
if os.path.split(frame['file_name_image'])[1] +
'.geometric.bin' in fnames
]
# 2d depth metrics
for i, frame in enumerate(frames):
if i % 25 == 0:
print(scene, i, len(fnames))
fname_trgt = frame['file_name_depth']
fname_pred = os.path.join(
pathout, dataset, scene, 'stereo', 'depth_maps',
os.path.split(frame['file_name_image'])[1] + '.geometric.bin')
depth_trgt = imageio.imread(fname_trgt).astype('float32') / 1000
depth_pred = read_array(fname_pred)
depth_pred[
depth_pred >
5] = 0 # ignore depth beyond 5 meters as it is probably wrong
depth_pred = resize(depth_pred, depth_trgt.shape)
temp = eval_depth(depth_pred, depth_trgt)
if i == 0:
metrics_depth = temp
else:
metrics_depth = {
key: value + temp[key]
for key, value in metrics_depth.items()
}
metrics_depth = {
key: value / len(frames)
for key, value in metrics_depth.items()
}
# 3d point metrics
fname_pred = os.path.join(pathout, dataset, scene, 'fused.ply')
fname_trgt = info['file_name_mesh_gt']
metrics_mesh = eval_mesh(fname_pred, fname_trgt)
metrics = {**metrics_depth, **metrics_mesh}
print(metrics)
rslt_file = os.path.join(pathout, dataset, scene, 'metrics.json')
json.dump(metrics, open(rslt_file, 'w'))
return metrics
