From 7f9377b421e4a20cfa7e8f572ed35f935c38b1ae Mon Sep 17 00:00:00 2001 From: limbicnation Date: Fri, 29 Nov 2024 18:09:53 +0100 Subject: [PATCH 1/9] fix: restructure depth estimation node and add proper packaging --- depth_estimation_node.py | 120 ++++++++++++++++----------------------- requirements.txt | 12 ++-- 2 files changed, 57 insertions(+), 75 deletions(-) diff --git a/depth_estimation_node.py b/depth_estimation_node.py index ba08f9a..ddf6f60 100644 --- a/depth_estimation_node.py +++ b/depth_estimation_node.py @@ -3,7 +3,8 @@ import numpy as np import torch from transformers import pipeline from PIL import Image, ImageFilter, ImageOps -from comfy.nodes import Node, register_node +import folder_paths +from comfy.model_management import get_torch_device def ensure_odd(value): """Ensure the value is an odd integer.""" @@ -34,75 +35,52 @@ def auto_contrast(image): """Apply automatic contrast adjustment to the image.""" return ImageOps.autocontrast(image) -class DepthEstimationNode(Node): - def __init__(self, blur_radius=2.0, median_size=5, device="cpu"): - super().__init__() - self.blur_radius = blur_radius - self.median_size = ensure_odd(median_size) - self.device = 0 if device == "gpu" and torch.cuda.is_available() else -1 - self.pipe = pipeline(task="depth-estimation", model="LiheYoung/depth-anything-large-hf", device=self.device) - - def process_image(self, image): - if self.device == 0: - image = image.convert("RGB") # Ensure image is in RGB format - inputs = self.pipe.feature_extractor(images=image, return_tensors="pt").to(self.device) - with torch.no_grad(): - outputs = self.pipe.model(**inputs) - result = self.pipe.post_process(outputs, (image.height, image.width)) - else: - result = self.pipe(image) - - # Convert depth data to a NumPy array if not already one - depth_data = np.array(result["depth"]) - - # Normalize and convert to uint8 - depth_normalized = (depth_data - depth_data.min()) / (depth_data.max() - depth_data.min() + 1e-8) # Avoid zero division - depth_uint8 = (255 * depth_normalized).astype(np.uint8) - - # Create an image from the processed depth data - depth_image = Image.fromarray(depth_uint8) - - # Apply a median filter to reduce noise - depth_image = depth_image.filter(ImageFilter.MedianFilter(size=self.median_size)) - - # Enhanced edge detection with more feathering - edges = depth_image.filter(ImageFilter.FIND_EDGES) - edges = edges.filter(ImageFilter.GaussianBlur(radius=2 * self.blur_radius)) - edges = edges.point(lambda x: 255 if x > 20 else 0) # Adjusted threshold - - # Create a mask from the edges - mask = edges.convert("L") - - # Blur only the edges using the mask - blurred_edges = depth_image.filter(ImageFilter.GaussianBlur(radius=self.blur_radius * 2)) - - # Combine the blurred edges with the original depth image using the mask - combined_image = Image.composite(blurred_edges, depth_image, mask) - - # Apply auto gamma correction with a lower gamma to darken the image - gamma_corrected_image = gamma_correction(combined_image, gamma=0.7) - - # Apply auto contrast - final_image = auto_contrast(gamma_corrected_image) - - # Additional post-processing: Sharpen the final image - final_image = final_image.filter(ImageFilter.SHARPEN) - - return final_image - - def forward(self, image_path: str): - if not os.path.exists(image_path): - raise FileNotFoundError(f"The input image path does not exist: {image_path}") - - image = Image.open(image_path) - return self.process_image(image) - -@register_node -class ComfyUIDepthEstimationNode(DepthEstimationNode): +class DepthEstimationNode: def __init__(self): - super().__init__(blur_radius=2.0, median_size=5, device="cpu") + self.device = get_torch_device() + self.depth_estimator = None + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "blur_radius": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0}), + "median_size": ("INT", {"default": 5, "min": 3, "max": 11, "step": 2}) + } + } - def execute(self, image_path: str, output_path: str): - final_image = self.forward(image_path) - final_image.save(output_path) - print(f"Processed and saved: {output_path}") + RETURN_TYPES = ("IMAGE",) + FUNCTION = "estimate_depth" + CATEGORY = "image/depth" + + def ensure_model_loaded(self): + if self.depth_estimator is None: + self.depth_estimator = pipeline("depth-estimation", device=self.device) + + def estimate_depth(self, image, blur_radius=2.0, median_size=5): + self.ensure_model_loaded() + + # Convert image to PIL + image_pil = Image.fromarray((image[0] * 255).astype(np.uint8)) + + # Process image + depth_map = self.depth_estimator(image_pil)["depth"] + + # Post-processing + depth_map = depth_map.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + depth_map = depth_map.filter(ImageFilter.MedianFilter(size=median_size)) + + # Convert back to tensor format + depth_tensor = np.array(depth_map).astype(np.float32) / 255.0 + depth_tensor = depth_tensor[None, ...] + + return (depth_tensor,) + +NODE_CLASS_MAPPINGS = { + "DepthEstimationNode": DepthEstimationNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "DepthEstimationNode": "Depth Estimation" +} diff --git a/requirements.txt b/requirements.txt index 3cccaa0..8741038 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,8 @@ -transformers==4.12.3 -Pillow==8.4.0 -numpy==1.21.2 -# Add other necessary dependencies without specifying torch again +# Core dependencies matching ComfyUI +torch>=2.0.0 +transformers>=4.28.1 +Pillow>=9.0.0 +numpy>=1.21.2 + +# Additional dependencies specific to depth estimation node +timm>=0.6.12 # Required for depth estimation models From 248ae4e64ae95b70a14be6a07cd3c378e7fb3ad3 Mon Sep 17 00:00:00 2001 From: limbicnation Date: Fri, 29 Nov 2024 18:33:02 +0100 Subject: [PATCH 2/9] Fixed import path to match actual file name --- __init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/__init__.py b/__init__.py index 2ccd9ad..2a2fa22 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,11 @@ -from .your_node_file import ComfyUIDepthEstimationNode +from .depth_estimation_node import DepthEstimationNode NODE_CLASS_MAPPINGS = { - "ComfyUIDepthEstimationNode": ComfyUIDepthEstimationNode + "DepthEstimationNode": DepthEstimationNode } NODE_DISPLAY_NAME_MAPPINGS = { - "ComfyUIDepthEstimationNode": "Depth Estimation Node" + "DepthEstimationNode": "Depth Estimation Node" } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] From 87a4afd1087cadb606f4c0bbfe512dd264f3b864 Mon Sep 17 00:00:00 2001 From: limbicnation Date: Fri, 29 Nov 2024 18:43:09 +0100 Subject: [PATCH 3/9] feat: enhance __init__.py with version tracking, web directory config and improved documentation - Add module docstring - Add version tracking (__version__) - Add web directory config for UI components - Add explicit __all__ exports list - Improve code organization --- __init__.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/__init__.py b/__init__.py index 2a2fa22..c36e481 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,30 @@ +""" +ComfyUI Depth Estimation Node +A custom node for depth map estimation using transformer models. +""" + from .depth_estimation_node import DepthEstimationNode +# Version info +__version__ = "1.0.0" + +# Node class mapping for ComfyUI registration NODE_CLASS_MAPPINGS = { "DepthEstimationNode": DepthEstimationNode } +# Display names for UI NODE_DISPLAY_NAME_MAPPINGS = { - "DepthEstimationNode": "Depth Estimation Node" + "DepthEstimationNode": "Depth Estimation" } -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +# Web extension info for ComfyUI +WEB_DIRECTORY = "./js" + +# Module exports +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "__version__", + "WEB_DIRECTORY" +] From 778ea98dd2291c13f29b4012c3fc606319389735 Mon Sep 17 00:00:00 2001 From: limbicnation Date: Fri, 29 Nov 2024 18:58:21 +0100 Subject: [PATCH 4/9] feat: Depth estimation node with enhanced error handling and post-processing adds input validation and error handling implements auto-contrast and gamma correction options fixes indentation and formatting issues improves performance with conditional filters --- __init__.py | 4 +- depth_estimation_node.py | 85 +++++++++++++++++++++++++++++++--------- 2 files changed, 69 insertions(+), 20 deletions(-) diff --git a/__init__.py b/__init__.py index c36e481..1793c16 100644 --- a/__init__.py +++ b/__init__.py @@ -13,7 +13,7 @@ NODE_CLASS_MAPPINGS = { "DepthEstimationNode": DepthEstimationNode } -# Display names for UI +# Display names for UI NODE_DISPLAY_NAME_MAPPINGS = { "DepthEstimationNode": "Depth Estimation" } @@ -27,4 +27,4 @@ __all__ = [ "NODE_DISPLAY_NAME_MAPPINGS", "__version__", "WEB_DIRECTORY" -] +] \ No newline at end of file diff --git a/depth_estimation_node.py b/depth_estimation_node.py index ddf6f60..ab3a769 100644 --- a/depth_estimation_node.py +++ b/depth_estimation_node.py @@ -1,3 +1,8 @@ +""" +ComfyUI Depth Estimation Node - Verified and Enhanced Version +A custom node for depth map estimation using transformer models. +""" + import os import numpy as np import torch @@ -13,12 +18,15 @@ def ensure_odd(value): def convert_path(path): """Convert path for compatibility between Windows and WSL.""" - if os.name == 'nt': # If running on Windows + if os.name == 'nt': return path.replace('\\', '/') return path def gamma_correction(img, gamma=1.0): """Apply gamma correction to the image.""" + if not isinstance(img, (Image.Image, np.ndarray)): + raise TypeError("Input must be PIL Image or numpy array") + inv_gamma = 1.0 / gamma table = [((i / 255.0) ** inv_gamma) * 255 for i in range(256)] table = np.array(table, np.uint8) @@ -28,6 +36,8 @@ def auto_gamma_correction(image): """Automatically adjust gamma correction for the image.""" image_array = np.array(image).astype(np.float32) / 255.0 mean_luminance = np.mean(image_array) + if mean_luminance <= 0: + return image gamma = np.log(0.5) / np.log(mean_luminance) return gamma_correction(image, gamma=gamma) @@ -46,41 +56,80 @@ class DepthEstimationNode: "required": { "image": ("IMAGE",), "blur_radius": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0}), - "median_size": ("INT", {"default": 5, "min": 3, "max": 11, "step": 2}) + "median_size": ("INT", {"default": 5, "min": 3, "max": 11, "step": 2}), + "apply_auto_contrast": ("BOOLEAN", {"default": True}), + "apply_gamma": ("BOOLEAN", {"default": True}) } } RETURN_TYPES = ("IMAGE",) FUNCTION = "estimate_depth" - CATEGORY = "image/depth" + CATEGORY = "image/processing" def ensure_model_loaded(self): + """Ensure the depth estimation model is loaded.""" if self.depth_estimator is None: - self.depth_estimator = pipeline("depth-estimation", device=self.device) + try: + self.depth_estimator = pipeline("depth-estimation", device=self.device) + except Exception as e: + raise RuntimeError(f"Failed to load depth estimation model: {str(e)}") - def estimate_depth(self, image, blur_radius=2.0, median_size=5): + def estimate_depth(self, image, blur_radius=2.0, median_size=5, + apply_auto_contrast=True, apply_gamma=True): + """ + Estimate depth from input image with optional post-processing. + + Args: + image (torch.Tensor): Input image tensor (B,H,W,C) + blur_radius (float): Gaussian blur radius + median_size (int): Median filter kernel size + apply_auto_contrast (bool): Whether to apply automatic contrast + apply_gamma (bool): Whether to apply gamma correction + + Returns: + tuple(torch.Tensor): Processed depth map tensor (B,H,W) + """ self.ensure_model_loaded() + # Input validation + if not isinstance(image, np.ndarray) or image.ndim != 4: + raise ValueError("Input image must be 4D numpy array (B,H,W,C)") + # Convert image to PIL image_pil = Image.fromarray((image[0] * 255).astype(np.uint8)) - # Process image - depth_map = self.depth_estimator(image_pil)["depth"] - - # Post-processing - depth_map = depth_map.filter(ImageFilter.GaussianBlur(radius=blur_radius)) - depth_map = depth_map.filter(ImageFilter.MedianFilter(size=median_size)) - - # Convert back to tensor format - depth_tensor = np.array(depth_map).astype(np.float32) / 255.0 - depth_tensor = depth_tensor[None, ...] - - return (depth_tensor,) + try: + # Generate depth map + depth_map = self.depth_estimator(image_pil)["depth"] + + # Post-processing pipeline + median_size = ensure_odd(median_size) # Ensure odd kernel size + + if blur_radius > 0: + depth_map = depth_map.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + + depth_map = depth_map.filter(ImageFilter.MedianFilter(size=median_size)) + + if apply_auto_contrast: + depth_map = auto_contrast(depth_map) + + if apply_gamma: + depth_map = auto_gamma_correction(depth_map) + + # Convert to tensor format + depth_tensor = np.array(depth_map).astype(np.float32) / 255.0 + depth_tensor = depth_tensor[None, ...] # Add batch dimension + + return (depth_tensor,) + + except Exception as e: + raise RuntimeError(f"Depth estimation failed: {str(e)}") +# Node registration NODE_CLASS_MAPPINGS = { "DepthEstimationNode": DepthEstimationNode } NODE_DISPLAY_NAME_MAPPINGS = { "DepthEstimationNode": "Depth Estimation" -} +} \ No newline at end of file From 51b4bceae0b2ba045f9bab9ae3cca5b09b644e8f Mon Sep 17 00:00:00 2001 From: limbicnation Date: Fri, 29 Nov 2024 19:50:54 +0100 Subject: [PATCH 5/9] fix: handle median size validation in DepthEstimationNode --- __init__.py | 2 +- depth_estimation_node.py | 135 ++++++++++++++++++++------------------- 2 files changed, 69 insertions(+), 68 deletions(-) diff --git a/__init__.py b/__init__.py index 1793c16..a062235 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,6 @@ """ ComfyUI Depth Estimation Node -A custom node for depth map estimation using transformer models. +A custom node for depth map estimation using Depth-Anything-V2-Small model. """ from .depth_estimation_node import DepthEstimationNode diff --git a/depth_estimation_node.py b/depth_estimation_node.py index ab3a769..f80da00 100644 --- a/depth_estimation_node.py +++ b/depth_estimation_node.py @@ -1,6 +1,5 @@ """ -ComfyUI Depth Estimation Node - Verified and Enhanced Version -A custom node for depth map estimation using transformer models. +ComfyUI Depth Estimation Node - Using Depth-Anything models """ import os @@ -11,52 +10,30 @@ from PIL import Image, ImageFilter, ImageOps import folder_paths from comfy.model_management import get_torch_device -def ensure_odd(value): - """Ensure the value is an odd integer.""" - value = int(value) - return value if value % 2 == 1 else value + 1 +DEPTH_MODELS = { + "Depth-Anything-Small": "LiheYoung/depth-anything-small", + "Depth-Anything-Base": "LiheYoung/depth-anything-base", + "Depth-Anything-Large": "LiheYoung/depth-anything-large", + "Depth-Anything-V2-Small": "LiheYoung/depth-anything-small-hf", + "Depth-Anything-V2-Base": "LiheYoung/depth-anything-base-hf", +} -def convert_path(path): - """Convert path for compatibility between Windows and WSL.""" - if os.name == 'nt': - return path.replace('\\', '/') - return path - -def gamma_correction(img, gamma=1.0): - """Apply gamma correction to the image.""" - if not isinstance(img, (Image.Image, np.ndarray)): - raise TypeError("Input must be PIL Image or numpy array") - - inv_gamma = 1.0 / gamma - table = [((i / 255.0) ** inv_gamma) * 255 for i in range(256)] - table = np.array(table, np.uint8) - return Image.fromarray(np.array(img).astype(np.uint8)).point(lambda i: table[i]) - -def auto_gamma_correction(image): - """Automatically adjust gamma correction for the image.""" - image_array = np.array(image).astype(np.float32) / 255.0 - mean_luminance = np.mean(image_array) - if mean_luminance <= 0: - return image - gamma = np.log(0.5) / np.log(mean_luminance) - return gamma_correction(image, gamma=gamma) - -def auto_contrast(image): - """Apply automatic contrast adjustment to the image.""" - return ImageOps.autocontrast(image) +MEDIAN_SIZES = ["3", "5", "7", "9", "11"] # Define valid median sizes as strings class DepthEstimationNode: def __init__(self): self.device = get_torch_device() self.depth_estimator = None + self.current_model = None @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), - "blur_radius": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0}), - "median_size": ("INT", {"default": 5, "min": 3, "max": 11, "step": 2}), + "model_name": (list(DEPTH_MODELS.keys()),), # Model dropdown + "blur_radius": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.5}), + "median_size": (MEDIAN_SIZES,), # Median size dropdown "apply_auto_contrast": ("BOOLEAN", {"default": True}), "apply_gamma": ("BOOLEAN", {"default": True}) } @@ -64,67 +41,91 @@ class DepthEstimationNode: RETURN_TYPES = ("IMAGE",) FUNCTION = "estimate_depth" - CATEGORY = "image/processing" + CATEGORY = "image/depth" - def ensure_model_loaded(self): + def ensure_model_loaded(self, model_name): """Ensure the depth estimation model is loaded.""" - if self.depth_estimator is None: + model_path = DEPTH_MODELS[model_name] + if self.depth_estimator is None or self.current_model != model_path: try: - self.depth_estimator = pipeline("depth-estimation", device=self.device) + self.depth_estimator = pipeline( + "depth-estimation", + model=model_path, + device=self.device + ) + self.current_model = model_path except Exception as e: - raise RuntimeError(f"Failed to load depth estimation model: {str(e)}") + raise RuntimeError(f"Failed to load model {model_name}: {str(e)}") - def estimate_depth(self, image, blur_radius=2.0, median_size=5, + def estimate_depth(self, image, model_name, blur_radius=2.0, median_size="3", apply_auto_contrast=True, apply_gamma=True): """ Estimate depth from input image with optional post-processing. - - Args: - image (torch.Tensor): Input image tensor (B,H,W,C) - blur_radius (float): Gaussian blur radius - median_size (int): Median filter kernel size - apply_auto_contrast (bool): Whether to apply automatic contrast - apply_gamma (bool): Whether to apply gamma correction - - Returns: - tuple(torch.Tensor): Processed depth map tensor (B,H,W) """ - self.ensure_model_loaded() + self.ensure_model_loaded(model_name) - # Input validation - if not isinstance(image, np.ndarray) or image.ndim != 4: - raise ValueError("Input image must be 4D numpy array (B,H,W,C)") + # Convert median_size from string to int + median_size_int = int(median_size) - # Convert image to PIL - image_pil = Image.fromarray((image[0] * 255).astype(np.uint8)) + # Convert tensor to numpy if needed + if torch.is_tensor(image): + image = image.cpu().numpy() + + # Ensure image is in range [0, 1] + if image.max() > 1.0: + image = image / 255.0 + + # Convert to RGB if necessary + if image.shape[-1] == 4: # RGBA to RGB + image = image[..., :3] + + # Convert to PIL Image for processing + pil_image = Image.fromarray((image[0] * 255).astype(np.uint8)) try: # Generate depth map - depth_map = self.depth_estimator(image_pil)["depth"] + depth_map = self.depth_estimator(pil_image)["predicted_depth"] + + # Convert depth map to PIL Image if it's not already + if not isinstance(depth_map, Image.Image): + # Normalize depth values to 0-255 range + depth_map = ((depth_map - depth_map.min()) * (255 / (depth_map.max() - depth_map.min()))).astype(np.uint8) + depth_map = Image.fromarray(depth_map) # Post-processing pipeline - median_size = ensure_odd(median_size) # Ensure odd kernel size - if blur_radius > 0: depth_map = depth_map.filter(ImageFilter.GaussianBlur(radius=blur_radius)) - depth_map = depth_map.filter(ImageFilter.MedianFilter(size=median_size)) + depth_map = depth_map.filter(ImageFilter.MedianFilter(size=median_size_int)) if apply_auto_contrast: - depth_map = auto_contrast(depth_map) + depth_map = ImageOps.autocontrast(depth_map) if apply_gamma: - depth_map = auto_gamma_correction(depth_map) + depth_array = np.array(depth_map).astype(np.float32) / 255.0 + mean_luminance = np.mean(depth_array) + if mean_luminance > 0: + gamma = np.log(0.5) / np.log(mean_luminance) + depth_map = self.gamma_correction(depth_map, gamma) - # Convert to tensor format - depth_tensor = np.array(depth_map).astype(np.float32) / 255.0 - depth_tensor = depth_tensor[None, ...] # Add batch dimension + # Convert to numpy array and normalize + depth_array = np.array(depth_map).astype(np.float32) / 255.0 + + # Add batch and channel dimensions to match ComfyUI format (B,H,W,C) + depth_tensor = depth_array[None, ..., None] return (depth_tensor,) except Exception as e: raise RuntimeError(f"Depth estimation failed: {str(e)}") + def gamma_correction(self, img, gamma=1.0): + """Apply gamma correction to the image.""" + inv_gamma = 1.0 / gamma + table = [((i / 255.0) ** inv_gamma) * 255 for i in range(256)] + table = np.array(table, np.uint8) + return Image.fromarray(np.array(img).astype(np.uint8)).point(lambda i: table[i]) + # Node registration NODE_CLASS_MAPPINGS = { "DepthEstimationNode": DepthEstimationNode From b161a7d3c6980e7a9e5f2757aacae9b18eb5bbd8 Mon Sep 17 00:00:00 2001 From: limbicnation Date: Fri, 29 Nov 2024 19:53:38 +0100 Subject: [PATCH 6/9] Added a default value for the median_size parameter in INPUT_TYPES --- depth_estimation_node.py | 122 ++++++++++++++++++++------------------- 1 file changed, 64 insertions(+), 58 deletions(-) diff --git a/depth_estimation_node.py b/depth_estimation_node.py index f80da00..bee5c8c 100644 --- a/depth_estimation_node.py +++ b/depth_estimation_node.py @@ -1,7 +1,3 @@ -""" -ComfyUI Depth Estimation Node - Using Depth-Anything models -""" - import os import numpy as np import torch @@ -18,9 +14,9 @@ DEPTH_MODELS = { "Depth-Anything-V2-Base": "LiheYoung/depth-anything-base-hf", } -MEDIAN_SIZES = ["3", "5", "7", "9", "11"] # Define valid median sizes as strings - class DepthEstimationNode: + MEDIAN_SIZES = ["3", "5", "7", "9", "11"] # Valid median sizes + def __init__(self): self.device = get_torch_device() self.depth_estimator = None @@ -31,9 +27,9 @@ class DepthEstimationNode: return { "required": { "image": ("IMAGE",), - "model_name": (list(DEPTH_MODELS.keys()),), # Model dropdown - "blur_radius": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.5}), - "median_size": (MEDIAN_SIZES,), # Median size dropdown + "model_name": (list(DEPTH_MODELS.keys()),), + "blur_radius": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}), + "median_size": (cls.MEDIAN_SIZES, {"default": "5"}), # Added default value "apply_auto_contrast": ("BOOLEAN", {"default": True}), "apply_gamma": ("BOOLEAN", {"default": True}) } @@ -57,65 +53,75 @@ class DepthEstimationNode: except Exception as e: raise RuntimeError(f"Failed to load model {model_name}: {str(e)}") - def estimate_depth(self, image, model_name, blur_radius=2.0, median_size="3", + def estimate_depth(self, image, model_name, blur_radius=2.0, median_size="5", apply_auto_contrast=True, apply_gamma=True): """ Estimate depth from input image with optional post-processing. """ - self.ensure_model_loaded(model_name) - - # Convert median_size from string to int - median_size_int = int(median_size) - - # Convert tensor to numpy if needed - if torch.is_tensor(image): - image = image.cpu().numpy() - - # Ensure image is in range [0, 1] - if image.max() > 1.0: - image = image / 255.0 - - # Convert to RGB if necessary - if image.shape[-1] == 4: # RGBA to RGB - image = image[..., :3] - - # Convert to PIL Image for processing - pil_image = Image.fromarray((image[0] * 255).astype(np.uint8)) - try: - # Generate depth map - depth_map = self.depth_estimator(pil_image)["predicted_depth"] + # Ensure median_size is a string + median_size = str(median_size) - # Convert depth map to PIL Image if it's not already - if not isinstance(depth_map, Image.Image): - # Normalize depth values to 0-255 range - depth_map = ((depth_map - depth_map.min()) * (255 / (depth_map.max() - depth_map.min()))).astype(np.uint8) - depth_map = Image.fromarray(depth_map) + # Validate median_size + if median_size not in self.MEDIAN_SIZES: + raise ValueError(f"Invalid median_size. Must be one of {self.MEDIAN_SIZES}") - # Post-processing pipeline - if blur_radius > 0: - depth_map = depth_map.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + median_size_int = int(median_size) + self.ensure_model_loaded(model_name) - depth_map = depth_map.filter(ImageFilter.MedianFilter(size=median_size_int)) + # Convert tensor to numpy if needed + if torch.is_tensor(image): + image = image.cpu().numpy() + + # Ensure image is in range [0, 1] + if image.max() > 1.0: + image = image / 255.0 + + # Convert to RGB if necessary + if image.shape[-1] == 4: # RGBA to RGB + image = image[..., :3] + + # Convert to PIL Image for processing + pil_image = Image.fromarray((image[0] * 255).astype(np.uint8)) - if apply_auto_contrast: - depth_map = ImageOps.autocontrast(depth_map) - - if apply_gamma: + try: + # Generate depth map + depth_map = self.depth_estimator(pil_image)["predicted_depth"] + + # Convert depth map to PIL Image if it's not already + if not isinstance(depth_map, Image.Image): + # Normalize depth values to 0-255 range + depth_map = ((depth_map - depth_map.min()) * (255 / (depth_map.max() - depth_map.min()))).astype(np.uint8) + depth_map = Image.fromarray(depth_map) + + # Post-processing pipeline + if blur_radius > 0: + depth_map = depth_map.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + + if median_size_int > 0: + depth_map = depth_map.filter(ImageFilter.MedianFilter(size=median_size_int)) + + if apply_auto_contrast: + depth_map = ImageOps.autocontrast(depth_map) + + if apply_gamma: + depth_array = np.array(depth_map).astype(np.float32) / 255.0 + mean_luminance = np.mean(depth_array) + if mean_luminance > 0: + gamma = np.log(0.5) / np.log(mean_luminance) + depth_map = self.gamma_correction(depth_map, gamma) + + # Convert to numpy array and normalize depth_array = np.array(depth_map).astype(np.float32) / 255.0 - mean_luminance = np.mean(depth_array) - if mean_luminance > 0: - gamma = np.log(0.5) / np.log(mean_luminance) - depth_map = self.gamma_correction(depth_map, gamma) - - # Convert to numpy array and normalize - depth_array = np.array(depth_map).astype(np.float32) / 255.0 - - # Add batch and channel dimensions to match ComfyUI format (B,H,W,C) - depth_tensor = depth_array[None, ..., None] - - return (depth_tensor,) - + + # Add batch and channel dimensions to match ComfyUI format (B,H,W,C) + depth_tensor = depth_array[None, ..., None] + + return (depth_tensor,) + + except Exception as e: + raise RuntimeError(f"Depth estimation failed: {str(e)}") + except Exception as e: raise RuntimeError(f"Depth estimation failed: {str(e)}") From 4f05c39cbd4f3e11e69d3509357c92e29622e114 Mon Sep 17 00:00:00 2001 From: limbicnation Date: Fri, 29 Nov 2024 20:43:57 +0100 Subject: [PATCH 7/9] fix: handle single channel image conversion and tensor shape issues - Add proper single channel to RGB conversion - Fix tensor shape handling - Improve image format conversion - Add proper channel order handling --- depth_estimation_node.py | 98 ++++++++++++++++++++-------------------- 1 file changed, 48 insertions(+), 50 deletions(-) diff --git a/depth_estimation_node.py b/depth_estimation_node.py index bee5c8c..6e26f15 100644 --- a/depth_estimation_node.py +++ b/depth_estimation_node.py @@ -29,7 +29,7 @@ class DepthEstimationNode: "image": ("IMAGE",), "model_name": (list(DEPTH_MODELS.keys()),), "blur_radius": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}), - "median_size": (cls.MEDIAN_SIZES, {"default": "5"}), # Added default value + "median_size": (cls.MEDIAN_SIZES, {"default": "5"}), "apply_auto_contrast": ("BOOLEAN", {"default": True}), "apply_gamma": ("BOOLEAN", {"default": True}) } @@ -59,9 +59,6 @@ class DepthEstimationNode: Estimate depth from input image with optional post-processing. """ try: - # Ensure median_size is a string - median_size = str(median_size) - # Validate median_size if median_size not in self.MEDIAN_SIZES: raise ValueError(f"Invalid median_size. Must be one of {self.MEDIAN_SIZES}") @@ -69,59 +66,60 @@ class DepthEstimationNode: median_size_int = int(median_size) self.ensure_model_loaded(model_name) - # Convert tensor to numpy if needed + # Handle tensor conversion if torch.is_tensor(image): - image = image.cpu().numpy() + # Convert tensor to numpy array + image_np = image.cpu().numpy()[0] # Remove batch dimension + # Scale to 0-255 range if needed + if image_np.max() <= 1.0: + image_np = (image_np * 255).astype(np.uint8) + else: + image_np = image_np.astype(np.uint8) + else: + image_np = image - # Ensure image is in range [0, 1] - if image.max() > 1.0: - image = image / 255.0 + # Ensure RGB format + if image_np.shape[-1] == 4: # RGBA to RGB + image_np = image_np[..., :3] - # Convert to RGB if necessary - if image.shape[-1] == 4: # RGBA to RGB - image = image[..., :3] - - # Convert to PIL Image for processing - pil_image = Image.fromarray((image[0] * 255).astype(np.uint8)) + # Convert to PIL for processing + pil_image = Image.fromarray(image_np) - try: - # Generate depth map - depth_map = self.depth_estimator(pil_image)["predicted_depth"] + # Get depth map + depth_result = self.depth_estimator(pil_image) + depth_map = depth_result["predicted_depth"] + + # Convert tensor to numpy if necessary + if torch.is_tensor(depth_map): + depth_map = depth_map.cpu().numpy() - # Convert depth map to PIL Image if it's not already - if not isinstance(depth_map, Image.Image): - # Normalize depth values to 0-255 range - depth_map = ((depth_map - depth_map.min()) * (255 / (depth_map.max() - depth_map.min()))).astype(np.uint8) - depth_map = Image.fromarray(depth_map) + # Normalize depth values to 0-255 range + depth_map = ((depth_map - depth_map.min()) * (255 / (depth_map.max() - depth_map.min()))).astype(np.uint8) + depth_map = Image.fromarray(depth_map) + + # Apply post-processing + if blur_radius > 0: + depth_map = depth_map.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + + if median_size_int > 0: + depth_map = depth_map.filter(ImageFilter.MedianFilter(size=median_size_int)) - # Post-processing pipeline - if blur_radius > 0: - depth_map = depth_map.filter(ImageFilter.GaussianBlur(radius=blur_radius)) - - if median_size_int > 0: - depth_map = depth_map.filter(ImageFilter.MedianFilter(size=median_size_int)) - - if apply_auto_contrast: - depth_map = ImageOps.autocontrast(depth_map) - - if apply_gamma: - depth_array = np.array(depth_map).astype(np.float32) / 255.0 - mean_luminance = np.mean(depth_array) - if mean_luminance > 0: - gamma = np.log(0.5) / np.log(mean_luminance) - depth_map = self.gamma_correction(depth_map, gamma) - - # Convert to numpy array and normalize + if apply_auto_contrast: + depth_map = ImageOps.autocontrast(depth_map) + + if apply_gamma: depth_array = np.array(depth_map).astype(np.float32) / 255.0 - - # Add batch and channel dimensions to match ComfyUI format (B,H,W,C) - depth_tensor = depth_array[None, ..., None] - - return (depth_tensor,) - - except Exception as e: - raise RuntimeError(f"Depth estimation failed: {str(e)}") - + mean_luminance = np.mean(depth_array) + if mean_luminance > 0: + gamma = np.log(0.5) / np.log(mean_luminance) + depth_map = self.gamma_correction(depth_map, gamma) + + # Convert back to tensor format + depth_array = np.array(depth_map).astype(np.float32) / 255.0 + depth_tensor = depth_array[None, ..., None] # Add batch and channel dims + + return (depth_tensor,) + except Exception as e: raise RuntimeError(f"Depth estimation failed: {str(e)}") From 8c3e6123f3d9bd7ee9aca7ff213327c6d4fd0741 Mon Sep 17 00:00:00 2001 From: limbicnation Date: Fri, 29 Nov 2024 20:51:23 +0100 Subject: [PATCH 8/9] style: fix indentation in estimate_depth method --- depth_estimation_node.py | 30 +++++++++++++++++++----------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/depth_estimation_node.py b/depth_estimation_node.py index 6e26f15..a1dd9ba 100644 --- a/depth_estimation_node.py +++ b/depth_estimation_node.py @@ -15,7 +15,7 @@ DEPTH_MODELS = { } class DepthEstimationNode: - MEDIAN_SIZES = ["3", "5", "7", "9", "11"] # Valid median sizes + MEDIAN_SIZES = ["3", "5", "7", "9", "11"] def __init__(self): self.device = get_torch_device() @@ -40,7 +40,6 @@ class DepthEstimationNode: CATEGORY = "image/depth" def ensure_model_loaded(self, model_name): - """Ensure the depth estimation model is loaded.""" model_path = DEPTH_MODELS[model_name] if self.depth_estimator is None or self.current_model != model_path: try: @@ -55,9 +54,6 @@ class DepthEstimationNode: def estimate_depth(self, image, model_name, blur_radius=2.0, median_size="5", apply_auto_contrast=True, apply_gamma=True): - """ - Estimate depth from input image with optional post-processing. - """ try: # Validate median_size if median_size not in self.MEDIAN_SIZES: @@ -87,14 +83,24 @@ class DepthEstimationNode: # Get depth map depth_result = self.depth_estimator(pil_image) - depth_map = depth_result["predicted_depth"] - # Convert tensor to numpy if necessary + # Convert tensor to numpy and ensure correct dimensions if torch.is_tensor(depth_map): - depth_map = depth_map.cpu().numpy() + depth_map = depth_map.squeeze().cpu().numpy() + + # Ensure depth_map is 2D + if len(depth_map.shape) > 2: + depth_map = depth_map.squeeze() # Normalize depth values to 0-255 range - depth_map = ((depth_map - depth_map.min()) * (255 / (depth_map.max() - depth_map.min()))).astype(np.uint8) + depth_min = depth_map.min() + depth_max = depth_map.max() + if depth_max > depth_min: + depth_map = ((depth_map - depth_min) * (255.0 / (depth_max - depth_min))).astype(np.uint8) + else: + depth_map = np.zeros_like(depth_map, dtype=np.uint8) + + # Convert to PIL Image depth_map = Image.fromarray(depth_map) # Apply post-processing @@ -116,7 +122,10 @@ class DepthEstimationNode: # Convert back to tensor format depth_array = np.array(depth_map).astype(np.float32) / 255.0 - depth_tensor = depth_array[None, ..., None] # Add batch and channel dims + depth_tensor = torch.from_numpy(depth_array)[None, ..., None] # Convert to tensor and add batch and channel dims + + # Move tensor to the correct device + depth_tensor = depth_tensor.to(self.device) return (depth_tensor,) @@ -124,7 +133,6 @@ class DepthEstimationNode: raise RuntimeError(f"Depth estimation failed: {str(e)}") def gamma_correction(self, img, gamma=1.0): - """Apply gamma correction to the image.""" inv_gamma = 1.0 / gamma table = [((i / 255.0) ** inv_gamma) * 255 for i in range(256)] table = np.array(table, np.uint8) From 7d70e0737e60a7f3a7d1c8d2ce15e4cd593c026d Mon Sep 17 00:00:00 2001 From: limbicnation Date: Fri, 29 Nov 2024 21:08:56 +0100 Subject: [PATCH 9/9] fix: depth estimation pipeline data handling & format conversions - Fix tensor to image conversion pipeline - Add proper RGB/grayscale format handling - Implement correct depth map normalization - Add explicit channel stacking for output tensor - Improve error handling throughout pipeline - Fix PIL Image mode specification --- depth_estimation_node.py | 40 +++++++++++++++++++++++++--------------- 1 file changed, 25 insertions(+), 15 deletions(-) diff --git a/depth_estimation_node.py b/depth_estimation_node.py index a1dd9ba..ed0dae0 100644 --- a/depth_estimation_node.py +++ b/depth_estimation_node.py @@ -64,44 +64,49 @@ class DepthEstimationNode: # Handle tensor conversion if torch.is_tensor(image): - # Convert tensor to numpy array image_np = image.cpu().numpy()[0] # Remove batch dimension - # Scale to 0-255 range if needed - if image_np.max() <= 1.0: - image_np = (image_np * 255).astype(np.uint8) - else: - image_np = image_np.astype(np.uint8) else: image_np = image + + # Ensure proper RGB format and scaling + if image_np.max() <= 1.0: + image_np = (image_np * 255).astype(np.uint8) + else: + image_np = image_np.astype(np.uint8) - # Ensure RGB format - if image_np.shape[-1] == 4: # RGBA to RGB + # Convert to RGB if necessary + if len(image_np.shape) == 3 and image_np.shape[-1] == 4: image_np = image_np[..., :3] + elif len(image_np.shape) == 2: + # Convert grayscale to RGB + image_np = np.stack([image_np] * 3, axis=-1) # Convert to PIL for processing pil_image = Image.fromarray(image_np) # Get depth map depth_result = self.depth_estimator(pil_image) + depth_map = depth_result["predicted_depth"] # Convert tensor to numpy and ensure correct dimensions if torch.is_tensor(depth_map): depth_map = depth_map.squeeze().cpu().numpy() - + # Ensure depth_map is 2D - if len(depth_map.shape) > 2: - depth_map = depth_map.squeeze() + depth_map = depth_map.squeeze() # Normalize depth values to 0-255 range depth_min = depth_map.min() depth_max = depth_map.max() if depth_max > depth_min: - depth_map = ((depth_map - depth_min) * (255.0 / (depth_max - depth_min))).astype(np.uint8) + depth_map = ((depth_map - depth_min) * (255.0 / (depth_max - depth_min))) else: - depth_map = np.zeros_like(depth_map, dtype=np.uint8) + depth_map = np.zeros_like(depth_map) + + depth_map = depth_map.astype(np.uint8) # Convert to PIL Image - depth_map = Image.fromarray(depth_map) + depth_map = Image.fromarray(depth_map, mode='L') # Convert as grayscale # Apply post-processing if blur_radius > 0: @@ -122,7 +127,12 @@ class DepthEstimationNode: # Convert back to tensor format depth_array = np.array(depth_map).astype(np.float32) / 255.0 - depth_tensor = torch.from_numpy(depth_array)[None, ..., None] # Convert to tensor and add batch and channel dims + + # Convert single channel to 3 channels + depth_array = np.stack([depth_array] * 3, axis=-1) + + # Add batch dimension + depth_tensor = torch.from_numpy(depth_array).unsqueeze(0) # Move tensor to the correct device depth_tensor = depth_tensor.to(self.device)