31 lines
1.1 KiB
Python
31 lines
1.1 KiB
Python
from __future__ import annotations
|
||
|
||
import io
|
||
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
|
||
def depth_to_png16_bytes(depth: np.ndarray, target_hw: tuple[int, int] | None = None) -> bytes:
|
||
depth = np.asarray(depth, dtype=np.float32).squeeze()
|
||
if depth.ndim != 2:
|
||
raise ValueError(f"depth 应为二维数组，实际 shape={depth.shape}")
|
||
if target_hw is not None:
|
||
th, tw = int(target_hw[0]), int(target_hw[1])
|
||
if th > 0 and tw > 0 and (depth.shape[0] != th or depth.shape[1] != tw):
|
||
im = Image.fromarray(depth, mode="F")
|
||
im = im.resize((tw, th), Image.Resampling.BILINEAR)
|
||
depth = np.asarray(im, dtype=np.float32)
|
||
dmin = float(depth.min())
|
||
dmax = float(depth.max())
|
||
if dmax > dmin:
|
||
norm = (depth - dmin) / (dmax - dmin)
|
||
else:
|
||
norm = np.zeros_like(depth, dtype=np.float32)
|
||
# 前后端约定：最远=255，最近=0（8-bit）
|
||
u8 = ((1.0 - norm) * 255.0).clip(0, 255).astype(np.uint8)
|
||
img = Image.fromarray(u8, mode="L")
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="PNG")
|
||
return buf.getvalue()
|