Spaces:
Sleeping
Sleeping
File size: 7,180 Bytes
ac305e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 |
import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image
import comfy.model_management
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
from comfy_api.latest import io
class IntensityDepthEstimation(ComfyNodeABC):
"""
Simple intensity-based depth estimation node.
Converts grayscale intensity to depth values using various methods.
"""
@classmethod
def INPUT_TYPES(cls) -> InputTypeDict:
return {
"required": {
"image": (IO.IMAGE, {"tooltip": "Input image for depth estimation"}),
"method": (["intensity", "inverted_intensity", "gradient", "sobel"], {
"default": "intensity",
"tooltip": "Depth estimation method"
}),
"depth_range": (IO.FLOAT, {
"default": 1.0,
"min": 0.1,
"max": 10.0,
"step": 0.1,
"tooltip": "Maximum depth range"
}),
"normalize": (IO.BOOLEAN, {
"default": True,
"tooltip": "Normalize depth values to 0-1 range"
}),
},
"optional": {
"blur_radius": (IO.FLOAT, {
"default": 1.0,
"min": 0.0,
"max": 10.0,
"step": 0.1,
"tooltip": "Gaussian blur radius for smoothing"
}),
}
}
RETURN_TYPES = (IO.IMAGE,)
RETURN_NAMES = ("depth_image",)
FUNCTION = "estimate_depth"
CATEGORY = "image/processing"
DESCRIPTION = "Simple intensity-based depth estimation using various methods"
def estimate_depth(self, image, method, depth_range, normalize, blur_radius=1.0):
"""
Estimate depth from image intensity using various methods.
Args:
image: Input image tensor [B, H, W, C]
method: Depth estimation method
depth_range: Maximum depth range
normalize: Whether to normalize output
blur_radius: Gaussian blur radius for smoothing
Returns:
depth_image: Estimated depth map [B, H, W, C]
"""
device = comfy.model_management.get_torch_device()
image = image.to(device)
# Convert to grayscale if needed
if image.shape[-1] > 1:
# Convert RGB to grayscale using standard weights
gray = 0.299 * image[..., 0:1] + 0.587 * image[..., 1:2] + 0.114 * image[..., 2:3]
else:
gray = image
# Apply Gaussian blur if specified
if blur_radius > 0:
gray = self._apply_gaussian_blur(gray, blur_radius)
# Apply depth estimation method
if method == "intensity":
depth = gray
elif method == "inverted_intensity":
depth = 1.0 - gray
elif method == "gradient":
depth = self._compute_gradient_depth(gray)
elif method == "sobel":
depth = self._compute_sobel_depth(gray)
else:
depth = gray
# Scale by depth range
depth = depth * depth_range
# Normalize if requested
if normalize:
depth_min = depth.min()
depth_max = depth.max()
if depth_max > depth_min:
depth = (depth - depth_min) / (depth_max - depth_min)
else:
depth = torch.zeros_like(depth)
# Convert back to RGB for visualization
depth_rgb = depth.repeat(1, 1, 1, 3)
# Move to intermediate device
depth_rgb = depth_rgb.to(comfy.model_management.intermediate_device())
return (depth_rgb,)
def _apply_gaussian_blur(self, image, radius):
"""Apply Gaussian blur to the image."""
if radius <= 0:
return image
# Create Gaussian kernel
kernel_size = int(2 * radius + 1)
if kernel_size % 2 == 0:
kernel_size += 1
# Generate Gaussian kernel
sigma = radius / 3.0
x = torch.arange(kernel_size, dtype=torch.float32, device=image.device)
x = x - kernel_size // 2
kernel_1d = torch.exp(-(x ** 2) / (2 * sigma ** 2))
kernel_1d = kernel_1d / kernel_1d.sum()
# Create 2D kernel
kernel_2d = kernel_1d[:, None] * kernel_1d[None, :]
kernel_2d = kernel_2d / kernel_2d.sum()
# Reshape for convolution
kernel_2d = kernel_2d.view(1, 1, kernel_size, kernel_size)
# Apply convolution
batch_size, height, width, channels = image.shape
image_reshaped = image.permute(0, 3, 1, 2) # [B, C, H, W]
# Pad image
pad_size = kernel_size // 2
padded = F.pad(image_reshaped, (pad_size, pad_size, pad_size, pad_size), mode='reflect')
# Apply convolution
blurred = F.conv2d(padded, kernel_2d, padding=0)
# Reshape back
blurred = blurred.permute(0, 2, 3, 1) # [B, H, W, C]
return blurred
def _compute_gradient_depth(self, image):
"""Compute depth using gradient magnitude."""
# Compute gradients
grad_x = torch.diff(image, dim=2, prepend=image[:, :, :1, :])
grad_y = torch.diff(image, dim=1, prepend=image[:, :1, :, :])
# Compute gradient magnitude
grad_mag = torch.sqrt(grad_x ** 2 + grad_y ** 2)
# Invert so edges have higher depth
depth = 1.0 - grad_mag
return depth
def _compute_sobel_depth(self, image):
"""Compute depth using Sobel edge detection."""
# Sobel kernels
sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32, device=image.device)
sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32, device=image.device)
# Reshape kernels for convolution
sobel_x = sobel_x.view(1, 1, 3, 3)
sobel_y = sobel_y.view(1, 1, 3, 3)
# Reshape image for convolution
batch_size, height, width, channels = image.shape
image_reshaped = image.permute(0, 3, 1, 2) # [B, C, H, W]
# Pad image
padded = F.pad(image_reshaped, (1, 1, 1, 1), mode='reflect')
# Apply Sobel filters
grad_x = F.conv2d(padded, sobel_x, padding=0)
grad_y = F.conv2d(padded, sobel_y, padding=0)
# Compute gradient magnitude
grad_mag = torch.sqrt(grad_x ** 2 + grad_y ** 2)
# Reshape back
grad_mag = grad_mag.permute(0, 2, 3, 1) # [B, H, W, C]
# Invert so edges have higher depth
depth = 1.0 - grad_mag
return depth
# Node class mappings
NODE_CLASS_MAPPINGS = {
"IntensityDepthEstimation": IntensityDepthEstimation,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"IntensityDepthEstimation": "Intensity Depth Estimation",
}
|