mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-05-25 00:17:23 +08:00
Move turbo to colormap.py and add it to DepthAnything3Render.
This commit is contained in:
parent
83e2321d50
commit
f99c526a33
25
comfy/ldm/colormap.py
Normal file
25
comfy/ldm/colormap.py
Normal file
@ -0,0 +1,25 @@
|
|||||||
|
"""Colormap utilities for depth and geometry visualisation."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
|
||||||
|
def turbo(x: torch.Tensor) -> torch.Tensor:
|
||||||
|
"""Anton Mikhailov polynomial approximation of the Turbo colormap.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
x: Float tensor with values in [0, 1].
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
RGB tensor of the same shape as ``x`` with a trailing size-3 dimension.
|
||||||
|
"""
|
||||||
|
x = x.clamp(0.0, 1.0)
|
||||||
|
x2 = x * x
|
||||||
|
x3 = x2 * x
|
||||||
|
x4 = x2 * x2
|
||||||
|
x5 = x4 * x
|
||||||
|
r = 0.13572138 + 4.61539260*x - 42.66032258*x2 + 132.13108234*x3 - 152.94239396*x4 + 59.28637943*x5
|
||||||
|
g = 0.09140261 + 2.19418839*x + 4.84296658*x2 - 14.18503333*x3 + 4.27729857*x4 + 2.82956604*x5
|
||||||
|
b = 0.10667330 + 12.64194608*x - 60.58204836*x2 + 110.36276771*x3 - 89.90310912*x4 + 27.34824973*x5
|
||||||
|
return torch.stack([r, g, b], dim=-1).clamp(0.0, 1.0)
|
||||||
@ -33,6 +33,7 @@ import torch
|
|||||||
import comfy.model_management as mm
|
import comfy.model_management as mm
|
||||||
import comfy.sd
|
import comfy.sd
|
||||||
import folder_paths
|
import folder_paths
|
||||||
|
from comfy.ldm.colormap import turbo as _turbo
|
||||||
from comfy.ldm.depth_anything_3 import preprocess as da3_preprocess
|
from comfy.ldm.depth_anything_3 import preprocess as da3_preprocess
|
||||||
from comfy_api.latest import ComfyExtension, io
|
from comfy_api.latest import ComfyExtension, io
|
||||||
|
|
||||||
@ -150,7 +151,7 @@ class DepthAnything3Inference(io.ComfyNode):
|
|||||||
@classmethod
|
@classmethod
|
||||||
def define_schema(cls):
|
def define_schema(cls):
|
||||||
return io.Schema(
|
return io.Schema(
|
||||||
node_id="DepthAnything3",
|
node_id="DepthAnything3Inference",
|
||||||
search_aliases=["depth", "geometry", "da3", "depth anything", "monocular", "pointmap", "sky", "3d", "metric depth", "disparity"],
|
search_aliases=["depth", "geometry", "da3", "depth anything", "monocular", "pointmap", "sky", "3d", "metric depth", "disparity"],
|
||||||
display_name="Run Depth Anything 3",
|
display_name="Run Depth Anything 3",
|
||||||
category="image/geometry_estimation",
|
category="image/geometry_estimation",
|
||||||
@ -321,6 +322,8 @@ class DepthAnything3Inference(io.ComfyNode):
|
|||||||
return io.NodeOutput(geometry)
|
return io.NodeOutput(geometry)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
class DepthAnything3Render(io.ComfyNode):
|
class DepthAnything3Render(io.ComfyNode):
|
||||||
"""Visualise a DA3_GEOMETRY packet as a single image.
|
"""Visualise a DA3_GEOMETRY packet as a single image.
|
||||||
|
|
||||||
@ -328,6 +331,19 @@ class DepthAnything3Render(io.ComfyNode):
|
|||||||
Use multiple nodes in parallel to get depth + sky + confidence simultaneously.
|
Use multiple nodes in parallel to get depth + sky + confidence simultaneously.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
_DEPTH_RENDER_INPUTS = [
|
||||||
|
io.Combo.Input("normalization",
|
||||||
|
options=["v2_style", "min_max", "raw"],
|
||||||
|
default="v2_style",
|
||||||
|
tooltip="'v2_style': mean/std normalisation for perceptually balanced results (default). "
|
||||||
|
"'min_max': stretches the full depth range to [0, 1] for maximum contrast. "
|
||||||
|
"'raw': no scaling — preserves metric units for DA3-Metric-Large."),
|
||||||
|
io.Boolean.Input("apply_sky_clip", default=False,
|
||||||
|
tooltip="Clip sky-region depth to the 99th percentile of foreground depth before "
|
||||||
|
"normalisation. Requires a 'sky' tensor in the geometry "
|
||||||
|
"(DA3-Mono-Large or DA3-Metric-Large); raises an error otherwise."),
|
||||||
|
]
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def define_schema(cls):
|
def define_schema(cls):
|
||||||
return io.Schema(
|
return io.Schema(
|
||||||
@ -338,22 +354,13 @@ class DepthAnything3Render(io.ComfyNode):
|
|||||||
inputs=[
|
inputs=[
|
||||||
DA3Geometry.Input("geometry"),
|
DA3Geometry.Input("geometry"),
|
||||||
io.DynamicCombo.Input("output",
|
io.DynamicCombo.Input("output",
|
||||||
tooltip="depth: normalised depth image. "
|
tooltip="depth: normalised greyscale depth image. "
|
||||||
|
"depth_colored: depth mapped through the Turbo colormap. "
|
||||||
"sky_mask: sky probability in [0, 1] (Mono/Metric variants only). "
|
"sky_mask: sky probability in [0, 1] (Mono/Metric variants only). "
|
||||||
"confidence: normalised depth confidence (Small/Base variants only).",
|
"confidence: normalised depth confidence (Small/Base variants only).",
|
||||||
options=[
|
options=[
|
||||||
io.DynamicCombo.Option("depth", [
|
io.DynamicCombo.Option("depth", cls._DEPTH_RENDER_INPUTS),
|
||||||
io.Combo.Input("normalization",
|
io.DynamicCombo.Option("depth_colored", cls._DEPTH_RENDER_INPUTS),
|
||||||
options=["v2_style", "min_max", "raw"],
|
|
||||||
default="v2_style",
|
|
||||||
tooltip="'v2_style': mean/std normalisation for perceptually balanced results (default). "
|
|
||||||
"'min_max': stretches the full depth range to [0, 1] for maximum contrast. "
|
|
||||||
"'raw': no scaling — preserves metric units for DA3-Metric-Large."),
|
|
||||||
io.Boolean.Input("apply_sky_clip", default=False,
|
|
||||||
tooltip="Clip sky-region depth to the 99th percentile of foreground depth before "
|
|
||||||
"normalisation. Requires a 'sky' tensor in the geometry "
|
|
||||||
"(DA3-Mono-Large or DA3-Metric-Large); raises an error otherwise."),
|
|
||||||
]),
|
|
||||||
io.DynamicCombo.Option("sky_mask", []),
|
io.DynamicCombo.Option("sky_mask", []),
|
||||||
io.DynamicCombo.Option("confidence", []),
|
io.DynamicCombo.Option("confidence", []),
|
||||||
]),
|
]),
|
||||||
@ -365,7 +372,7 @@ class DepthAnything3Render(io.ComfyNode):
|
|||||||
def execute(cls, geometry, output) -> io.NodeOutput:
|
def execute(cls, geometry, output) -> io.NodeOutput:
|
||||||
output_val = output["output"]
|
output_val = output["output"]
|
||||||
|
|
||||||
if output_val == "depth":
|
if output_val in ("depth", "depth_colored"):
|
||||||
normalization = output["normalization"]
|
normalization = output["normalization"]
|
||||||
apply_sky_clip = output["apply_sky_clip"]
|
apply_sky_clip = output["apply_sky_clip"]
|
||||||
if apply_sky_clip and "sky" not in geometry:
|
if apply_sky_clip and "sky" not in geometry:
|
||||||
@ -380,7 +387,8 @@ class DepthAnything3Render(io.ComfyNode):
|
|||||||
da3_preprocess.apply_sky_aware_clip(depth[i], sky[i])
|
da3_preprocess.apply_sky_aware_clip(depth[i], sky[i])
|
||||||
for i in range(depth.shape[0])
|
for i in range(depth.shape[0])
|
||||||
], dim=0)
|
], dim=0)
|
||||||
result = cls._depth_to_image(depth, sky, normalization)
|
grey = cls._depth_to_image(depth, sky, normalization) # (B,H,W,3) greyscale
|
||||||
|
result = _turbo(grey[..., 0]) if output_val == "depth_colored" else grey
|
||||||
|
|
||||||
elif output_val == "sky_mask":
|
elif output_val == "sky_mask":
|
||||||
if "sky" not in geometry:
|
if "sky" not in geometry:
|
||||||
|
|||||||
@ -9,6 +9,7 @@ import folder_paths
|
|||||||
from comfy_api.latest import ComfyExtension, Types, io
|
from comfy_api.latest import ComfyExtension, Types, io
|
||||||
from typing_extensions import override
|
from typing_extensions import override
|
||||||
|
|
||||||
|
from comfy.ldm.colormap import turbo as _turbo
|
||||||
from comfy.ldm.moge.model import MoGeModel
|
from comfy.ldm.moge.model import MoGeModel
|
||||||
from comfy.ldm.moge.geometry import triangulate_grid_mesh
|
from comfy.ldm.moge.geometry import triangulate_grid_mesh
|
||||||
from comfy.ldm.moge.panorama import get_panorama_cameras, split_panorama_image, merge_panorama_depth, spherical_uv_to_directions, _uv_grid
|
from comfy.ldm.moge.panorama import get_panorama_cameras, split_panorama_image, merge_panorama_depth, spherical_uv_to_directions, _uv_grid
|
||||||
@ -31,18 +32,6 @@ DA3Geometry = io.Custom("DA3_GEOMETRY")
|
|||||||
# "image": torch.Tensor (B, H, W, 3) in [0, 1], CPU (always present)
|
# "image": torch.Tensor (B, H, W, 3) in [0, 1], CPU (always present)
|
||||||
|
|
||||||
|
|
||||||
def _turbo(x: torch.Tensor) -> torch.Tensor:
|
|
||||||
"""Anton Mikhailov polynomial approximation of the turbo colormap."""
|
|
||||||
x = x.clamp(0.0, 1.0)
|
|
||||||
x2 = x * x
|
|
||||||
x3 = x2 * x
|
|
||||||
x4 = x2 * x2
|
|
||||||
x5 = x4 * x
|
|
||||||
r = 0.13572138 + 4.61539260*x - 42.66032258*x2 + 132.13108234*x3 - 152.94239396*x4 + 59.28637943*x5
|
|
||||||
g = 0.09140261 + 2.19418839*x + 4.84296658*x2 - 14.18503333*x3 + 4.27729857*x4 + 2.82956604*x5
|
|
||||||
b = 0.10667330 + 12.64194608*x - 60.58204836*x2 + 110.36276771*x3 - 89.90310912*x4 + 27.34824973*x5
|
|
||||||
return torch.stack([r, g, b], dim=-1).clamp(0.0, 1.0)
|
|
||||||
|
|
||||||
|
|
||||||
def _normals_from_points(points: torch.Tensor) -> torch.Tensor:
|
def _normals_from_points(points: torch.Tensor) -> torch.Tensor:
|
||||||
"""Camera-space surface normals from a (B, H, W, 3) point map (v1 fallback)."""
|
"""Camera-space surface normals from a (B, H, W, 3) point map (v1 fallback)."""
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user