200 lines
6.8 KiB
Python
200 lines
6.8 KiB
Python
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
|
|
|