complex_nodes.py

import math
import torch
import torch.nn.functional as F
import numpy as np
from typing import Dict, Tuple, Any
from PIL import Image
from tqdm import tqdm

# reuse existing nodes
from .reprojection_nodes import ReprojectImage, ReprojectDepth
from .metric_depth_nodes import (
    DepthEstimatorNode,
    ZDepthToRayDepthNode,
    CombineDepthsNode,
    DepthRenormalizer
)
from .pointcloud_nodes import (
    ProjectPointCloud,
    DepthToPointCloud,
    TransformPointCloud,
    PointCloudCleaner,
    ProjectAndClean
)
from .flux_fisheye_filling_nodes import OutpaintAnyProjection

# HuggingFace models
possible_models = [
    "Depth-Anything-V2-Metric-Indoor-Base-hf",
    "Depth-Anything-V2-Metric-Indoor-Small-hf",
    "Depth-Anything-V2-Metric-Indoor-Large-hf",
    "Depth-Anything-V2-Metric-Outdoor-Base-hf",
    "Depth-Anything-V2-Metric-Outdoor-Small-hf",
    "Depth-Anything-V2-Metric-Outdoor-Large-hf",
]

class FisheyeDepthEstimator:
    """
    Estimates a full 180° fisheye depthmap by:
      1. Estimating depth on the full fisheye image
      2. Extracting 90° pinhole views (front, left, right, up, down)
      3. Running depth estimation per view
      4. Reprojecting back to fisheye
      5. Renormalizing overlaps
      6. Merging all six maps with selectable mode
    Outputs:
      - depthmap: [B,H,W,1]
      - mask: [B,H,W], circular mask of fisheye region
    """
    @classmethod
    def INPUT_TYPES(cls) -> Dict[str, Any]:
        return {
            "required": {
                "image":              ("IMAGE",),
                "model_name":         (possible_models, {"default": possible_models[0]}),
                "depth_scale":        ("FLOAT",  {"default":1.0, "min":0.0, "max":1000.0, "step":0.01}),
                "pinhole_fov":        ("FLOAT",  {"default":90.0, "min":1.0, "max":179.0}),
                "pinhole_resolution": ("INT",    {"default":1024, "min":64}),
                "fisheye_resolution": ("INT",    {"default":4096, "min":64, "max":8192}),
                "mode":               (["SRC","DST","AVERAGE","SOFTMERGE","DISTANCE_AWARE"], {"default":"AVERAGE"}),
                "softmerge_radius":   ("INT",    {"default":25,    "min":1, "tooltip":"Gaussian radius for merging"}),
                "median_blur_kernel": ("INT",    {"default":1,     "min":1, "max":31, "tooltip":"Kernel size for median blur of depthmap"}),
            }
        }
    RETURN_TYPES = ("TENSOR","MASK")
    RETURN_NAMES = ("depthmap","mask")
    FUNCTION = "estimate_fisheye_depth"
    CATEGORY = "Camera/Depth"

    def estimate_fisheye_depth(
        self,
        image: torch.Tensor,
        model_name: str,
        depth_scale: float,
        pinhole_fov: float,
        pinhole_resolution: int,
        fisheye_resolution: int,
        mode: str,
        softmerge_radius: int,
        median_blur_kernel: int,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # helper nodes
        de_node   = DepthEstimatorNode()
        z2r_node  = ZDepthToRayDepthNode()
        ri_node   = ReprojectImage()
        rd_node   = ReprojectDepth()
        ren_node  = DepthRenormalizer()
        comb_node = CombineDepthsNode()

        # constants
        fisheye_fov = 180.0
        pin_w = pin_h = pinhole_resolution
        fish_w = fish_h = fisheye_resolution

        # 1) Full fisheye depth + mask
        depth_full, = de_node.estimate_depth(image, model_name, depth_scale)
        mask_full   = (depth_full > 0).float()

        # 2) Pinhole orientations (5 views)
        rotations = [
            (0,   0,  0),    # front
            (0,  45,  0),    # right
            (0, -45,  0),    # left
            (45,  0,  0),    # up
            (-45, 0,  0),    # down
        ]

        fisheye_depths = []
        fisheye_masks  = []

        # euler → matrix
        def euler_to_matrix(pitch, yaw, roll):
            p, y, r = map(math.radians, (pitch, yaw, roll))
            Rx = torch.tensor([[1,0,0],[0,math.cos(p),-math.sin(p)],[0,math.sin(p),math.cos(p)]], dtype=torch.float32)
            Ry = torch.tensor([[math.cos(y),0,math.sin(y)],[0,1,0],[-math.sin(y),0,math.cos(y)]], dtype=torch.float32)
            Rz = torch.tensor([[math.cos(r),-math.sin(r),0],[math.sin(r),math.cos(r),0],[0,0,1]], dtype=torch.float32)
            R = Rz @ Ry @ Rx
            M = torch.eye(4, dtype=torch.float32)
            M[:3, :3] = R
            return M

        # 3) Process each orientation
        for pitch, yaw, roll in rotations:
            M     = euler_to_matrix(pitch, yaw, roll)
            M_np  = M.numpy()
            M_inv = torch.inverse(M).numpy()

            # fisheye → pinhole
            img_pin, mask_pin = ri_node.reproject_image(
                image,
                input_horiszontal_fov = fisheye_fov,
                output_horiszontal_fov= pinhole_fov,
                input_projection     = "FISHEYE",
                output_projection    = "PINHOLE",
                output_width         = pin_w,
                output_height        = pin_h,
                transform_matrix     = M_np,
                feathering           = 0,
            )

            # estimate pinhole depth
            depth_pin, = de_node.estimate_depth(img_pin, model_name, depth_scale, median_blur_kernel=median_blur_kernel)
            depth_pin, = z2r_node.depth_to_ray_depth(
                depth_pin,
                pinhole_fov,
            )
            # pinhole → fisheye
            fish_depth, fish_mask = rd_node.reproject_depth(
                depth_pin,
                input_horizontal_fov = pinhole_fov,
                output_horizontal_fov= fisheye_fov,
                input_projection     = "PINHOLE",
                output_projection    = "FISHEYE",
                output_width         = fish_w,
                output_height        = fish_h,
                transform_matrix     = M_inv,
            )
            # squeeze mask to [B,H,W]
            fish_mask = fish_mask.squeeze(1)

            fisheye_depths.append(fish_depth) # [B,H,W]
            fisheye_masks.append(fish_mask)
        fisheye_depths.append(depth_full) # [B,H,W]
        fisheye_masks.append(mask_full.squeeze(-1)) # [B,H,W 1]
        # merged mask
        merged_mask = torch.sum(torch.stack(fisheye_masks), dim=0) > 0.5
        # print(fisheye_depths[0].shape, fisheye_depths[-1].shape, merged_mask.shape)
        # 4) Merge in sequence
        d_acc = fisheye_depths[0]
        m_acc = fisheye_masks[0]
        for d_new, m_new in zip(fisheye_depths[1:-1], fisheye_masks[1:-1]):
            d_norm, = ren_node.renormalize_depth(d_new, d_acc, m_new, m_acc, use_inverse=False)
            d_acc, m_acc = comb_node.combine_depths(
                d_acc,
                m_acc,
                d_norm,
                m_new,
                mode             = mode,
                invert_mask      = False,
                softmerge_radius = softmerge_radius
            )
        # last depth is full fisheye -mode is SRC to prevent resolution drawbacks
        m_new_last = fisheye_masks[-1]
        d_new_last = fisheye_depths[-1]
        d_norm, = ren_node.renormalize_depth(d_new_last, d_acc, m_new_last, m_acc, use_inverse=False)
        d_acc,m_acc = comb_node.combine_depths(
            d_acc,
            m_acc,
            d_norm,
            m_new_last,
            mode             = "SRC",
            invert_mask      = False,
            softmerge_radius = softmerge_radius
        )

        # 5) Circular mask
        ys = torch.arange(fish_h, device=d_acc.device).view(1, fish_h, 1)
        xs = torch.arange(fish_w, device=d_acc.device).view(1, 1, fish_w)
        cy = (fish_h - 1) / 2.0
        cx = (fish_w - 1) / 2.0
        dist2   = (ys - cy)**2 + (xs - cx)**2
        radius2 = (min(fish_w, fish_h) / 2.0)**2
        circ_mask = (dist2 <= radius2).float()
        return d_acc, circ_mask

class PointcloudTrajectoryEnricher:
    """
    Enriches a pointcloud along a camera trajectory by outpainting missing regions per view.

    Returns:
      - enriched_pointcloud: [N,4+] tensor
      - debug_image:         [1,H,W,3] tensor
      - debug_depth:         [1,H,W,1] tensor
    """
    @classmethod
    def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]:
        return {
            "required": {
                "pointcloud":            ("TENSOR",),
                "trajectory":            ("TENSOR",),  # (K,4,4)
                "camera_type":           (("PINHOLE","FISHEYE","EQUIRECTANGULAR"),),
                "horizontal_fov":        ("FLOAT", {"default":90.0}),
                "width":                 ("INT",   {"default":512}),
                "height":                ("INT",   {"default":512}),
                # outpainting params
                "patch_projection":      (("PINHOLE","FISHEYE","EQUIRECTANGULAR"),),
                "patch_horiz_fov":       ("FLOAT", {"default":90.0}),
                "patch_res":             ("INT",   {"default":512}),
                "patch_phi":             ("FLOAT", {"default":0.0}),
                "patch_theta":           ("FLOAT", {"default":0.0}),
                "prompt":                ("STRING",{"default":""}),
                "num_inference_steps":   ("INT",   {"default":50}),
                "guidance_scale":        ("FLOAT", {"default":7.5}),
                "mask_blur":             ("INT",   {"default":5}),
                # cleaning params
                "voxel_size":            ("FLOAT", {"default":0.07}),
                "min_points_per_voxel":  ("INT",   {"default":3}),
                # depth estimation model
                "model_name":            ("STRING", {"choices": possible_models, "default": possible_models[0]}),
            }
        }
    RETURN_TYPES = ("TENSOR","IMAGE","TENSOR")
    RETURN_NAMES = ("enriched_pointcloud","debug_image","debug_depth")
    FUNCTION = "enrich_trajectory"
    CATEGORY = "Camera/Trajectory"

    def enrich_trajectory(
        self,
        pointcloud: torch.Tensor,
        trajectory: torch.Tensor,
        camera_type: str,
        horizontal_fov: float,
        width: int,
        height: int,
        patch_projection: str,
        patch_horiz_fov: float,
        patch_res: int,
        patch_phi: float,
        patch_theta: float,
        prompt: str,
        num_inference_steps: int,
        guidance_scale: float,
        mask_blur: int,
        voxel_size: float,
        min_points_per_voxel: int,
        model_name: str,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        device = pointcloud.device

        # reuse stateless nodes
        proj_node      = ProjectPointCloud()
        outpaint_node  = OutpaintAnyProjection()
        depth_node     = DepthEstimatorNode()
        renorm_node    = DepthRenormalizer()
        depth2pc_node  = DepthToPointCloud()
        transform_node = TransformPointCloud()
        clean_node     = PointCloudCleaner()
        zdepth_node   = ZDepthToRayDepthNode()

        # initial clean on GPU
        with torch.no_grad():
            pointcloud, = clean_node.clean_pointcloud(
                pointcloud,
                voxel_size=voxel_size,
                min_points_per_voxel=min_points_per_voxel,
                width=4096,
                height=4096,
            )
        debug_img = torch.zeros((1, height, width, 3), device=device)
        debug_depth = torch.zeros((1, height, width, 1), device=device)
        enriched_pc = pointcloud
        # loop over trajectory (limit or full)
        for M in tqdm(trajectory[:15], desc="Enriching trajectory"):
            M_np  = M.cpu().numpy()
            M_inv = np.linalg.inv(M_np)

            # transform and select front points
            rotated, = transform_node.transform_pointcloud(enriched_pc, M_np)
            pc_front = rotated[rotated[:, 2] > 0]

            # clean front points
            pc_front, = clean_node.clean_pointcloud(
                pc_front,
                voxel_size=voxel_size,
                min_points_per_voxel=min_points_per_voxel,
                width=4096,
                height=4096,
            )

            # project to image + depth
            img, mask, depth_map = proj_node.project_pointcloud(
                pc_front,
                camera_type,
                horizontal_fov,
                width,
                height,
                point_size=3,
                return_inverse_depth=False,
            )
            debug_img = img
            # fill nan in depthmap with (-1)
            # outpaint missing regions
            hole_mask = (mask < 0.5).float()
            out_img, out_mask = outpaint_node.outpaint_any(
                img,
                input_projection    = camera_type,
                input_horiz_fov     = horizontal_fov,
                output_projection   = camera_type,
                output_horiz_fov    = horizontal_fov,
                output_width        = width,
                output_height       = height,
                patch_projection    = patch_projection,
                patch_horiz_fov     = patch_horiz_fov,
                patch_res           = patch_res,
                patch_phi           = patch_phi,
                patch_theta         = patch_theta,
                prompt              = prompt,
                num_inference_steps = num_inference_steps,
                cached              = False,
                guidance_scale      = guidance_scale,
                mask_blur           = mask_blur,
                mask                = hole_mask,
                debug               = False,
            )
            debug_img = out_img
            # estimate and renormalize depth
            nan_mask = torch.isnan(depth_map)
            # …and replace them with –1.0 (in-place)
            depth_map[nan_mask] = 0
            # clip from -1 to 1000
            depth_map = torch.clamp(depth_map, 0, 1000.0)
            new_depth, = depth_node.estimate_depth(out_img, model_name, depth_scale=1.0)
            new_depth, = zdepth_node.depth_to_ray_depth(
                new_depth,
                horizontal_fov,
            )
            # renormalize depth
            norm_depth, = renorm_node.renormalize_depth(
                new_depth,
                depth_map,
                depth_mask=(mask>=0.5)*1,
                guidance_mask=(mask<0.5)*1,
                use_inverse=False,
            )
            # median blur on depth
            k = 5
            d = norm_depth.permute(0,3,1,2)  # [B,1,H,W]
            pad = k//2
            pd = F.pad(d, (pad, pad, pad, pad), mode='reflect')
            patches = pd.unfold(2, k, 1).unfold(3, k, 1)
            patches = patches.contiguous().view(d.shape[0], d.shape[1], d.shape[2], d.shape[3], k*k)
            d, _ = patches.median(dim=-1)
            norm_depth = d.permute(0,2,3,1)  # [B,H,W,1]
            debug_depth = norm_depth*hole_mask.unsqueeze(0).unsqueeze(-1)+depth_map*(1-hole_mask.unsqueeze(0).unsqueeze(-1))

            # back to pointcloud
            pc_new, = depth2pc_node.depth_to_pointcloud(
                out_img,
                camera_type,
                horizontal_fov,
                depth_scale=1.0,
                invert_depth=False,
                depthmap=norm_depth,
                mask=hole_mask,
            )

            pc_world, = transform_node.transform_pointcloud(pc_new, M_inv)
            # enriched_pc is not rotated
            enriched_pc = torch.cat([enriched_pc, pc_world.to(device)], dim=0)
        return enriched_pc, debug_img, norm_depth
    

NODE_CLASS_MAPPINGS = {"FisheyeDepthEstimator": FisheyeDepthEstimator,
                       "PointcloudTrajectoryEnricher": PointcloudTrajectoryEnricher}