convert_mat_to_img.py 文件源码

def convert_nyu_16bit_image(path, target_dataset_directory):
    print("load dataset: %s" % (path))
    print("target dataset dir: %s" % (target_dataset_directory))

    f = h5py.File(path)

    for i, (image, depth) in enumerate(zip(f['images'], f['depths'])):

        ra_image = image.transpose(2, 1, 0)
        ra_depth = depth.transpose(1, 0) * 100

        re_depth = (ra_depth/np.max(ra_depth))*255.0

        image_pil = Image.fromarray(np.uint8(ra_image))
        depth_pil = Image.fromarray(np.uint8(re_depth))
        image_name = os.path.join(target_dataset_directory, "%05d_img.png" % (i))
        image_pil.save(image_name)
        depth_name = os.path.join(target_dataset_directory, "%05d_dep.png" % (i))
        depth_pil.save(depth_name)

        depth_meters_name = os.path.join(target_dataset_directory, "%05d_dep_meters.png" % (i))
        with open(depth_meters_name, 'wb') as f:
            writer = png.Writer(width=ra_depth.shape[1], height=ra_depth.shape[0], bitdepth=16, greyscale=True)
            zgray2list = ra_depth.tolist()
            writer.write(f, zgray2list)

        print(i)