nodes/Limbicnation-ComfyUIDepthEstimation
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: handle single channel image conversion and tensor shape issues

Browse Source 
- Add proper single channel to RGB conversion
- Fix tensor shape handling
- Improve image format conversion
- Add proper channel order handling
This commit is contained in:
limbicnation
2024-11-29 20:43:57 +01:00
parent b161a7d3c6
commit 4f05c39cbd
1 changed files with 48 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
depth_estimation_node.py
+48 -50
View File
@@ -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)}")