136 lines
4.9 KiB
Python
136 lines
4.9 KiB
Python
"""Shared depth normalization methods for DepthAnythingV3 nodes."""
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def apply_edge_antialiasing(mask):
|
|
"""Apply minimal anti-aliasing ONLY to border pixels (1-2px transition)."""
|
|
# Ensure mask is in correct format [B, 1, H, W]
|
|
if mask.dim() == 2:
|
|
mask = mask.unsqueeze(0).unsqueeze(0)
|
|
elif mask.dim() == 3:
|
|
mask = mask.unsqueeze(0)
|
|
|
|
# Very small 3x3 averaging kernel for minimal smoothing
|
|
kernel = torch.ones((1, 1, 3, 3), device=mask.device, dtype=mask.dtype) / 9.0
|
|
|
|
# Apply minimal blur
|
|
mask_blurred = F.conv2d(mask, kernel, padding=1)
|
|
|
|
# Detect edges: where original mask has transitions
|
|
mask_dilated = F.max_pool2d(mask, kernel_size=3, stride=1, padding=1)
|
|
mask_eroded = 1.0 - F.max_pool2d(1.0 - mask, kernel_size=3, stride=1, padding=1)
|
|
|
|
# Edge zone is where dilated and eroded differ
|
|
edge_zone = (mask_dilated - mask_eroded).abs()
|
|
edge_zone = (edge_zone > 0.01).float()
|
|
|
|
# Apply anti-aliasing ONLY in edge zone
|
|
mask_aa = mask * (1.0 - edge_zone) + mask_blurred * edge_zone
|
|
|
|
return mask_aa
|
|
|
|
|
|
def apply_standard_normalization(depth, invert_depth):
|
|
"""
|
|
Standard min-max normalization (original V3 approach).
|
|
|
|
By default, inverts to match V2-Style convention (close=bright).
|
|
Standard normalization naturally outputs far=bright, so we invert by default.
|
|
"""
|
|
depth_norm = (depth - depth.min()) / (depth.max() - depth.min() + 1e-8)
|
|
|
|
# Invert by default to match V2-Style (close = bright)
|
|
if not invert_depth:
|
|
depth_norm = 1.0 - depth_norm
|
|
|
|
return depth_norm
|
|
|
|
|
|
def apply_v2_style_normalization(depth, sky, device, invert_depth=False):
|
|
"""
|
|
V2-Style disparity normalization (Ltamann/TBG approach).
|
|
|
|
- Converts depth to disparity (1/depth) so sky becomes black
|
|
- Uses content-only normalization (excludes sky)
|
|
- Applies percentile-based contrast enhancement (1st-99th percentile)
|
|
- Adds contrast boost via power transform (hardcoded to 2.0)
|
|
- Applies edge anti-aliasing for natural transitions (hardcoded to True)
|
|
- Optionally inverts the final output if invert_depth is True
|
|
"""
|
|
epsilon = 1e-6
|
|
contrast_boost = 2.0 # Hardcoded default
|
|
edge_soften = True # Hardcoded default
|
|
|
|
# 1. Create HARD binary content mask
|
|
if sky.max() > 0.1:
|
|
# Threshold of 0.3 for aggressive sky detection
|
|
content_mask_binary = (sky < 0.3).float()
|
|
|
|
# Apply edge anti-aliasing (affects 1-2px border only)
|
|
if edge_soften:
|
|
content_mask_smooth = apply_edge_antialiasing(content_mask_binary)
|
|
else:
|
|
content_mask_smooth = content_mask_binary
|
|
else:
|
|
content_mask_binary = torch.ones_like(depth)
|
|
content_mask_smooth = content_mask_binary
|
|
|
|
# Ensure same shape as depth
|
|
while content_mask_binary.dim() < depth.dim():
|
|
content_mask_binary = content_mask_binary.unsqueeze(0)
|
|
while content_mask_smooth.dim() < depth.dim():
|
|
content_mask_smooth = content_mask_smooth.unsqueeze(0)
|
|
|
|
# 2. Convert depth to disparity (inverse depth) like V2
|
|
disparity = 1.0 / (depth + epsilon)
|
|
|
|
# 3. Use HARD mask for normalization calculations
|
|
disparity_masked = disparity * content_mask_binary
|
|
|
|
# 4. Extract ONLY content pixels for normalization
|
|
content_pixels = disparity_masked[content_mask_binary > 0.5]
|
|
|
|
if content_pixels.numel() > 100:
|
|
# Get min/max from CONTENT ONLY
|
|
disp_min = content_pixels.min()
|
|
disp_max = content_pixels.max()
|
|
|
|
# Use percentile-based normalization for better contrast
|
|
if content_pixels.numel() > 1000:
|
|
sorted_pixels = torch.sort(content_pixels.flatten())[0]
|
|
p1_idx = int(sorted_pixels.numel() * 0.01)
|
|
p99_idx = int(sorted_pixels.numel() * 0.99)
|
|
disp_min = sorted_pixels[p1_idx]
|
|
disp_max = sorted_pixels[p99_idx]
|
|
|
|
# Normalize using content-only range
|
|
disparity_norm = (disparity - disp_min) / (disp_max - disp_min + epsilon)
|
|
disparity_norm = torch.clamp(disparity_norm, 0.0, 1.0)
|
|
else:
|
|
# Fallback for very small content regions
|
|
disp_min = disparity.min()
|
|
disp_max = disparity.max()
|
|
disparity_norm = (disparity - disp_min) / (disp_max - disp_min + epsilon)
|
|
|
|
# 5. Apply contrast boost
|
|
disparity_contrast = torch.pow(disparity_norm, 1.0 / contrast_boost)
|
|
|
|
# 6. Apply SMOOTH mask for final output (with anti-aliased edges)
|
|
disparity_final = disparity_contrast * content_mask_smooth
|
|
|
|
# 7. Apply inversion if requested (do this AFTER all processing)
|
|
if invert_depth:
|
|
disparity_final = 1.0 - disparity_final
|
|
|
|
return disparity_final
|
|
|
|
|
|
def apply_raw_normalization(depth, invert_depth):
|
|
"""Raw/metric depth - no normalization (for 3D reconstruction)."""
|
|
if invert_depth:
|
|
# For raw metric depth, invert as max - depth
|
|
depth = depth.max() - depth
|
|
|
|
return depth
|