convert_stanford_to_npy.py

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)