depth_io.py

"""
Space efficient save and load of depth maps
"""

from ctypes import c_ubyte, c_float
import numpy as np
from imageio.v3 import imread, imwrite


def depth_encode(dmap):
    h, w = dmap.shape[:2]
    all_bytes = []
    for y in range(h):
        for x in range(w):
            d_val = dmap[y, x, ...]
            d_cubyte = (c_ubyte * 4).from_buffer(c_float(d_val))
            d_b = bytes(d_cubyte)
            d_np = np.frombuffer(d_b, dtype=np.ubyte)
            all_bytes += [d_np]
    all_bytes = np.concatenate(all_bytes, axis=0)
    all_bytes = all_bytes.reshape(h, w, 4)
    return all_bytes


def depth_decode(dmap_img):
    h, w = dmap_img.shape[:2]
    all_floats = []
    for y in range(h):
        for x in range(w):
            d_b = dmap_img[y, x, ...]
            d_b = d_b.tobytes()
            d_ubyte = (c_ubyte * 4).from_buffer_copy(d_b)
            d_float = (c_float).from_buffer(d_ubyte).value
            all_floats += [d_float]
    all_floats = np.array(all_floats, dtype=np.float32).reshape(h, w, 1)
    return all_floats

dmap = np.random.rand(200, 100, 1)
dmap_ubyte = depth_encode(dmap)
imwrite("depth.png", dmap_ubyte)
np.save("depth.npy", dmap)
dmap_ubyte2 = imread("depth.png")
dmap_float = depth_decode(dmap_ubyte2)
print(np.allclose(dmap, dmap_float))