Created
April 11, 2018 16:02
-
-
Save N-McA/f8a34a46449994f8efbaef74c1536911 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import numpy as np | |
from pathlib import Path | |
import h5py | |
def resolve_name(f, name_ref): | |
return ''.join([chr(i) for i in f[name_ref]]) | |
def flat(x): | |
return np.array(x).flatten() | |
def get_object_ids_from_pointcloud_files(data_loc): | |
full_obj_names = [] | |
for area_n in [1, 2, 3, 4]: | |
filename = data_loc / 'area_{}/3d/pointcloud.mat'.format(area_n) | |
f = h5py.File(filename) | |
area_name = 'Area_{}'.format(area_n) | |
ds_name_refs = flat(f[area_name]['Disjoint_Space']['name']) | |
ds_object_refs = flat(f[area_name]['Disjoint_Space']['object']) | |
for ds_name_ref, ds_obj_ref in zip(ds_name_refs, ds_object_refs): | |
ds_name = resolve_name(f, ds_name_ref) | |
objs = f[ds_obj_ref] | |
name_refs = flat(objs['name']) | |
for i in range(len(name_refs)): | |
obj_name = resolve_name(f, name_refs[i]) | |
full_obj_name = '{}_{}_{}'.format( | |
area_name, ds_name, obj_name | |
).lower() | |
full_obj_names.append(full_obj_name) | |
full_obj_names = sorted(full_obj_names) | |
assert len(full_obj_names) == len(np.unique(full_obj_names)) | |
return full_obj_names | |
def load_points(area_n, obj_id_to_int): | |
points = [] | |
colors = [] | |
obj_ids = [] | |
filename = data_loc / 'area_{}/3d/pointcloud.mat'.format(area_n) | |
f = h5py.File(filename) | |
area_name = 'Area_{}'.format(area_n) | |
ds_name_refs = flat(f[area_name]['Disjoint_Space']['name']) | |
ds_object_refs = flat(f[area_name]['Disjoint_Space']['object']) | |
for ds_name_ref, ds_obj_ref in zip(ds_name_refs, ds_object_refs): | |
ds_name = resolve_name(f, ds_name_ref) | |
objs = f[ds_obj_ref] | |
name_refs = flat(objs['name']) | |
point_refs = flat(objs['points']) | |
colour_refs = flat(objs['RGB_color']) | |
for i in range(len(name_refs)): | |
obj_name = resolve_name(f, name_refs[i]) | |
full_obj_name = '{}_{}_{}'.format( | |
area_name, ds_name, obj_name | |
).lower() | |
obj_id_n = obj_id_to_int[full_obj_name] | |
colors.append( | |
np.array(f[colour_refs[i]], dtype=np.uint8).T | |
) | |
points.append( | |
np.array(f[point_refs[i]], dtype=np.float16).T | |
) | |
idx = object_id_to_int[full_obj_name] | |
obj_ids.append( | |
idx * np.ones(len(points[-1]), dtype=np.uint16) | |
) | |
points = np.vstack(points) | |
colors = np.vstack(colors) | |
obj_ids = np.concatenate(obj_ids) | |
return points, colors, obj_ids | |
def generate_and_save_object_ids(data_loc, pydata_loc): | |
object_ids = get_object_ids_from_pointcloud_files(data_loc) | |
with (pydata_loc / 'object_ids').open('w') as f: | |
for obj_id in object_ids: | |
print(obj_id, file=f) | |
def load_object_ids(pydata_loc): | |
with (pydata_loc / 'object_ids').open() as f: | |
r = [line.strip() for line in f] | |
return r | |
data_loc = Path('/local/sdd/nm583/stanford_2d3ds/') | |
pydata_loc = Path('/local/sdd/nm583/stanford_2d3ds/python_data') | |
generate_and_save_object_ids(data_loc, pydata_loc) | |
obj_ids = load_object_ids(pydata_loc) | |
object_id_to_int = {ob_id: i for i, ob_id in enumerate(obj_ids)} | |
for area_n in [1, 2, 3, 4]: | |
points, colors, obj_ids = load_points(area_n, object_id_to_int) | |
np.save(pydata_loc / 'area_{}.points.npy'.format(area_n), points) | |
np.save(pydata_loc / 'area_{}.colors.npy'.format(area_n), colors) | |
np.save(pydata_loc / 'area_{}.obj_ids.npy'.format(area_n), obj_ids) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment