src/model/encoder/encoder_multy_scale.py

from dataclasses import dataclass
from typing import Literal, Optional, List

import torch
from einops import rearrange, repeat
from jaxtyping import Float
from torch import Tensor, nn
import MinkowskiEngine as ME
import torch.nn.init as init

from ...dataset.shims.patch_shim import apply_patch_shim
from ...dataset.types import BatchedExample, DataShim
from ...geometry.projection import sample_image_grid
from ..types import Gaussians

########形成高斯点#########
# from .common.gaussian_adapter_revise import GaussianAdapter_revise, GaussianAdapterCfg
from .common.guassian_adapter_depth import GaussianAdapter_depth, GaussianAdapterCfg


from .encoder import Encoder
from .visualization.encoder_visualizer_depthsplat_cfg import EncoderVisualizerDepthSplatCfg

import torchvision.transforms as T
import torch.nn.functional as F

from .unimatch.mv_unimatch import MultiViewUniMatch
from .unimatch.dpt_head import DPTHead

from .common.voxel_feature import project_features_to_3d, project_features_to_voxel, adapte_features_to_voxel, adapte_project_features_to_3d
from .common.me_fea import project_features_to_me

from ...geometry.projection import get_world_rays
from .common.sparse_net import SparseGaussianHead, SparseUNetWithAttention
from .common.mink_resnet import  MultiScaleSparseHead

from ...test.export_ply import save_point_cloud_to_ply

# import debugpy
# try:
#     # 5678 is the default attach port in the VS Code debug configurations. Unless a host and port are specified, host defaults to 127.0.0.1
#     debugpy.listen(("localhost", 9326))
#     print("Waiting for debugger attach")
#     debugpy.wait_for_client()
# except Exception as e:
#     pass


# 内存打印工具（轻量）
def print_mem(tag: str = ""):
    if not torch.cuda.is_available():
        print(f"[MEM] {tag} - no CUDA")
        return
    allocated = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    print(f"[MEM] {tag} | allocated={allocated:.1f} MB reserved={reserved:.1f} MB")

@dataclass
class EncoderDepthSplatCfg:
    name: Literal["depthsplat"]
    d_feature: int
    num_depth_candidates: int
    num_surfaces: int
    visualizer: EncoderVisualizerDepthSplatCfg
    gaussian_adapter: GaussianAdapterCfg
    gaussians_per_pixel: int
    unimatch_weights_path: str | None
    downscale_factor: int
    shim_patch_size: int
    multiview_trans_attn_split: int
    costvolume_unet_feat_dim: int
    costvolume_unet_channel_mult: List[int]
    costvolume_unet_attn_res: List[int]
    depth_unet_feat_dim: int
    depth_unet_attn_res: List[int]
    depth_unet_channel_mult: List[int]

    # mv_unimatch
    num_scales: int
    upsample_factor: int
    lowest_feature_resolution: int
    depth_unet_channels: int
    grid_sample_disable_cudnn: bool

    # depthsplat color branch
    large_gaussian_head: bool
    color_large_unet: bool
    init_sh_input_img: bool
    feature_upsampler_channels: int
    gaussian_regressor_channels: int

    # loss config
    supervise_intermediate_depth: bool
    return_depth: bool

    # only depth
    train_depth_only: bool

    # monodepth config
    monodepth_vit_type: str

    # multi-view matching
    local_mv_match: int


class EncoderDepthSplat_test(Encoder[EncoderDepthSplatCfg]):
    def __init__(self, cfg: EncoderDepthSplatCfg) -> None:
        super().__init__(cfg)

        self.depth_predictor = MultiViewUniMatch(
            num_scales=cfg.num_scales,
            upsample_factor=cfg.upsample_factor,
            lowest_feature_resolution=cfg.lowest_feature_resolution,
            vit_type=cfg.monodepth_vit_type,
            unet_channels=cfg.depth_unet_channels,
            grid_sample_disable_cudnn=cfg.grid_sample_disable_cudnn,
        )

        if self.cfg.train_depth_only:
            return

        # upsample features to the original resolution
        model_configs = {
            'vits': {'in_channels': 384, 'features': 64, 'out_channels': [48, 96, 192, 384]},
            'vitb': {'in_channels': 768, 'features': 96, 'out_channels': [96, 192, 384, 768]},
            'vitl': {'in_channels': 1024, 'features': 128, 'out_channels': [128, 256, 512, 1024]},
        }

        self.feature_upsampler = DPTHead(**model_configs[cfg.monodepth_vit_type],
                                        downsample_factor=cfg.upsample_factor,
                                        return_feature=True,
                                        num_scales=cfg.num_scales,
                                        )
        feature_upsampler_channels = model_configs[cfg.monodepth_vit_type]["features"]
        
        # gaussians adapter
        self.gaussian_adapter = GaussianAdapter_depth(cfg.gaussian_adapter)

        # concat(img, depth, match_prob, features)
        in_channels = 3 + 1 + 1 + feature_upsampler_channels
        channels = self.cfg.gaussian_regressor_channels

        # conv regressor
        modules = [
                    nn.Conv2d(in_channels, channels, 3, 1, 1),
                    nn.GELU(),
                    nn.Conv2d(channels, channels, 3, 1, 1),
                ]

        self.gaussian_regressor = nn.Sequential(*modules)

        # predict gaussian parameters: scale, q, sh, offset, opacity
        # num_gaussian_parameters = self.gaussian_adapter.d_in + 2 + 1  # 34 + 2(x,y) + 1(o) = 37
        num_gaussian_parameters = self.gaussian_adapter.d_in + 3 + 1  # 34 + 3(x,y,z) + 1(o) = 38
        # num_gaussian_parameters = self.gaussian_adapter.d_in + 1  # 34 +  + 1(o) = 35

        # concat(img, features, regressor_out, match_prob)
        in_channels = 3 + feature_upsampler_channels + channels + 1   
        
        #3D稀疏UNet
        self.spare_unet =SparseUNetWithAttention(
                            in_channels=in_channels, 
                            out_channels=in_channels, 
                            num_blocks=3, 
                            use_attention=False
                            )
        
        # 创建高斯头
        self.gaussian_head = SparseGaussianHead(in_channels, num_gaussian_parameters)
        
        # self.depth_fuse = nn.Sequential(nn.Conv2d(2, 4, kernel_size=1, padding=0),
        #     nn.ReLU(),
        #     nn.Conv2d(4, 1, kernel_size=1, padding=0)
        # )
        
        # # —— 初始化 depth_fuse 中所有 Conv2d 的权重和偏置 —— #
        # for m in self.depth_fuse.modules():
        #     if isinstance(m, nn.Conv2d):
        #         init.constant_(m.weight, 0.5)
        #         if m.bias is not None:
        #             init.constant_(m.bias, 0.5)
                    
                    
                    
        #######体素分辨率预测######
        # self.feature_extractor = nn.Sequential(
        #     nn.Conv2d(1, 2, 3, padding=1),
        #     nn.ReLU(),
        #     nn.MaxPool2d(2),  # 1/2
        #     nn.Conv2d(2, 4, 3, padding=1),
        #     nn.ReLU(),
        #     nn.MaxPool2d(2),  # 1/4
        #     nn.Conv2d(4, 8, 3, padding=1),
        #     nn.ReLU(),
        #     nn.AdaptiveAvgPool2d((1, 1))  # 全局特征聚合
        # )
        
        # # 回归预测头
        # self.regressor = nn.Sequential(
        #     nn.Flatten(),
        #     nn.Linear(8, 4),
        #     nn.ReLU(),
        #     nn.Linear(4, 1),
        #     nn.Sigmoid()  # 输出0-1范围
        # )
        
        
        # # —— 初始化 depth_fuse 中所有 Conv2d 的权重和偏置 —— #
        # for m in self.feature_extractor.modules():
        #     if isinstance(m, nn.Conv2d):
        #         init.constant_(m.weight, 0.5)
        #         if m.bias is not None:
        #             init.constant_(m.bias, 0.5)
        
        
        # 输出缩放参数 (0.01 + 0.04*sigmoid_output)
        self.scale = 0.04
        self.shift = 0.01

    def forward(
        self,
        context: dict,
        global_step: int,
        deterministic: bool = False,
        visualization_dump: Optional[dict] = None,
        scene_names: Optional[list] = None,
        ues_voxelnet: bool = True,
    ):
        device = context["image"].device
        b, v, _, h, w = context["image"].shape

        if v > 3:
            with torch.no_grad():
                xyzs = context["extrinsics"][:, :, :3, -1].detach()
                cameras_dist_matrix = torch.cdist(xyzs, xyzs, p=2)
                cameras_dist_index = torch.argsort(cameras_dist_matrix)

                cameras_dist_index = cameras_dist_index[:, :, :(self.cfg.local_mv_match + 1)]
        else:
            cameras_dist_index = None


        results_dict = self.depth_predictor(
            context["image"],
            attn_splits_list=[2],
            min_depth=1. / context["far"],
            max_depth=1. / context["near"],
            intrinsics=context["intrinsics"],
            extrinsics=context["extrinsics"],
            nn_matrix=cameras_dist_index,
        )

        # list of [B, V, H, W], with all the intermediate depths
        depth_preds = results_dict['depth_preds']
        
        depth = depth_preds[-1] 
        
        ########预测体素分辨率###########
        
        # if ues_voxelnet:
            # # [B*V, 1, H, W] -> [B*V, 16, 1, 1]
            # depth_features = self.feature_extractor(depth_fused)

            # # 回归预测 [B*V, 1]
            # resolution_raw = self.regressor(depth_features)
            
            # resolution_mean = resolution_raw.mean().unsqueeze(0)
            # # 缩放到目标范围 [0.01, 0.05]
            # voxel_resolution = self.shift + self.scale * resolution_mean
            # voxel_resolution = voxel_resolution.item()
            # print(f"预测体素分辨率: {voxel_resolution} m")
        # else:
        voxel_resolution =  0.02 #1mm体素
        # voxel_resolution =  0.01 #1cm体素
        
        
        
        if self.cfg.train_depth_only:
            # convert format
            # [B, V, H*W, 1, 1]
            depths = rearrange(depth, "b v h w -> b v (h w) () ()")

            if self.cfg.supervise_intermediate_depth and len(depth_preds) > 1:
                # supervise all the intermediate depth predictions
                num_depths = len(depth_preds)

                # [B, V, H*W, 1, 1]
                intermediate_depths = torch.cat(
                    depth_preds[:(num_depths - 1)], dim=0)
                intermediate_depths = rearrange(
                    intermediate_depths, "b v h w -> b v (h w) () ()")

                # concat in the batch dim
                depths = torch.cat((intermediate_depths, depths), dim=0)

                b *= num_depths

            # return depth prediction for supervision
            depths = rearrange(
                depths, "b v (h w) srf s -> b v h w srf s", h=h, w=w
            ).squeeze(-1).squeeze(-1)
            # print(depths.shape)  # [B, V, H, W]

            return {
                "gaussians": None,
                "depths": depths
            }

        # features [BV, C, H, W]
        features = self.feature_upsampler(results_dict["features_mono_intermediate"],
                                          cnn_features=results_dict["features_cnn_all_scales"][::-1],
                                          mv_features=results_dict["features_mv"][
                                          0] if self.cfg.num_scales == 1 else results_dict["features_mv"][::-1]
                                          )
        
        # match prob from softmax
        # [BV, D, H, W] in feature resolution
        match_prob = results_dict['match_probs'][-1]
        match_prob = torch.max(match_prob, dim=1, keepdim=True)[
            0]  # [BV, 1, H, W]
        match_prob = F.interpolate(
            match_prob, size=depth.shape[-2:], mode='nearest')
        
        
        # unet input [BV, C, H, W]  [6, 101, 256, 448]
        concat = torch.cat((
            rearrange(context["image"], "b v c h w -> (b v) c h w"),
            rearrange(depth, "b v h w -> (b v) () h w"),
            match_prob,
            features,
        ), dim=1)
        # [BV, C, H, W]
        out = self.gaussian_regressor(concat)
        concat = [out,
                    rearrange(context["image"],
                            "b v c h w -> (b v) c h w"),
                    features,
                    match_prob]
        # [BV, C, H, W]   [6, 164, 256, 448]    
        out = torch.cat(concat, dim=1) 
  
        sparse_input, aggregated_points, counts = project_features_to_me(
                context["intrinsics"],
                context["extrinsics"],
                out,
                depth=depth, 
                voxel_resolution=voxel_resolution,
                b=b, v=v
                )

        sparse_out = self.spare_unet(sparse_input)   #3D稀疏UNet
        # refine with residual
        if torch.equal(sparse_out.C, sparse_input.C) and sparse_out.F.shape[1] == sparse_input.F.shape[1]:
            # 创建新的特征张量
            new_features = sparse_out.F + sparse_input.F
            
            # 创建新的 SparseTensor
            sparse_out_with_residual = ME.SparseTensor(
                features=new_features,
                coordinate_map_key=sparse_out.coordinate_map_key,
                coordinate_manager=sparse_out.coordinate_manager
            )
        else:
            # 处理坐标不一致的情况
            print("警告：输入和输出坐标不一致，跳过残差连接")
            sparse_out_with_residual = sparse_out

        #([1, 128, 80, 80, 80])   -> [N, 38]
        gaussians = self.gaussian_head(sparse_out_with_residual)

        # 及时释放不再需要的变量
        del sparse_out_with_residual,sparse_out,sparse_input,new_features

        
        # [B, V, H*W, 1, 1]
        depths = rearrange(depth, "b v h w -> b v (h w) () ()")
        
        
        # 输出也是稀疏张量
        print(f"输出稀疏张量: {gaussians.F.shape[0]}个体素")
        
        gaussian_params = gaussians.F.unsqueeze(0).unsqueeze(0) #[N, 38] -> [1, 1, N, 38]
        
        # 分离不透明度和其他参数
        opacities = gaussian_params[..., :1].sigmoid().unsqueeze(-1)  #[1, 1, 256000, 1, 1]
        raw_gaussians = gaussian_params[..., 1:]    #[1, 1, 256000, 37]
        raw_gaussians = rearrange(
        raw_gaussians,
        "... (srf c) -> ... srf c",
        srf=self.cfg.num_surfaces,
        )
        
        try:
            # 将raw_gaussians转换成gaussians参数
            gaussians = self.gaussian_adapter.forward(
                extrinsics = context["extrinsics"],
                intrinsics = context["intrinsics"],
                opacities = opacities,
                raw_gaussians = rearrange(raw_gaussians,"b v r srf c -> b v r srf () c"),
                input_images =rearrange(context["image"], "b v c h w -> (b v) c h w"),   #[6, 3, 256, 448]
                depth = depth,
                coordidate = gaussians.C,
                points = aggregated_points,
                voxel_resolution = voxel_resolution
            )
        except Exception as e:
            import traceback; traceback.print_exc()
            raise

        

        if self.cfg.supervise_intermediate_depth and len(depth_preds) > 1:
            intermediate_depth = depth_preds[0]
            #得到voxel_feature
            intermediate_voxel_feature, median_points, counts = project_features_to_me(
                context["intrinsics"],
                context["extrinsics"],
                out,
                depth=intermediate_depth, 
                voxel_resolution=voxel_resolution,
                b=b, v=v
                )
            #############################经过U-net进行finture#######################
            intermediate_out = self.spare_unet(intermediate_voxel_feature)   #3D稀疏UNet
            # refine with residual
            if torch.equal(intermediate_out.C, intermediate_voxel_feature.C) and intermediate_out.F.shape[1] == intermediate_voxel_feature.F.shape[1]:
                # 创建新的特征张量
                new_inter_features = intermediate_out.F + intermediate_voxel_feature.F
                
                # 创建新的 SparseTensor
                intermedian_out_with_residual = ME.SparseTensor(
                    features=new_inter_features,
                    coordinate_map_key=intermediate_voxel_feature.coordinate_map_key,
                    coordinate_manager=intermediate_voxel_feature.coordinate_manager
                )
            else:
                # 处理坐标不一致的情况
                print("警告：输入和输出坐标不一致，跳过残差连接")
                intermedian_out_with_residual = intermediate_voxel_feature

            #([1, 128, 80, 80, 80])   -> [N, 38]
            intermediate_gaussians = self.gaussian_head(intermedian_out_with_residual)

            # 及时释放不再需要的变量
            del intermediate_voxel_feature,intermediate_out,intermedian_out_with_residual

            # intermediate_gaussians = self.gaussian_head(intermediate_voxel_feature)  #[N, 38]
            # print_mem("after media_depth gaussian_head")
            
            gaussian_params = intermediate_gaussians.F.unsqueeze(0).unsqueeze(0) #[N, 38] -> [1, 1, N, 38]
       
            # 分离不透明度和其他参数
            intermediate_opacities = gaussian_params[..., :1].sigmoid().unsqueeze(-1)  #[1, 1, 256000, 1, 1]
            intermediate_raw_gaussians = gaussian_params[..., 1:]    #[1, 1, 256000, 37]
            intermediate_raw_gaussians = rearrange(
            intermediate_raw_gaussians,
            "... (srf c) -> ... srf c",
            srf=self.cfg.num_surfaces,
            )
            
 
            # 将raw_gaussians转换成gaussians参数            
            intermediate_gaussians = self.gaussian_adapter.forward(
                extrinsics = context["extrinsics"],
                intrinsics = context["intrinsics"],
                opacities = intermediate_opacities,
                raw_gaussians = rearrange(intermediate_raw_gaussians,"b v r srf c -> b v r srf () c"),
                input_images =rearrange(context["image"], "b v c h w -> (b v) c h w"),   #[6, 3, 256, 448]
                depth = intermediate_depth,
                coordidate = intermediate_gaussians.C,
                points = median_points,
                voxel_resolution = voxel_resolution
            )
        
            intermediate_gaussians = Gaussians(
                rearrange(
                    intermediate_gaussians.means,   #[2, 1, 256000, 1, 1, 3]
                    "b v r srf spp xyz -> b (v r srf spp) xyz",   #[2, 256000, 3]
                ),
                rearrange(
                    intermediate_gaussians.covariances,  #[2, 1, 256000, 1, 1, 3, 3]
                    "b v r srf spp i j -> b (v r srf spp) i j",  #[2, 256000, 3, 3]
                ),
                rearrange(
                    intermediate_gaussians.harmonics, #[2, 1, 256000, 1, 1, 3, 9]
                    "b v r srf spp c d_sh -> b (v r srf spp) c d_sh",  #[2, 256000, 3, 9]
                ),
                rearrange(
                    intermediate_gaussians.opacities,  #[2, 1, 256000, 1, 1]
                    "b v r srf spp -> b (v r srf spp)",  #[2, 256000]
                ),
            )
        else:
            intermediate_gaussians = None


        
        gaussians = Gaussians(
            rearrange(
                gaussians.means,   #[2, 1, 256000, 1, 1, 3]
                "b v r srf spp xyz -> b (v r srf spp) xyz",   #[2, 256000, 3]
            ),
            rearrange(
                gaussians.covariances,  #[2, 1, 256000, 1, 1, 3, 3]
                "b v r srf spp i j -> b (v r srf spp) i j",  #[2, 256000, 3, 3]
            ),
            rearrange(
                gaussians.harmonics, #[2, 1, 256000, 1, 1, 3, 9]
                "b v r srf spp c d_sh -> b (v r srf spp) c d_sh",  #[2, 256000, 3, 9]
            ),
            rearrange(
                gaussians.opacities,  #[2, 1, 256000, 1, 1]
                "b v r srf spp -> b (v r srf spp)",  #[2, 256000]
            ),
        )

        # print_mem("end forward")
        if self.cfg.return_depth:
            # return depth prediction for supervision
            # depths  = torch.cat(depth_preds, dim=0)
            depths = rearrange(
                depths, "b v (h w) srf s -> b v h w srf s", h=h, w=w
            ).squeeze(-1).squeeze(-1)
            
            # print(depths.shape)  # [B, V, H, W]  [2, 6, 256, 448]
            if intermediate_gaussians is not None:
                return {
                    "gaussians": gaussians,
                    "depths": depths,
                    "intermediate_gaussians": intermediate_gaussians
                }
            else:
                return {
                    "gaussians": gaussians,
                    "depths": depths,
                }

        return gaussians

    def get_data_shim(self) -> DataShim:
        def data_shim(batch: BatchedExample) -> BatchedExample:
            batch = apply_patch_shim(
                batch,
                patch_size=self.cfg.shim_patch_size
                * self.cfg.downscale_factor,
            )

            return batch

        return data_shim

    @property
    def sampler(self):
        return None