parser = argparse.ArgumentParser('Generate 2D bbox GroundTruth from frame-level instance segmentation annotation')
parser.add_argument('--data_dir', type=str, default='/mnt/Data/Datasets/ScanNet_v2/scans/',
help='The path to annotations')
parser.add_argument('--frame_skip', type=int, default=15,
help='the number of frames to skip in extracting instance annotation images')
parser.add_argument('--scene_name', type=str, default=None, help='specific scene name to process')
parser.add_argument('--min_ratio', type=int, default=10, help='The minimum ratio between the object dimension '
'and the image dimension to filter objects')
parser.add_argument('--to_visu', type=bool, default=True, help='Visualize 2D bboxes')
parser.add_argument('--to_save', type=bool, default=True, help='Save the 2D bbox into files..')
opt = parser.parse_args()
# map original class name into nyu40 class ids, and extract 18 target classes
LABEL_MAP_FILE = '/mnt/Data/Datasets/ScanNet_v2/scannetv2-labels.combined.tsv'
LABEL_MAP = scannet_utils.read_label_mapping(LABEL_MAP_FILE, label_from='raw_category', label_to='nyu40class')
TARGET_CLASS_NAMES = cfg.SCANNET.CLASSES
if opt.scene_name is None:
SCAN_NAMES = [line.rstrip() for line in open('/mnt/Data/Datasets/ScanNet_v1/sceneid_sort.txt')]
else:
SCAN_NAMES = [opt.scene_name]
for scan_id, scan_name in enumerate(SCAN_NAMES):
print('====== Process {0}-th scan [{1}] ======'.format(scan_id, scan_name))
scan_path = os.path.join(opt.data_dir, scan_name)
if not os.path.exists(os.path.join(scan_path, 'instance-filt')): unzip_instace_file(scan_path, scan_name)
objectID2label, colour_code = get_objectID2label_and_color(scan_path, scan_name)
frame_names = os.listdir(os.path.join(scan_path, 'color'))
def collect_point_data(scene_name):
# read label mapping file
label_map = scannet_utils.read_label_mapping(opt.label_map_file,
label_from='raw_category',
label_to='nyu40id')
# Over-segmented segments: maps from segment to vertex/point IDs
data_folder = os.path.join(opt.scannet_path, scene_name)
# Read segmentation label
seg_filename = os.path.join(
data_folder, '%s_vh_clean_2.0.010000.segs.json' % (scene_name))
seg_to_verts, num_verts = scannet_utils.read_segmentation(seg_filename)
# Read Instances segmentation label
agg_filename = os.path.join(data_folder,
'%s.aggregation.json' % (scene_name))
object_id_to_segs, label_to_segs = scannet_utils.read_aggregation(
agg_filename)
# Raw points in XYZRGBA
ply_filename = os.path.join(data_folder,
'%s_vh_clean_2.ply' % (scene_name))
#points = pc_util.read_ply_rgba(ply_filename)
points = pc_utils.read_ply_rgba_normal(ply_filename)
label_ids = np.zeros(shape=(num_verts), dtype=np.uint32) # 0: unannotated
for label, segs in label_to_segs.items():
# convert scannet raw label to nyu40 label (1~40), 0 for unannotated, 41 for unknown
label_id = label_map[label]
# only evaluate 20 class in nyu40 label
# map nyu40 to 1~21, 0 for unannotated, unknown and not evalutated
if label_id in g_label_ids: # IDS for 20 classes in nyu40 for evaluation (1~21)
eval_label_id = g_label_ids.index(label_id)
else: # IDS unannotated, unknow or not for evaluation go to unannotate label (0)
eval_label_id = g_label_names.index('unannotate')
for seg in segs:
verts = seg_to_verts[seg]
label_ids[verts] = eval_label_id
#for i in range(20):
# print(label_ids[i])
instance_ids = np.zeros(shape=(num_verts),
dtype=np.uint32) # 0: unannotated
for object_id, segs in object_id_to_segs.items():
for seg in segs:
verts = seg_to_verts[seg]
instance_ids[verts] = object_id
points = np.delete(points, 6, 1) # only RGB, ignoring A
label_ids = np.expand_dims(label_ids, 1)
instance_ids = np.expand_dims(instance_ids, 1)
#print(points.shape, label_ids.shape, instance_ids.shape)
# order is critical, do not change the order
data = np.concatenate((points, instance_ids, label_ids), 1)
#print(data.shape)
#for i in range(20):
# print(data[i, 10])
out_filename = os.path.join(data_folder, scene_name +
'.npy') # scene0000_00/scene0000_00.npy
np.save(out_filename, data)
print(scene_name, ' points shape:', data.shape)
print(opt)
'''
tools for export and conver label
'''
label_map = None
g_label_names = None
g_label_ids = None
if opt.export_label_images:
try:
sys.path.append('../utils')
import scannet_utils
except:
print('Failed to import ScanNet code toolbox util')
sys.exit(-1)
label_map = scannet_utils.read_label_mapping(opt.label_map_file,
label_from='id',
label_to='nyu40id')
g_label_names = scannet_utils.g_label_names
g_label_ids = scannet_utils.g_label_ids
'''
tools for export .sens
'''
try:
from SensorData import SensorData
except:
print('Failed to import SensorData (from ScanNet code toolbox)')
sys.exit(-1)
def print_error(message):
sys.stderr.write('ERROR: ' + str(message) + '\n')
def export(mesh_file,
agg_file,
seg_file,
meta_file,
label_map_file,
output_file=None):
"""Export original files to vert, ins_label, sem_label and bbox file.
Args:
mesh_file (str): Path of the mesh_file.
agg_file (str): Path of the agg_file.
seg_file (str): Path of the seg_file.
meta_file (str): Path of the meta_file.
label_map_file (str): Path of the label_map_file.
output_file (str): Path of the output folder.
Default: None.
It returns a tuple, which containts the the following things:
np.ndarray: Vertices of points data.
np.ndarray: Indexes of label.
np.ndarray: Indexes of instance.
np.ndarray: Instance bboxes.
dict: Map from object_id to label_id.
"""
label_map = scannet_utils.read_label_mapping(label_map_file,
label_from='raw_category',
label_to='nyu40id')
mesh_vertices = scannet_utils.read_mesh_vertices_rgb(mesh_file)
# Load scene axis alignment matrix
lines = open(meta_file).readlines()
for line in lines:
if 'axisAlignment' in line:
axis_align_matrix = [
float(x)
for x in line.rstrip().strip('axisAlignment = ').split(' ')
]
break
axis_align_matrix = np.array(axis_align_matrix).reshape((4, 4))
pts = np.ones((mesh_vertices.shape[0], 4))
pts[:, 0:3] = mesh_vertices[:, 0:3]
pts = np.dot(pts, axis_align_matrix.transpose()) # Nx4
mesh_vertices[:, 0:3] = pts[:, 0:3]
# Load semantic and instance labels
object_id_to_segs, label_to_segs = read_aggregation(agg_file)
seg_to_verts, num_verts = read_segmentation(seg_file)
label_ids = np.zeros(shape=(num_verts), dtype=np.uint32)
object_id_to_label_id = {}
for label, segs in label_to_segs.items():
label_id = label_map[label]
for seg in segs:
verts = seg_to_verts[seg]
label_ids[verts] = label_id
instance_ids = np.zeros(shape=(num_verts),
dtype=np.uint32) # 0: unannotated
num_instances = len(np.unique(list(object_id_to_segs.keys())))
for object_id, segs in object_id_to_segs.items():
for seg in segs:
verts = seg_to_verts[seg]
instance_ids[verts] = object_id
if object_id not in object_id_to_label_id:
object_id_to_label_id[object_id] = label_ids[verts][0]
instance_bboxes = np.zeros((num_instances, 7))
for obj_id in object_id_to_segs:
label_id = object_id_to_label_id[obj_id]
obj_pc = mesh_vertices[instance_ids == obj_id, 0:3]
if len(obj_pc) == 0:
continue
xmin = np.min(obj_pc[:, 0])
ymin = np.min(obj_pc[:, 1])
zmin = np.min(obj_pc[:, 2])
xmax = np.max(obj_pc[:, 0])
ymax = np.max(obj_pc[:, 1])
zmax = np.max(obj_pc[:, 2])
bbox = np.array([(xmin + xmax) / 2, (ymin + ymax) / 2,
(zmin + zmax) / 2, xmax - xmin, ymax - ymin,
zmax - zmin, label_id])
# NOTE: this assumes obj_id is in 1,2,3,.,,,.NUM_INSTANCES
instance_bboxes[obj_id - 1, :] = bbox
if output_file is not None:
np.save(output_file + '_vert.npy', mesh_vertices)
np.save(output_file + '_sem_label.npy', label_ids)
np.save(output_file + '_ins_label.npy', instance_ids)
np.save(output_file + '_bbox.npy', instance_bboxes)
return mesh_vertices, label_ids, instance_ids,\
instance_bboxes, object_id_to_label_id
def main(opt):
def log_string(out_str):
LOG_FOUT.write(out_str + '\n')
LOG_FOUT.flush()
print(out_str)
def select_frames_for_one_scan(scan_path, objectID2label):
# Round1: check invalid pose file (discard those frames contain inf) e.g., scene0067_02/pose/133.txx
num_frames = len(os.listdir(os.path.join(scan_path, 'color')))
selected_framenames_r1 = []
for i in range(num_frames):
if find_valid_pose(scan_path, i):
selected_framenames_r1.append(i)
log_string('\tAfter filtering invalid pose files, remain {0} [origin {1}] frames'.format(
len(selected_framenames_r1), num_frames))
# Round2: check the 2D segmentation annotation (discard those frames without target classes)
selected_framenames_r2 = []
for frame_name in selected_framenames_r1:
if find_valid_objects(scan_path, frame_name, objectID2label, LABEL_MAP, TARGET_CLASS_NAMES):
selected_framenames_r2.append(frame_name)
log_string('\tAfter filtering frames without the interested objects, remain {0} [origin {1}] frames'.format(
len(selected_framenames_r2), len(selected_framenames_r1)))
# Round3: check the camera rotation angle between two consecutive frames
# (discard those frames with small view change)
selected_framenames_r3 = find_valid_rotation_angle(selected_framenames_r2, scan_path,
opt.min_angle, opt.visu_validRot)
log_string('\tAfter filtering frames with small viewpoint change, remain {0} [origin {1}] frames'.format(
len(selected_framenames_r3), len(selected_framenames_r2)))
return selected_framenames_r3
LOG_FOUT = open(os.path.join(opt.data_dir, 'log_process_data.txt'), 'w')
LOG_FOUT.write(str(opt) + '\n')
# map original class name into nyu40 class ids, and extract target classes
LABEL_MAP_FILE = os.path.join(opt.data_dir, 'scannetv2-labels.combined.tsv')
LABEL_MAP = scannet_utils.read_label_mapping(LABEL_MAP_FILE, label_from='raw_category', label_to='nyu40class')
TARGET_CLASS_NAMES = cfg.SCANNET.CLASSES
if opt.scene_name is None:
SCAN_NAMES = [line.rstrip() for line in open('/mnt/Data/Datasets/ScanNet_v1/sceneid_sort.txt')]
# import random
# random.shuffle(SCAN_NAMES)
else:
SCAN_NAMES = [opt.scene_name]
valid_scannames = []
for scan_id, scan_name in enumerate(SCAN_NAMES):
log_string('\n====== Process {0}-th scan [{1}] ======'.format(scan_id, scan_name))
scan_path = os.path.join(opt.data_dir, 'scans', scan_name)
if not os.path.exists(os.path.join(scan_path, 'instance-filt')): unzip_instace_file(scan_path, scan_name)
objectID2label, colour_code = get_objectID2label_and_color(scan_path, scan_name)
sel_framenames = select_frames_for_one_scan(scan_path, objectID2label)
if len(sel_framenames) >= opt.min_frames:
# During 2D bboxes extraction, the boxes whose dimension ratio to the image dimension is less than
# min_ratio will be removed. This may lead to invalid without bbox..
valid_framenames = extract_bbox2d_for_one_scan(scan_path, scan_name, sel_framenames, objectID2label,
LABEL_MAP, TARGET_CLASS_NAMES, colour_code,
opt.min_ratio, opt.visu_box2d, opt.save_box2d)
log_string('\tAfter filtering objects by checking its dimension ratio to the image dimension, '
'remain {0} [origin {1}] frames'.format(len(valid_framenames), len(sel_framenames)))
if len(valid_framenames) > opt.min_frames:
valid_scannames.append(scan_name)
log_string('=======================================================')
log_string('{0} scans are kept after processing...'.format(len(valid_scannames)))
# save the valid scans..
with open(os.path.join(opt.data_dir, 'sceneid_valid.txt'), 'w') as f:
for scanname in valid_scannames:
f.write('%s\n' %scanname)
LOG_FOUT.close()
def export(mesh_file,
agg_file,
seg_file,
meta_file,
label_map_file,
output_file=None):
""" points are XYZ RGB (RGB in 0-255),
semantic label as nyu40 ids,
instance label as 1-#instance,
box as (cx,cy,cz,dx,dy,dz,semantic_label)
"""
label_map = scannet_utils.read_label_mapping(label_map_file,
label_from='raw_category',
label_to='nyu40id')
# mesh_vertices = scannet_utils.read_mesh_vertices_rgb(mesh_file)
mesh_vertices = scannet_utils.read_mesh_vertices_rgb_normal(mesh_file)
# Load scene axis alignment matrix
lines = open(meta_file).readlines()
axis_align_matrix = None
for line in lines:
if 'axisAlignment' in line:
axis_align_matrix = [
float(x)
for x in line.rstrip().strip('axisAlignment = ').split(' ')
]
if axis_align_matrix != None:
axis_align_matrix = np.array(axis_align_matrix).reshape((4, 4))
pts = np.ones((mesh_vertices.shape[0], 4))
pts[:, 0:3] = mesh_vertices[:, 0:3]
pts = np.dot(pts, axis_align_matrix.transpose()) # Nx4
aligned_vertices = np.copy(mesh_vertices)
aligned_vertices[:, 0:3] = pts[:, 0:3]
else:
print("No axis alignment matrix found")
aligned_vertices = mesh_vertices
# Load semantic and instance labels
if os.path.isfile(agg_file):
object_id_to_segs, label_to_segs = read_aggregation(agg_file)
seg_to_verts, num_verts = read_segmentation(seg_file)
label_ids = np.zeros(shape=(num_verts),
dtype=np.uint32) # 0: unannotated
object_id_to_label_id = {}
for label, segs in label_to_segs.items():
label_id = label_map[label]
for seg in segs:
verts = seg_to_verts[seg]
label_ids[verts] = label_id
instance_ids = np.zeros(shape=(num_verts),
dtype=np.uint32) # 0: unannotated
num_instances = len(np.unique(list(object_id_to_segs.keys())))
for object_id, segs in object_id_to_segs.items():
for seg in segs:
verts = seg_to_verts[seg]
instance_ids[verts] = object_id
if object_id not in object_id_to_label_id:
object_id_to_label_id[object_id] = label_ids[verts][0]
instance_bboxes = np.zeros(
(num_instances, 8)) # also include object id
aligned_instance_bboxes = np.zeros(
(num_instances, 8)) # also include object id
for obj_id in object_id_to_segs:
label_id = object_id_to_label_id[obj_id]
# bboxes in the original meshes
obj_pc = mesh_vertices[instance_ids == obj_id, 0:3]
if len(obj_pc) == 0: continue
# Compute axis aligned box
# An axis aligned bounding box is parameterized by
# (cx,cy,cz) and (dx,dy,dz) and label id
# where (cx,cy,cz) is the center point of the box,
# dx is the x-axis length of the box.
xmin = np.min(obj_pc[:, 0])
ymin = np.min(obj_pc[:, 1])
zmin = np.min(obj_pc[:, 2])
xmax = np.max(obj_pc[:, 0])
ymax = np.max(obj_pc[:, 1])
zmax = np.max(obj_pc[:, 2])
bbox = np.array([(xmin + xmax) / 2, (ymin + ymax) / 2,
(zmin + zmax) / 2, xmax - xmin, ymax - ymin,
zmax - zmin, label_id,
obj_id - 1]) # also include object id
# NOTE: this assumes obj_id is in 1,2,3,.,,,.NUM_INSTANCES
instance_bboxes[obj_id - 1, :] = bbox
# bboxes in the aligned meshes
obj_pc = aligned_vertices[instance_ids == obj_id, 0:3]
if len(obj_pc) == 0: continue
# Compute axis aligned box
# An axis aligned bounding box is parameterized by
# (cx,cy,cz) and (dx,dy,dz) and label id
# where (cx,cy,cz) is the center point of the box,
# dx is the x-axis length of the box.
xmin = np.min(obj_pc[:, 0])
ymin = np.min(obj_pc[:, 1])
zmin = np.min(obj_pc[:, 2])
xmax = np.max(obj_pc[:, 0])
ymax = np.max(obj_pc[:, 1])
zmax = np.max(obj_pc[:, 2])
bbox = np.array([(xmin + xmax) / 2, (ymin + ymax) / 2,
(zmin + zmax) / 2, xmax - xmin, ymax - ymin,
zmax - zmin, label_id,
obj_id - 1]) # also include object id
# NOTE: this assumes obj_id is in 1,2,3,.,,,.NUM_INSTANCES
aligned_instance_bboxes[obj_id - 1, :] = bbox
else:
# use zero as placeholders for the test scene
print("use placeholders")
num_verts = mesh_vertices.shape[0]
label_ids = np.zeros(shape=(num_verts),
dtype=np.uint32) # 0: unannotated
instance_ids = np.zeros(shape=(num_verts),
dtype=np.uint32) # 0: unannotated
instance_bboxes = np.zeros((1, 8)) # also include object id
aligned_instance_bboxes = np.zeros((1, 8)) # also include object id
if output_file is not None:
np.save(output_file + '_vert.npy', mesh_vertices)
np.save(output_file + '_aligned_vert.npy', aligned_vertices)
np.save(output_file + '_sem_label.npy', label_ids)
np.save(output_file + '_ins_label.npy', instance_ids)
np.save(output_file + '_bbox.npy', instance_bboxes)
np.save(output_file + '_aligned_bbox.npy', instance_bboxes)
return mesh_vertices, aligned_vertices, label_ids, instance_ids, instance_bboxes, aligned_instance_bboxes
def export(mesh_file,
agg_file,
seg_file,
meta_file,
label_map_file,
output_file=None,
test_mode=False):
"""Export original files to vert, ins_label, sem_label and bbox file.
Args:
mesh_file (str): Path of the mesh_file.
agg_file (str): Path of the agg_file.
seg_file (str): Path of the seg_file.
meta_file (str): Path of the meta_file.
label_map_file (str): Path of the label_map_file.
output_file (str): Path of the output folder.
Default: None.
test_mode (bool): Whether is generating test data without labels.
Default: False.
It returns a tuple, which containts the the following things:
np.ndarray: Vertices of points data.
np.ndarray: Indexes of label.
np.ndarray: Indexes of instance.
np.ndarray: Instance bboxes.
dict: Map from object_id to label_id.
"""
label_map = scannet_utils.read_label_mapping(label_map_file,
label_from='raw_category',
label_to='nyu40id')
mesh_vertices = scannet_utils.read_mesh_vertices_rgb(mesh_file)
# Load scene axis alignment matrix
lines = open(meta_file).readlines()
# test set data doesn't have align_matrix
axis_align_matrix = np.eye(4)
for line in lines:
if 'axisAlignment' in line:
axis_align_matrix = [
float(x)
for x in line.rstrip().strip('axisAlignment = ').split(' ')
]
break
axis_align_matrix = np.array(axis_align_matrix).reshape((4, 4))
# perform global alignment of mesh vertices
pts = np.ones((mesh_vertices.shape[0], 4))
pts[:, 0:3] = mesh_vertices[:, 0:3]
pts = np.dot(pts, axis_align_matrix.transpose()) # Nx4
aligned_mesh_vertices = np.concatenate([pts[:, 0:3], mesh_vertices[:, 3:]],
axis=1)
# Load semantic and instance labels
if not test_mode:
object_id_to_segs, label_to_segs = read_aggregation(agg_file)
seg_to_verts, num_verts = read_segmentation(seg_file)
label_ids = np.zeros(shape=(num_verts), dtype=np.uint32)
object_id_to_label_id = {}
for label, segs in label_to_segs.items():
label_id = label_map[label]
for seg in segs:
verts = seg_to_verts[seg]
label_ids[verts] = label_id
instance_ids = np.zeros(shape=(num_verts),
dtype=np.uint32) # 0: unannotated
for object_id, segs in object_id_to_segs.items():
for seg in segs:
verts = seg_to_verts[seg]
instance_ids[verts] = object_id
if object_id not in object_id_to_label_id:
object_id_to_label_id[object_id] = label_ids[verts][0]
unaligned_bboxes = extract_bbox(mesh_vertices, object_id_to_segs,
object_id_to_label_id, instance_ids)
aligned_bboxes = extract_bbox(aligned_mesh_vertices, object_id_to_segs,
object_id_to_label_id, instance_ids)
else:
label_ids = None
instance_ids = None
unaligned_bboxes = None
aligned_bboxes = None
object_id_to_label_id = None
if output_file is not None:
np.save(output_file + '_vert.npy', mesh_vertices)
if not test_mode:
np.save(output_file + '_sem_label.npy', label_ids)
np.save(output_file + '_ins_label.npy', instance_ids)
np.save(output_file + '_unaligned_bbox.npy', unaligned_bboxes)
np.save(output_file + '_aligned_bbox.npy', aligned_bboxes)
np.save(output_file + '_axis_align_matrix.npy', axis_align_matrix)
return mesh_vertices, label_ids, instance_ids, unaligned_bboxes, \
aligned_bboxes, object_id_to_label_id, axis_align_matrix