Gökay Aydoğan
a9e20b7764
* feat: transform depth nodes into a complete v3 depth toolkit * fix(ci): validate metadata on Python 3.10
85 lines
3.2 KiB
Python
85 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import struct
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
import depth_nodes as nodes
|
|
|
|
|
|
def depth(height=24, width=32):
|
|
x = torch.linspace(0.05, 1, width).view(1, 1, width, 1)
|
|
return x.repeat(1, height, 1, 3)
|
|
|
|
|
|
def color(height=24, width=32):
|
|
image = torch.zeros((1, height, width, 3))
|
|
image[..., 0] = torch.linspace(0, 1, width)
|
|
image[..., 1] = torch.linspace(0, 1, height).view(height, 1)
|
|
image[..., 2] = 0.4
|
|
return image
|
|
|
|
|
|
def test_normalize_colormap_normal_and_masks():
|
|
normalized, report = nodes.DepthNormalize().normalize(depth(), "Percentile", 1, 99, False, 1)
|
|
colored, = nodes.DepthColormap().colorize(normalized, "turbo", False, True)
|
|
normal, = nodes.DepthToNormal().convert(normalized, 1.0, 60, "OpenGL (+Y)", False)
|
|
inside, outside, masked = nodes.DepthRangeMask().mask(normalized, 0.2, 0.8, 0.05, False)
|
|
assert normalized.min() == 0 and normalized.max() == 1
|
|
assert json.loads(report)["method"] == "Percentile"
|
|
assert colored.shape == normal.shape == masked.shape == (1, 24, 32, 3)
|
|
assert torch.allclose(inside + outside, torch.ones_like(inside), atol=1e-6)
|
|
|
|
|
|
def test_cleanup_and_stats_emit_inspectable_outputs():
|
|
source = depth()
|
|
source[:, 8:12, 10:15] = 0
|
|
cleaned, repaired, report = nodes.DepthCleanup().clean(source, 0.001, 8, 1, 0.8)
|
|
stats, histogram = nodes.DepthStats().analyze(cleaned, 32)
|
|
assert repaired[:, 8:12, 10:15].max() == 1
|
|
assert json.loads(report)["changed_fraction"] > 0
|
|
assert json.loads(stats)["batch"][0]["max"] <= 1
|
|
assert histogram.shape == (1, 300, 640, 3)
|
|
|
|
|
|
def test_median_filter_preserves_sub_8_bit_depth_precision():
|
|
values = np.array(
|
|
[[[0.5001, 0.5002, 0.5003], [0.5004, 0.5005, 0.5006], [0.5007, 0.5008, 0.5009]]],
|
|
dtype=np.float32,
|
|
)
|
|
filtered = nodes._median_filter(values, 1)
|
|
assert np.isclose(filtered[0, 1, 1], 0.5005, atol=1e-6)
|
|
assert not np.isclose(filtered[0, 1, 1] * 255, round(filtered[0, 1, 1] * 255))
|
|
|
|
|
|
def test_binary_ply_and_glb_exports(monkeypatch, tmp_path):
|
|
import folder_paths
|
|
|
|
monkeypatch.setattr(folder_paths, "get_output_directory", lambda: str(tmp_path), raising=False)
|
|
ply_paths, ply_manifest = nodes.DepthToPointCloud().export(
|
|
depth(), 60, 2, 4, True, "depth_exports/cloud", color()
|
|
)
|
|
glb_paths, glb_manifest = nodes.DepthToMesh().export(
|
|
depth(), 60, 2, 4, 100, "depth_exports/mesh", color()
|
|
)
|
|
ply = Path(ply_paths)
|
|
glb = Path(glb_paths)
|
|
assert ply.read_bytes().startswith(b"ply\nformat binary_little_endian")
|
|
magic, version, total = struct.unpack("<4sII", glb.read_bytes()[:12])
|
|
assert (magic, version, total) == (b"glTF", 2, glb.stat().st_size)
|
|
assert json.loads(ply_manifest)["point_counts"][0] > 0
|
|
assert json.loads(glb_manifest)["triangle_counts"][0] > 0
|
|
|
|
|
|
def test_parallax_frames_and_validity_masks():
|
|
frames, masks, manifest = nodes.DepthParallaxFrames().animate(
|
|
color(), depth(), 8, 12, "ellipse", True, 0.5
|
|
)
|
|
assert frames.shape == (8, 24, 32, 3)
|
|
assert masks.shape == (8, 24, 32)
|
|
assert json.loads(manifest)["total_frames"] == 8
|
|
assert masks.min() == 0 and masks.max() == 1
|