cedarconnor-ComfyUI-DreamCube/core/depth_interface.py

"""
ComfyUI-DreamCube - Generic Depth Model Interface

This module provides a generic interface for applying depth estimation
to cubemap faces, compatible with any ComfyUI depth node output.

Author: Cedar
License: Apache 2.0
"""

import numpy as np
import torch
from typing import Optional, Callable, List, Dict

from .cubemap import CubemapData


class DepthModelInterface:
    """
    Generic interface for applying any depth model to cubemap faces.

    This class doesn't contain depth models itself - it provides utilities
    for working with depth outputs from ComfyUI depth nodes.
    """

    def __init__(self, normalization: str = 'global'):
        """
        Initialize depth interface.

        Args:
            normalization: 'global', 'per_face', or 'adaptive'
                - global: Normalize all faces to same depth range
                - per_face: Normalize each face independently
                - adaptive: Use model's native output range
        """
        self.normalization = normalization
        self.supported_formats = ['comfyui_depth_anything', 'midas', 'custom']

    def apply_depth_to_face(
        self,
        cubemap: CubemapData,
        face_name: str,
        depth_map: np.ndarray
    ) -> CubemapData:
        """
        Apply depth estimation to a single cubemap face.

        Args:
            cubemap: CubemapData object
            face_name: Which face to process ('front', 'back', etc.)
            depth_map: Depth map from any ComfyUI depth node [H, W] or [H, W, 1]

        Returns:
            Updated cubemap with depth assigned to specified face

        Raises:
            ValueError: If depth_map size doesn't match face resolution
        """
        # Handle different depth map formats
        if isinstance(depth_map, torch.Tensor):
            depth_map = depth_map.cpu().numpy()

        # Ensure 2D or 3D array
        if depth_map.ndim == 4:  # [B, H, W, C]
            depth_map = depth_map[0]  # Take first batch

        if depth_map.ndim == 3 and depth_map.shape[2] > 1:
            # If multi-channel, take first channel (grayscale depth)
            depth_map = depth_map[:, :, 0]

        # Normalize depth
        depth_map = self._normalize_depth(depth_map)

        # Set depth face
        cubemap.set_depth_face(face_name, depth_map)

        return cubemap

    def apply_depth_to_all_faces(
        self,
        cubemap: CubemapData,
        depth_maps: Dict[str, np.ndarray]
    ) -> CubemapData:
        """
        Apply depth maps to all 6 cubemap faces.

        Args:
            cubemap: CubemapData object
            depth_maps: Dictionary mapping face names to depth arrays

        Returns:
            Updated cubemap with all depth faces set
        """
        all_depths = []

        # First pass: collect all depths for global normalization
        for face_name in cubemap.get_face_names():
            if face_name not in depth_maps:
                raise ValueError(f"Missing depth for face: {face_name}")

            depth = depth_maps[face_name]
            if isinstance(depth, torch.Tensor):
                depth = depth.cpu().numpy()

            # Handle shape variations
            if depth.ndim == 4:
                depth = depth[0]
            if depth.ndim == 3 and depth.shape[2] > 1:
                depth = depth[:, :, 0]
            if depth.ndim == 3 and depth.shape[2] == 1:
                depth = depth[:, :, 0]

            all_depths.append(depth)

        # Global normalization if needed
        if self.normalization == 'global':
            global_min = min(d.min() for d in all_depths)
            global_max = max(d.max() for d in all_depths)

            # Normalize all depths to [0, 1]
            for i, depth in enumerate(all_depths):
                if global_max - global_min > 1e-6:
                    all_depths[i] = (depth - global_min) / (global_max - global_min)

        # Second pass: assign normalized depths
        for face_name, depth in zip(cubemap.get_face_names(), all_depths):
            if self.normalization == 'per_face':
                depth = self._normalize_depth(depth)

            cubemap.set_depth_face(face_name, depth)

        return cubemap

    def batch_process_faces(
        self,
        cubemap: CubemapData,
        depth_estimator: Callable[[np.ndarray], np.ndarray]
    ) -> CubemapData:
        """
        Process all 6 faces through a depth estimation callback.

        This is useful when you want to programmatically call a depth
        estimator on each face.

        Args:
            cubemap: Input cubemap with RGB faces
            depth_estimator: Function that takes RGB image and returns depth map

        Returns:
            CubemapData with depth_faces populated
        """
        depth_maps = {}

        for face_name in cubemap.get_face_names():
            face_rgb = cubemap.get_face(face_name)
            if face_rgb is None:
                raise ValueError(f"Face {face_name} is not set")

            # Call depth estimator
            depth_map = depth_estimator(face_rgb)

            depth_maps[face_name] = depth_map

        # Apply all depths
        return self.apply_depth_to_all_faces(cubemap, depth_maps)

    def _normalize_depth(self, depth: np.ndarray) -> np.ndarray:
        """
        Normalize depth to [0, 1] range.

        Args:
            depth: Raw depth map

        Returns:
            Normalized depth map
        """
        d_min, d_max = depth.min(), depth.max()

        # Avoid division by zero
        if d_max - d_min < 1e-6:
            return np.zeros_like(depth)

        return (depth - d_min) / (d_max - d_min)

    def convert_comfyui_depth(self, depth_tensor: torch.Tensor) -> np.ndarray:
        """
        Convert ComfyUI depth node output to numpy array.

        ComfyUI typically outputs depth as [B, H, W, C] tensors.

        Args:
            depth_tensor: Depth tensor from ComfyUI node

        Returns:
            Numpy depth array [H, W]
        """
        # Convert to numpy
        if isinstance(depth_tensor, torch.Tensor):
            depth_np = depth_tensor.cpu().numpy()
        else:
            depth_np = depth_tensor

        # Handle batch dimension
        if depth_np.ndim == 4:
            depth_np = depth_np[0]  # Take first batch item

        # Handle channel dimension
        if depth_np.ndim == 3:
            if depth_np.shape[2] == 1:
                depth_np = depth_np[:, :, 0]
            else:
                # Take first channel or average
                depth_np = depth_np[:, :, 0]

        return depth_np

    def merge_rgb_depth(
        self,
        cubemap_rgb: CubemapData,
        cubemap_depth: CubemapData
    ) -> CubemapData:
        """
        Merge RGB and depth cubemaps into a single RGBD cubemap.

        Args:
            cubemap_rgb: CubemapData with RGB faces
            cubemap_depth: CubemapData with depth faces

        Returns:
            New CubemapData with combined RGBD faces
        """
        if cubemap_rgb.resolution != cubemap_depth.resolution:
            raise ValueError("RGB and depth cubemaps must have same resolution")

        # Create new cubemap
        merged = CubemapData(resolution=cubemap_rgb.resolution)

        # Merge faces
        for face_name in cubemap_rgb.get_face_names():
            rgb_face = cubemap_rgb.get_face(face_name)
            depth_face = cubemap_depth.get_depth_face(face_name)

            if rgb_face is None or depth_face is None:
                raise ValueError(f"Missing data for face: {face_name}")

            # Copy RGB
            merged.set_face(face_name, rgb_face.copy())

            # Copy depth
            merged.set_depth_face(face_name, depth_face.copy())

        return merged

    def extract_depth_channel(self, cubemap: CubemapData) -> CubemapData:
        """
        Extract depth channel as separate cubemap for visualization.

        Args:
            cubemap: CubemapData with depth information

        Returns:
            New CubemapData with depth as RGB (for preview)
        """
        if not cubemap.has_depth:
            raise ValueError("Cubemap has no depth information")

        # Create new cubemap for depth visualization
        depth_cubemap = CubemapData(resolution=cubemap.resolution)

        for face_name in cubemap.get_face_names():
            depth_face = cubemap.get_depth_face(face_name)
            if depth_face is None:
                continue

            # Convert depth to RGB (grayscale visualization)
            if depth_face.ndim == 2:
                depth_rgb = np.stack([depth_face] * 3, axis=-1)
            else:
                # Already has channel dimension
                d = depth_face[:, :, 0]
                depth_rgb = np.stack([d] * 3, axis=-1)

            depth_cubemap.set_face(face_name, depth_rgb)

        return depth_cubemap