65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
from pathlib import Path
|
|
|
|
from .surface import write_normal_texture
|
|
|
|
|
|
def write_da3_depth(prediction, output_path):
|
|
import numpy as np
|
|
|
|
values = {"depth": np.round(prediction.depth, 8)}
|
|
for name in ("conf", "extrinsics", "intrinsics"):
|
|
value = getattr(prediction, name, None)
|
|
if value is not None:
|
|
values[name] = np.round(value, 2) if name == "conf" else value
|
|
output_path = Path(output_path)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
np.savez_compressed(output_path, **values)
|
|
return output_path
|
|
|
|
|
|
def write_moge_depth(output_directory, prediction):
|
|
import numpy as np
|
|
|
|
output_directory = Path(output_directory)
|
|
np.savez(
|
|
output_directory / "results.npz",
|
|
**prediction.archive_values(),
|
|
)
|
|
for index, (depth, normal, mask) in enumerate(
|
|
zip(prediction.depth, prediction.normal, prediction.mask)
|
|
):
|
|
_write_moge_maps(
|
|
output_directory,
|
|
None if len(prediction.depth) == 1 else index,
|
|
depth,
|
|
normal,
|
|
mask,
|
|
)
|
|
|
|
|
|
def _write_moge_maps(output_directory, index, depth, normal, mask):
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
suffix = "" if index is None else f".{index + 1:04d}"
|
|
valid = mask & np.isfinite(depth) & (depth > 0)
|
|
preview = np.zeros(depth.shape, dtype=np.uint8)
|
|
if np.any(valid):
|
|
low, high = np.percentile(depth[valid], (2.0, 98.0))
|
|
span = max(float(high - low), 1e-6)
|
|
preview[valid] = np.clip(
|
|
(depth[valid] - low) / span * 255.0,
|
|
0.0,
|
|
255.0,
|
|
).astype(np.uint8)
|
|
Image.fromarray(preview, mode="L").save(
|
|
output_directory / f"depth{suffix}.png"
|
|
)
|
|
write_normal_texture(
|
|
normal,
|
|
output_directory / f"normal{suffix}.png",
|
|
)
|
|
Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save(
|
|
output_directory / f"mask{suffix}.png"
|
|
)
|