Upload folder using huggingface_hub

config_slat_flow_128to512_pointnet_head.yaml

@@ -0,0 +1,122 @@
+model:
+  pred_direction: false
+  relative_embed: true
+  using_attn: false
+  add_block_embed: true
+  multires: 12
+  
+  embed_dim: 1024
+  in_channels: 1024
+  model_channels: 384
+  latent_dim: 16
+
+  block_size: 16
+  pos_encoding: 'nerf'
+  attn_first: false
+
+  add_edge_glb_feats: true
+  add_direction: false
+  
+  encoder_blocks:
+  - in_channels: 1024
+    model_channels: 512
+    num_blocks: 8
+    num_heads: 8
+    out_channels: 512
+
+  decoder_blocks_edge:
+  - in_channels: 512
+    model_channels: 512
+    num_blocks: 0
+    num_heads: 0
+    out_channels: 256
+    resolution: 128
+  - in_channels: 256
+    model_channels: 256
+    num_blocks: 0
+    num_heads: 0
+    out_channels: 128
+    resolution: 256
+  # - in_channels: 64
+  #   model_channels: 64
+  #   num_blocks: 0
+  #   num_heads: 0
+  #   out_channels: 32
+  #   resolution: 512
+  
+  decoder_blocks_vtx:
+  - in_channels: 512
+    model_channels: 512
+    num_blocks: 0
+    num_heads: 0
+    out_channels: 256
+    resolution: 128
+  - in_channels: 256
+    model_channels: 256
+    num_blocks: 0
+    num_heads: 0
+    out_channels: 128
+    resolution: 256
+  # - in_channels: 64
+  #   model_channels: 64
+  #   num_blocks: 0
+  #   num_heads: 0
+  #   out_channels: 32
+  #   resolution: 512
+
+"t_schedule": 
+    "name": "logitNormal"
+    "args": 
+        "mean": 1.0
+        "std": 1.0
+    
+"sigma_min": 1.e-5
+
+training:
+  batch_size: 1
+  lr: 1.e-4
+  step_size: 20
+  gamma: 0.95
+  save_every: 500
+  start_epoch: 0
+  max_epochs: 300
+  num_workers: 32
+  
+  output_dir: /home/tiger/yy/src/checkpoints/output_slat_flow_matching_active/40w_128to512_head_rope
+  clip_model_path: None
+  dinov2_model_path: None
+
+  vae_path: /home/tiger/yy/src/checkpoints/vae/unique_files_glb_under6000face_2degree_30ratio_0.01/shapenet_bs2_128to512_wolabel_dir_sorted_dora_small_lowlr/checkpoint_epoch1_batch12000_loss0.1140.pt
+  denoiser_checkpoint_path: /home/tiger/yy/src/checkpoints/output_slat_flow_matching_active/40w_128to512/checkpoint_step143000_loss0_736924.pt
+
+
+dataset:
+  path: /home/tiger/yy/src/unique_files_glb_under6000face_2degree_30ratio_0.01
+  cache_dir: /home/tiger/yy/src/dataset_cache/unique_files_glb_under6000face_2degree_30ratio_0.01
+
+  renders_dir: None
+  filter_active_voxels: true
+  cache_filter_path: /home/tiger/yy/src/40w_2000-100000edge_2000-100000active.txt
+  
+  base_resolution: 1024
+  min_resolution: 128
+
+  n_train_samples: 1024
+  sample_type: dora
+
+flow:
+  "resolution": 128
+  "in_channels": 16
+  "out_channels": 16
+  "model_channels": 768
+  "cond_channels": 1024
+  "num_blocks": 12
+  "num_heads": 12
+  "mlp_ratio": 4
+  "patch_size": 2
+  "num_io_res_blocks": 2
+  "io_block_channels": [128]
+  "pe_mode": "rope"
+  "qk_rms_norm": true
+  "qk_rms_norm_cross": false
+  "use_fp16": false

config_slat_flow_128to512_pointnet_head_test.yaml

@@ -0,0 +1,123 @@
+model:
+  pred_direction: false
+  relative_embed: true
+  using_attn: false
+  add_block_embed: true
+  multires: 12
+  
+  embed_dim: 1024
+  in_channels: 1024
+  model_channels: 384
+  latent_dim: 16
+
+  block_size: 16
+  pos_encoding: 'nerf'
+  attn_first: false
+
+  add_edge_glb_feats: true
+  add_direction: false
+  
+  encoder_blocks:
+  - in_channels: 1024
+    model_channels: 512
+    num_blocks: 8
+    num_heads: 8
+    out_channels: 512
+
+  decoder_blocks_edge:
+  - in_channels: 512
+    model_channels: 512
+    num_blocks: 0
+    num_heads: 0
+    out_channels: 256
+    resolution: 128
+  - in_channels: 256
+    model_channels: 256
+    num_blocks: 0
+    num_heads: 0
+    out_channels: 128
+    resolution: 256
+  # - in_channels: 64
+  #   model_channels: 64
+  #   num_blocks: 0
+  #   num_heads: 0
+  #   out_channels: 32
+  #   resolution: 512
+  
+  decoder_blocks_vtx:
+  - in_channels: 512
+    model_channels: 512
+    num_blocks: 0
+    num_heads: 0
+    out_channels: 256
+    resolution: 128
+  - in_channels: 256
+    model_channels: 256
+    num_blocks: 0
+    num_heads: 0
+    out_channels: 128
+    resolution: 256
+  # - in_channels: 64
+  #   model_channels: 64
+  #   num_blocks: 0
+  #   num_heads: 0
+  #   out_channels: 32
+  #   resolution: 512
+
+"t_schedule": 
+    "name": "logitNormal"
+    "args": 
+        "mean": 1.0
+        "std": 1.0
+    
+"sigma_min": 1.e-5
+
+training:
+  batch_size: 1
+  lr: 1.e-4
+  step_size: 20
+  gamma: 0.95
+  save_every: 500
+  start_epoch: 0
+  max_epochs: 300
+  num_workers: 32
+  
+  output_dir: /home/tiger/yy/src/checkpoints/output_slat_flow_matching_active/40w_128to512_head_rope
+  clip_model_path: None
+  dinov2_model_path: None
+
+  vae_path: /home/tiger/yy/src/checkpoints/vae/unique_files_glb_under6000face_2degree_30ratio_0.01/shapenet_bs2_128to512_wolabel_dir_sorted_dora_small_lowlr/checkpoint_epoch1_batch12000_loss0.1140.pt
+  denoiser_checkpoint_path: /home/tiger/yy/src/checkpoints/output_slat_flow_matching_active/40w_128to512/checkpoint_step143000_loss0_736924.pt
+
+
+dataset:
+  path: /home/tiger/yy/src/unique_files_glb_under6000face_2degree_30ratio_0.01
+  path: /home/tiger/yy/src/trellis_clean_mesh/mesh_data
+  cache_dir: /home/tiger/yy/src/dataset_cache/unique_files_glb_under6000face_2degree_30ratio_0.01
+
+  renders_dir: None
+  filter_active_voxels: false
+  cache_filter_path: /home/tiger/yy/src/40w_2000-100000edge_2000-100000active.txt
+  
+  base_resolution: 1024
+  min_resolution: 128
+
+  n_train_samples: 1024
+  sample_type: dora
+
+flow:
+  "resolution": 128
+  "in_channels": 16
+  "out_channels": 16
+  "model_channels": 768
+  "cond_channels": 1024
+  "num_blocks": 12
+  "num_heads": 12
+  "mlp_ratio": 4
+  "patch_size": 2
+  "num_io_res_blocks": 2
+  "io_block_channels": [128]
+  "pe_mode": "rope"
+  "qk_rms_norm": true
+  "qk_rms_norm_cross": false
+  "use_fp16": false

test_slat_flow_128to512_pointnet_head.py

@@ -0,0 +1,404 @@
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+import yaml
+import time
+from datetime import datetime
+from torch.utils.data import DataLoader
+from functools import partial
+import torch.nn.functional as F
+from torch.amp import GradScaler, autocast
+from typing import *
+from transformers import CLIPTextModel, AutoTokenizer, CLIPTextConfig, Dinov2Model, AutoImageProcessor, Dinov2Config
+import torch
+import re
+from utils import load_pretrained_woself
+
+from dataset_triposf_head import VoxelVertexDataset_edge, collate_fn_pointnet
+from triposf.models.triposf_vae.VoxelFeatureVAE_edge_woself_128to1024_decoder_head import VoxelVAE
+from vertex_encoder import VoxelFeatureEncoder_active_pointnet
+
+from trellis.models.structured_latent_flow import SLatFlowModel
+from trellis.trainers.flow_matching.sparse_flow_matching_alone import SparseFlowMatchingTrainer
+
+from trellis.pipelines.samplers import FlowEulerSampler
+from safetensors.torch import load_file
+import open3d as o3d
+from PIL import Image
+
+from triposf.modules.sparse.basic import SparseTensor
+from trellis.modules.sparse.basic import SparseTensor as SparseTensor_trellis
+
+from triposf.modules.utils import DiagonalGaussianDistribution
+
+from sklearn.decomposition import PCA
+import trimesh
+import torchvision.transforms as transforms
+
+# --- Helper Functions ---
+def save_colored_ply(points, colors, filename):
+    if len(points) == 0:
+        print(f"[Warning] No points to save for {filename}")
+        return
+    # Ensure colors are uint8
+    if colors.max() <= 1.0:
+        colors = (colors * 255).astype(np.uint8)
+    colors = colors.astype(np.uint8)
+    
+    # Add Alpha if missing
+    if colors.shape[1] == 3:
+        colors = np.hstack([colors, np.full((len(colors), 1), 255, dtype=np.uint8)])
+        
+    cloud = trimesh.PointCloud(points, colors=colors)
+    cloud.export(filename)
+    print(f"Saved colored point cloud to {filename}")
+
+def normalize_to_rgb(features_3d):
+    min_vals = features_3d.min(axis=0)
+    max_vals = features_3d.max(axis=0)
+    range_vals = max_vals - min_vals
+    range_vals[range_vals == 0] = 1 
+    normalized = (features_3d - min_vals) / range_vals
+    return (normalized * 255).astype(np.uint8)
+
+class SLatFlowMatchingTrainer(SparseFlowMatchingTrainer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.cfg = kwargs.pop('cfg', None)
+        if self.cfg is None:
+            raise ValueError("Configuration dictionary 'cfg' must be provided.")
+        
+        self.sampler = FlowEulerSampler(sigma_min=1.e-5)
+        self.device = torch.device("cuda")
+
+        # Based on PointNet Encoder setting
+        self.resolution = 128 
+
+        self.condition_type = 'image'
+        self.is_cond = False
+        
+        self.img_res = 518
+        self.feature_dim = self.cfg['model']['latent_dim']
+
+        self._init_components(
+            clip_model_path=self.cfg['training'].get('clip_model_path', None),
+            dinov2_model_path=self.cfg['training'].get('dinov2_model_path', None),
+            vae_path=self.cfg['training']['vae_path'],
+        )
+
+        # Classifier head removed as it is not part of the Active Voxel pipeline
+
+    def _load_denoiser(self, denoiser_checkpoint_path):
+        path = denoiser_checkpoint_path
+        if not path or not os.path.isfile(path):
+            print("No valid checkpoint path provided for fine-tuning. Starting from scratch.")
+            return
+
+        print(f"Loading checkpoint from: {path}")
+        checkpoint = torch.load(path, map_location=self.device)
+
+        try:
+            denoiser_state_dict = checkpoint['denoiser']
+            # Handle DDP prefix
+            if next(iter(denoiser_state_dict)).startswith('module.'):
+                denoiser_state_dict = {k[7:]: v for k, v in denoiser_state_dict.items()}
+            
+            self.denoiser.load_state_dict(denoiser_state_dict)
+            print("Denoiser weights loaded successfully.")
+        except KeyError:
+            print("[WARN] 'denoiser' key not found in checkpoint. Skipping.")
+        except Exception as e:
+            print(f"[ERROR] Failed to load denoiser state_dict: {e}")
+
+    def _init_components(self,
+            clip_model_path=None,
+            dinov2_model_path=None,
+            vae_path=None,
+        ):
+        
+        # 1. Initialize PointNet Voxel Encoder (Matches Training)
+        self.voxel_encoder = VoxelFeatureEncoder_active_pointnet(
+            in_channels=15, 
+            hidden_dim=256, 
+            out_channels=1024, 
+            scatter_type='mean', 
+            n_blocks=5, 
+            resolution=128,
+            add_label=False,
+        ).to(self.device)
+
+        # 2. Initialize VAE
+        self.vae = VoxelVAE(
+            in_channels=self.cfg['model']['in_channels'],
+            latent_dim=self.cfg['model']['latent_dim'],
+            encoder_blocks=self.cfg['model']['encoder_blocks'],
+            decoder_blocks_vtx=self.cfg['model']['decoder_blocks_vtx'],
+            decoder_blocks_edge=self.cfg['model']['decoder_blocks_edge'],
+            num_heads=8,
+            num_head_channels=64,
+            mlp_ratio=4.0,
+            attn_mode="swin",
+            window_size=8,
+            pe_mode="ape",
+            use_fp16=False,
+            use_checkpoint=False,
+            qk_rms_norm=False,
+            using_subdivide=True,
+            using_attn=self.cfg['model']['using_attn'],
+            attn_first=self.cfg['model'].get('attn_first', True),
+            pred_direction=self.cfg['model'].get('pred_direction', False),
+        ).to(self.device)
+
+        # 3. Initialize Dataset with collate_fn_pointnet
+        self.dataset = VoxelVertexDataset_edge(
+            root_dir=self.cfg['dataset']['path'],
+            base_resolution=self.cfg['dataset']['base_resolution'],
+            min_resolution=self.cfg['dataset']['min_resolution'],
+            cache_dir=self.cfg['dataset']['cache_dir'],
+            renders_dir=self.cfg['dataset']['renders_dir'],
+
+            process_img=False,
+
+            active_voxel_res=128,
+            filter_active_voxels=self.cfg['dataset']['filter_active_voxels'],
+            cache_filter_path=self.cfg['dataset']['cache_filter_path'],
+            sample_type=self.cfg['dataset'].get('sample_type', 'dora'),
+        )
+        
+        self.dataloader = DataLoader(
+            self.dataset,
+            batch_size=1,
+            shuffle=True,
+            collate_fn=partial(collate_fn_pointnet,), # Critical Change
+            num_workers=0,
+            pin_memory=True,
+            persistent_workers=False,
+        )
+        
+        # 4. Load Pretrained Weights
+        # Assuming vae_path contains 'voxel_encoder' and 'vae'
+        print(f"Loading VAE/Encoder from {vae_path}")
+        ckpt = torch.load(vae_path, map_location='cpu')
+        
+        # Load VAE
+        if 'vae' in ckpt:
+            self.vae.load_state_dict(ckpt['vae'], strict=False)
+        else:
+            self.vae.load_state_dict(ckpt) # Fallback
+            
+        # Load Encoder
+        if 'voxel_encoder' in ckpt:
+            self.voxel_encoder.load_state_dict(ckpt['voxel_encoder'])
+        else:
+            print("[WARN] 'voxel_encoder' not found in checkpoint, random init (BAD for inference).")
+
+        self.voxel_encoder.eval()
+        self.vae.eval()
+
+        # 5. Initialize Conditioning Model
+        if self.condition_type == 'text':
+            self.tokenizer = AutoTokenizer.from_pretrained(clip_model_path)
+            self.condition_model = CLIPTextModel.from_pretrained(clip_model_path)
+        elif self.condition_type == 'image':
+            model_name = 'dinov2_vitl14_reg'
+            local_repo_path = "/home/tiger/yy/src/gemini/user/private/zhaotianhao/dinov2_resources/dinov2-main"
+            weights_path = "/home/tiger/yy/src/gemini/user/private/zhaotianhao/dinov2_resources/dinov2_vitl14_reg4_pretrain.pth"
+
+            dinov2_model = torch.hub.load(
+                repo_or_dir=local_repo_path,
+                model=model_name,
+                source='local',
+                pretrained=False 
+            )
+            self.condition_model = dinov2_model
+            self.condition_model.load_state_dict(torch.load(weights_path))
+            
+            self.image_cond_model_transform = transforms.Compose([
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            ])
+        else:
+            raise ValueError(f"Unsupported condition type: {self.condition_type}")
+
+        self.condition_model.to(self.device).eval()
+
+    @torch.no_grad()
+    def encode_image(self, images) -> torch.Tensor:
+        if isinstance(images, torch.Tensor):
+            batch_tensor = images.to(self.device)
+        elif isinstance(images, list):
+            assert all(isinstance(i, Image.Image) for i in images)
+            image = [i.resize((518, 518), Image.LANCZOS) for i in images]
+            image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
+            image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
+            batch_tensor = torch.stack(image).to(self.device)
+        else:
+            raise ValueError(f"Unsupported type of image: {type(images)}")
+
+        if batch_tensor.shape[-2:] != (518, 518):
+             batch_tensor = F.interpolate(batch_tensor, (518, 518), mode='bicubic', align_corners=False)
+
+        features = self.condition_model(batch_tensor, is_training=True)['x_prenorm']
+        patchtokens = F.layer_norm(features, features.shape[-1:])
+        return patchtokens
+        
+    def process_batch(self, batch):
+        preprocessed_images = batch['image'] 
+        cond_ = self.encode_image(preprocessed_images)
+        return cond_
+
+    def eval(self):
+        # Unconditional Setup
+        if self.is_cond == False:
+            if self.condition_type == 'text':
+                txt = ['']
+                encoding = self.tokenizer(txt, max_length=77, padding='max_length', truncation=True, return_tensors='pt')
+                tokens = encoding['input_ids'].to(self.device)
+                with torch.no_grad():
+                    cond_ = self.condition_model(input_ids=tokens).last_hidden_state
+            else:
+                blank_img = Image.fromarray(np.zeros((self.img_res, self.img_res, 3), dtype=np.uint8))
+                with torch.no_grad():
+                    dummy_cond = self.encode_image([blank_img])
+                    cond_ = torch.zeros_like(dummy_cond)
+                print(f"Generated unconditional image prompt (zero tensor) with shape: {cond_.shape}")
+
+        self.denoiser.eval()
+
+        # Load Denoiser Checkpoint
+        # Update this path to your ACTIVE VOXEL trained checkpoint
+        checkpoint_path = '/home/tiger/yy/src/checkpoints/output_slat_flow_matching_active/40w_128to512_head_rope/checkpoint_step143500_loss0_766792.pt'
+        
+        self._load_denoiser(checkpoint_path)
+
+        filename = os.path.basename(checkpoint_path)
+        match = re.search(r'step(\d+)', filename)
+        step_str = match.group(1) if match else "eval"
+        save_dir = os.path.join(os.path.dirname(checkpoint_path), f"{step_str}_sample_active_vis_42seed_trellis")
+        # save_dir = os.path.join(os.path.dirname(checkpoint_path), f"{step_str}_sample_active_vis_42seed_40w_train")
+        os.makedirs(save_dir, exist_ok=True)
+        print(f"Results will be saved to: {save_dir}")
+        
+        for i, batch in enumerate(self.dataloader):
+            if i > 50: exit() # Visualize first 10
+            
+            if self.is_cond and self.condition_type == 'image':
+                cond_ = self.process_batch(batch)
+
+            if cond_.shape[0] != 1:
+                cond_ = cond_.expand(batch['active_voxels_128'].shape[0], -1, -1).contiguous().to(self.device)
+            else:
+                cond_ = cond_.to(self.device)
+
+            # --- Data Retrieval (Matches collate_fn_pointnet) ---
+            point_cloud = batch['point_cloud_128'].to(self.device)
+            active_coords = batch['active_voxels_128'].to(self.device) # [N, 4]
+            
+            with autocast(device_type='cuda', dtype=torch.bfloat16):
+                with torch.no_grad():
+                    # 1. Encode Ground Truth Latents
+                    active_voxel_feats = self.voxel_encoder(
+                        p=point_cloud,
+                        sparse_coords=active_coords,
+                        res=128,
+                        bbox_size=(-0.5, 0.5),
+                    )
+
+                    sparse_input = SparseTensor(
+                        feats=active_voxel_feats,
+                        coords=active_coords.int()
+                    )
+
+                    # Encode to get GT distribution
+                    gt_latents, posterior = self.vae.encode(sparse_input)
+                    
+                    print(f"Batch {i}: Active voxels: {active_coords.shape[0]}")
+
+                    # 2. Generation / Sampling
+                    # Generate noise on the SAME active coordinates
+                    noise = SparseTensor_trellis(
+                        coords=active_coords.int(),
+                        feats=torch.randn(
+                            active_coords.shape[0], 
+                            self.feature_dim,
+                            device=self.device,
+                        ),
+                    )
+
+                    sample_results = self.sampler.sample(
+                        model=self.denoiser.float(),
+                        noise=noise.to(self.device).float(),
+                        cond=cond_.to(self.device).float(),
+                        steps=50,
+                        rescale_t=1.0,
+                        verbose=True,
+                    )
+
+                    generated_sparse_tensor = sample_results.samples
+                    generated_coords = generated_sparse_tensor.coords
+                    generated_features = generated_sparse_tensor.feats
+
+                    print('Gen features mean:', generated_features.mean().item(), 'std:', generated_features.std().item())
+                    print('GT features mean:', gt_latents.feats.mean().item(), 'std:', gt_latents.feats.std().item())
+                    print('MSE:', F.mse_loss(generated_features, gt_latents.feats).item())
+
+            # --- Visualization (PCA) ---
+            gt_feats_np = gt_latents.feats.detach().cpu().numpy()
+            gen_feats_np = generated_features.detach().cpu().numpy()
+            coords_np = active_coords[:, 1:4].detach().cpu().numpy() # x, y, z
+
+            print("Visualizing features using PCA...")
+            pca = PCA(n_components=3)
+            
+            # Fit PCA on GT, transform both
+            pca.fit(gt_feats_np)
+            gt_feats_3d = pca.transform(gt_feats_np)
+            gen_feats_3d = pca.transform(gen_feats_np)
+            
+            gt_colors = normalize_to_rgb(gt_feats_3d)
+            gen_colors = normalize_to_rgb(gen_feats_3d)
+
+            # Save PLYs
+            save_colored_ply(coords_np, gt_colors, os.path.join(save_dir, f"batch_{i}_gt_pca.ply"))
+            save_colored_ply(coords_np, gen_colors, os.path.join(save_dir, f"batch_{i}_gen_pca.ply"))
+
+            # Save Tensors for further analysis
+            torch.save(gt_latents, os.path.join(save_dir, f"gt_latent_{i}.pt"))
+
+            torch.save(batch, os.path.join(save_dir, f"gt_data_batch_{i}.pt"))
+            torch.save(sample_results.samples, os.path.join(save_dir, f"sample_latent_{i}.pt"))
+
+if __name__ == '__main__':
+    torch.manual_seed(42)
+    config_path = "/home/tiger/yy/src/Michelangelo-master/config_slat_flow_128to512_pointnet_head_test.yaml"
+    with open(config_path) as f:
+        cfg = yaml.safe_load(f)
+
+    # Initialize Model on CPU first
+    diffusion_model = SLatFlowModel(
+        resolution=cfg['flow']['resolution'],
+        in_channels=cfg['flow']['in_channels'],
+        out_channels=cfg['flow']['out_channels'],
+        model_channels=cfg['flow']['model_channels'],
+        cond_channels=cfg['flow']['cond_channels'],
+        num_blocks=cfg['flow']['num_blocks'],
+        num_heads=cfg['flow']['num_heads'],
+        mlp_ratio=cfg['flow']['mlp_ratio'],
+        patch_size=cfg['flow']['patch_size'],
+        num_io_res_blocks=cfg['flow']['num_io_res_blocks'],
+        io_block_channels=cfg['flow']['io_block_channels'],
+        pe_mode=cfg['flow']['pe_mode'],
+        qk_rms_norm=cfg['flow']['qk_rms_norm'],
+        qk_rms_norm_cross=cfg['flow']['qk_rms_norm_cross'],
+        use_fp16=cfg['flow'].get('use_fp16', False),
+    ).to("cuda" if torch.cuda.is_available() else "cpu")
+    
+    trainer = SLatFlowMatchingTrainer(
+        denoiser=diffusion_model,
+        t_schedule=cfg['t_schedule'],
+        sigma_min=cfg['sigma_min'],
+        cfg=cfg,
+    )
+    
+    trainer.eval()

test_slat_flow_128to512_pointnet_head_tomesh.py

@@ -0,0 +1,1630 @@
+import os
+import yaml
+import torch
+import numpy as np
+import random
+from tqdm import tqdm
+from collections import defaultdict
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+from torch.utils.data import DataLoader, Subset
+from triposf.modules.sparse.basic import SparseTensor
+
+from triposf.models.triposf_vae.VoxelFeatureVAE_edge_woself_128to1024_decoder_head import VoxelVAE
+
+from vertex_encoder import VoxelFeatureEncoder_edge, VoxelFeatureEncoder_vtx, VoxelFeatureEncoder_active, VoxelFeatureEncoder_active_pointnet, ConnectionHead
+from utils import load_pretrained_woself
+
+from dataset_triposf_head import VoxelVertexDataset_edge, collate_fn_pointnet
+
+
+from functools import partial
+import itertools
+from typing import List, Tuple, Set
+from collections import OrderedDict
+from scipy.spatial import cKDTree
+from sklearn.neighbors import KDTree
+
+import trimesh
+
+import torch
+import torch.nn.functional as F
+import time
+
+from sklearn.decomposition import PCA
+import matplotlib.pyplot as plt
+
+import networkx as nx
+
+def predict_mesh_connectivity(
+    connection_head, 
+    vtx_feats, 
+    vtx_coords, 
+    batch_size=10000, 
+    threshold=0.5,
+    k_neighbors=64, # 限制每个点只检测最近的 K 个邻居，设为 -1 则全连接检测
+    device='cuda'
+):
+    """
+    Args:
+        connection_head: 训练好的 MLP 模型
+        vtx_feats: [N, C] 顶点特征
+        vtx_coords: [N, 3] 顶点坐标 (用于 KNN 筛选候选边)
+        batch_size: MLP 推理的 batch size
+        threshold: 判定连接的概率阈值
+        k_neighbors: K-NN 数量。如果是 None 或 -1，则检测所有 N*(N-1)/2 对。
+    """
+    num_verts = vtx_feats.shape[0]
+    if num_verts < 3:
+        return [], [] # 无法构成三角形
+
+    connection_head.eval()
+    
+    # --- 1. 生成候选边 (Candidate Edges) ---
+    if k_neighbors is not None and k_neighbors > 0 and k_neighbors < num_verts:
+        # 策略 A: 局部 KNN (推荐)
+        # 计算距离矩阵可能会 OOM，使用分块或 KDTree/Faiss，这里用 PyTorch 的 cdist 分块简化版
+        # 或者直接暴力 cdist 如果 N < 10000
+        
+        # 为了简单且高效，这里演示简单的 cdist (注意显存)
+        # 如果 N 很大 (>5000)，建议使用 faiss 或 scipy.spatial.cKDTree
+        dist_mat = torch.cdist(vtx_coords.float(), vtx_coords.float()) # [N, N]
+        
+        # 取 topk (smallest distance)，排除自己
+        # values: [N, K], indices: [N, K]
+        _, indices = torch.topk(dist_mat, k=k_neighbors + 1, dim=1, largest=False)
+        neighbor_indices = indices[:, 1:] # 去掉第一列（自己）
+        
+        # 构建 source, target 索引
+        src = torch.arange(num_verts, device=device).unsqueeze(1).repeat(1, k_neighbors).flatten()
+        dst = neighbor_indices.flatten()
+        
+        # 此时得到的边是双向的 (u->v 和 v->u 可能都存在)，为了效率可以去重
+        # 但为了利用你的 symmetric MLP，保留双向或者只保留 u < v 均可
+        # 这里为了简单，我们生成 u < v 的 mask
+        mask = src < dst
+        u_indices = src[mask]
+        v_indices = dst[mask]
+        
+    else:
+        # 策略 B: 全连接 (O(N^2)) - 仅当 N 较小时使用
+        u_indices, v_indices = torch.triu_indices(num_verts, num_verts, offset=1, device=device)
+
+    # --- 2. 批量推理 ---
+    all_probs = []
+    num_candidates = u_indices.shape[0]
+    
+    with torch.no_grad():
+        for i in range(0, num_candidates, batch_size):
+            end = min(i + batch_size, num_candidates)
+            batch_u = u_indices[i:end]
+            batch_v = v_indices[i:end]
+            
+            feat_u = vtx_feats[batch_u]
+            feat_v = vtx_feats[batch_v]
+            
+            # Symmetric Forward (和你训练时保持一致)
+            # A -> B
+            input_uv = torch.cat([feat_u, feat_v], dim=-1)
+            logits_uv = connection_head(input_uv)
+            
+            # B -> A
+            input_vu = torch.cat([feat_v, feat_u], dim=-1)
+            logits_vu = connection_head(input_vu)
+            
+            # Sum logits
+            logits = (logits_uv + logits_vu)
+            probs = torch.sigmoid(logits)
+            all_probs.append(probs)
+            
+    all_probs = torch.cat(all_probs).squeeze() # [M]
+    
+    # --- 3. 筛选连接边 ---
+    connected_mask = all_probs > threshold
+    final_u = u_indices[connected_mask].cpu().numpy()
+    final_v = v_indices[connected_mask].cpu().numpy()
+    
+    edges = np.stack([final_u, final_v], axis=1) # [E, 2]
+    
+    return edges
+
+def build_triangles_from_edges(edges, num_verts):
+    """
+    从边列表构建三角形。
+    寻找图中所有的 3-Cliques (三元环)。
+    这在图论中是一个经典问题，可以使用 networkx 库。
+    """
+    if len(edges) == 0:
+        return np.empty((0, 3), dtype=int)
+        
+    G = nx.Graph()
+    G.add_nodes_from(range(num_verts))
+    G.add_edges_from(edges)
+    
+    # 寻找所有的 3-cliques (三角形)
+    # enumerate_all_cliques 返回所有大小的 clique，我们需要过滤大小为 3 的
+    # 或者使用 nx.triangles ? 不，那个只返回数量
+    # 使用 nx.enumerate_all_cliques 效率可能较低，对于稀疏图还可以
+    
+    # 更快的方法：迭代每条边 (u, v)，查找 u 和 v 的公共邻居 w
+    triangles = []
+    adj = [set(G.neighbors(n)) for n in range(num_verts)]
+    
+    # 为了避免重复 (u, v, w), (v, w, u)... 我们可以强制 u < v < w
+    # 既然 edges 已经是 u < v (如果我们之前做了 triu)，则只需要找 w > v 且 w in adj[u]
+    
+    # 优化算法：
+    for u, v in edges:
+        if u > v: u, v = v, u # 确保有序
+        
+        # 找公共邻居
+        common = adj[u].intersection(adj[v])
+        for w in common:
+            if w > v: # 强制顺序 u < v < w 防止重复
+                triangles.append([u, v, w])
+                
+    return np.array(triangles)
+
+def downsample_voxels(
+    voxels: torch.Tensor,
+    input_resolution: int,
+    output_resolution: int
+) -> torch.Tensor:
+    if input_resolution % output_resolution != 0:
+        raise ValueError(f"input_resolution ({input_resolution}) must be divisible "
+                         f"by output_resolution ({output_resolution}).")
+
+    factor = input_resolution // output_resolution
+
+    downsampled_voxels = voxels.clone().to(torch.long)
+
+    downsampled_voxels[:, 1:] //= factor
+
+    unique_downsampled_voxels = torch.unique(downsampled_voxels, dim=0)
+    return unique_downsampled_voxels
+
+def visualize_colored_points_ply(coords, vectors, filename):
+    """
+    可视化点云，并用向量方向的颜色来表示，保存为 PLY 文件。
+    
+    Args:
+        coords (torch.Tensor or np.ndarray): 3D坐标，形状为 (N, 3)。
+        vectors (torch.Tensor or np.ndarray): 方向向量，形状为 (N, 3)。
+        filename (str): 保存输出文件的名称，必须是 .ply 格式。
+    """
+    # 确保输入是 numpy 数组
+    if isinstance(coords, torch.Tensor):
+        coords = coords.detach().cpu().numpy()
+    if isinstance(vectors, torch.Tensor):
+        vectors = vectors.detach().cpu().to(torch.float32).numpy()
+
+    # 检查输入数据是否为空，防止崩溃
+    if coords.size == 0 or vectors.size == 0:
+        print(f"警告：输入数据为空，未生成 {filename} 文件。")
+        return
+
+    # 将向量分量从 [-1, 1] 映射到 [0, 255]
+    # np.clip 用于将数值限制在 -1 和 1 之间，防止颜色溢出
+    # (vectors + 1) 将范围从 [-1, 1] 移动到 [0, 2]
+    # * 127.5 将范围从 [0, 2] 缩放到 [0, 255]
+    colors = np.clip((vectors + 1) * 127.5, 0, 255).astype(np.uint8)
+
+    # 创建一个点云对象，并传入颜色信息
+    # trimesh.PointCloud 能够自动处理带颜色的点
+    points = trimesh.points.PointCloud(coords, colors=colors)
+    # 导出为 PLY 文件
+    points.export(filename, file_type='ply')
+    print(f"可视化文件已成功保存为: {filename}")
+
+
+def compute_vertex_matching(pred_coords, gt_coords, threshold=1.0):
+    # 转换为整数坐标并去重
+    print('len(pred_coords)', len(pred_coords))
+
+    pred_array = np.unique(pred_coords.detach().to(torch.float32).cpu().numpy(), axis=0)
+    gt_array = np.unique(gt_coords.detach().cpu().to(torch.float32).numpy(), axis=0)
+    print('len(pred_array)', len(pred_array))
+    pred_total = len(pred_array)
+    gt_total = len(gt_array)
+
+    # 如果没有点，直接返回
+    if pred_total == 0 or gt_total == 0:
+        return 0, 0.0, pred_total, gt_total
+
+    # 建立 KDTree（以 gt 为基准）
+    tree = KDTree(gt_array)
+
+    # 查找预测点到最近的 gt 点
+    dist, indices = tree.query(pred_array, k=1)
+    dist = dist.squeeze()
+    indices = indices.squeeze()
+
+    # 贪心去重：确保 1 对 1
+    matches = 0
+    used_gt = set()
+    for d, idx in zip(dist, indices):
+        if d <= threshold and idx not in used_gt:
+            matches += 1
+            used_gt.add(idx)
+
+    match_rate = matches / max(gt_total, 1)
+
+    return matches, match_rate, pred_total, gt_total
+
+def flatten_coords_4d(coords_4d: torch.Tensor):
+    coords_4d_long = coords_4d.long()
+    
+    base_x = 512 
+    base_y = 512 * 512
+    base_z = 512 * 512 * 512 
+    
+    flat_coords = coords_4d_long[:, 0] * base_z + \
+                  coords_4d_long[:, 1] * base_y + \
+                  coords_4d_long[:, 2] * base_x + \
+                  coords_4d_long[:, 3]
+    return flat_coords
+
+class Tester:
+    def __init__(self, ckpt_path, config_path=None, dataset_path=None):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.ckpt_path = ckpt_path
+        
+        self.config = self._load_config(config_path)
+        self.dataset_path = dataset_path # or self.config['dataset']['path']
+        checkpoint = torch.load(self.ckpt_path, map_location='cpu')
+        self.epoch = checkpoint.get('epoch', 0)
+        
+        self._init_models()
+        self._init_dataset()
+        
+        self.result_dir = os.path.join(os.path.dirname(ckpt_path), "evaluation_results")
+        os.makedirs(self.result_dir, exist_ok=True)
+
+        dataset_name_clean = os.path.basename(self.dataset_path).replace('.npz', '').replace('.npy', '')
+        self.output_voxel_dir = os.path.join(os.path.dirname(ckpt_path), 
+                                             f"epoch_{self.epoch}_{dataset_name_clean}_voxels_0_gs")
+        os.makedirs(self.output_voxel_dir, exist_ok=True)
+
+        self.output_obj_dir = os.path.join(os.path.dirname(ckpt_path), 
+                                           f"epoch_{self.epoch}_{dataset_name_clean}_obj_0_gs")
+        os.makedirs(self.output_obj_dir, exist_ok=True)
+
+    def _save_logit_visualization(self, dense_vol, name, sample_name, ply_threshold=0.01):
+        """
+        保存 Logit 的 3D .npy 文件、2D 最大投影热力图，以及带颜色和透明度的 3D .ply 点云
+        
+        Args:
+            dense_vol: (H, W, D) numpy array, values in [0, 1]
+            name: str (e.g., "edge" or "vertex")
+            sample_name: str
+            ply_threshold: float, 只有概率大于此值的点才会被保存
+        """
+        # 1. 保存原始 Dense 数据 (可选)
+        npy_path = os.path.join(self.output_voxel_dir, f"{sample_name}_{name}_logits.npy")
+        # np.save(npy_path, dense_vol)
+
+        # 2. 生成 2D 投影热力图 (保持不变)
+        proj_x = np.max(dense_vol, axis=0)
+        proj_y = np.max(dense_vol, axis=1)
+        proj_z = np.max(dense_vol, axis=2)
+
+        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+        im0 = axes[0].imshow(proj_x, cmap='turbo', vmin=0, vmax=1, origin='lower')
+        axes[0].set_title(f"{name} Max-Proj (YZ)")
+        im1 = axes[1].imshow(proj_y, cmap='turbo', vmin=0, vmax=1, origin='lower')
+        axes[1].set_title(f"{name} Max-Proj (XZ)")
+        im2 = axes[2].imshow(proj_z, cmap='turbo', vmin=0, vmax=1, origin='lower')
+        axes[2].set_title(f"{name} Max-Proj (XY)")
+
+        fig.colorbar(im2, ax=axes, orientation='vertical', fraction=0.02, pad=0.04)
+        plt.suptitle(f"{sample_name} - {name} Occupancy Probability")
+        
+        png_path = os.path.join(self.output_voxel_dir, f"{sample_name}_{name}_heatmap.png")
+        plt.savefig(png_path, dpi=150)
+        plt.close(fig)
+
+        # ------------------------------------------------------------------
+        # 3. 保存为带颜色和透明度(RGBA)的 PLY 点云
+        # ------------------------------------------------------------------
+        # 筛选出概率大于阈值的点坐标
+        indices = np.argwhere(dense_vol > ply_threshold)
+        
+        if len(indices) > 0:
+            # 获取这些点的概率值 [0, 1]
+            values = dense_vol[indices[:, 0], indices[:, 1], indices[:, 2]]
+            
+            # 使用 matplotlib 的 colormap 进行颜色映射
+            import matplotlib.cm as cm
+            cmap = cm.get_cmap('turbo') 
+            
+            # map values [0, 1] to RGBA [0, 1] (N, 4)
+            colors_float = cmap(values) 
+            
+            # -------------------------------------------------------
+            # 【核心修改】：修改 Alpha 通道 (透明度)
+            # -------------------------------------------------------
+            # 让透明度直接等于概率值。
+            # 概率 1.0 -> Alpha 1.0 (完全不透明/颜色深)
+            # 概率 0.1 -> Alpha 0.1 (非常透明/颜色浅)
+            colors_float[:, 3] = values
+            
+            # 转换为 uint8 [0, 255]，保留 4 个通道 (R, G, B, A)
+            colors_uint8 = (colors_float * 255).astype(np.uint8)
+            
+            # 坐标转换
+            vertices = indices 
+            
+            ply_filename = f"{sample_name}_{name}_logits_colored.ply"
+            ply_save_path = os.path.join(self.output_voxel_dir, ply_filename)
+            
+            try:
+                # 使用 Trimesh 保存 (Trimesh 支持 (N, 4) 的 colors)
+                pcd = trimesh.points.PointCloud(vertices=vertices, colors=colors_uint8)
+                pcd.export(ply_save_path)
+                print(f"Saved colored RGBA logit PLY to {ply_save_path}")
+            except Exception as e:
+                print(f"Failed to save PLY with trimesh: {e}")
+                # Fallback: 手动写入 PLY (需要添加 alpha 属性)
+                with open(ply_save_path, 'w') as f:
+                    f.write("ply\n")
+                    f.write("format ascii 1.0\n")
+                    f.write(f"element vertex {len(vertices)}\n")
+                    f.write("property float x\n")
+                    f.write("property float y\n")
+                    f.write("property float z\n")
+                    f.write("property uchar red\n")
+                    f.write("property uchar green\n")
+                    f.write("property uchar blue\n")
+                    f.write("property uchar alpha\n") # 新增 Alpha 属性
+                    f.write("end_header\n")
+                    for i in range(len(vertices)):
+                        v = vertices[i]
+                        c = colors_uint8[i] # c is now (R, G, B, A)
+                        f.write(f"{v[0]} {v[1]} {v[2]} {c[0]} {c[1]} {c[2]} {c[3]}\n")
+
+    def _point_line_segment_distance(self, px, py, pz, x1, y1, z1, x2, y2, z2):
+        """
+        计算点 (px,py,pz) 到线段 (x1,y1,z1)-(x2,y2,z2) 的最短距离的平方。
+        全部输入为 Tensor，支持广播。
+        """
+        # 线段向量 AB
+        ABx = x2 - x1
+        ABy = y2 - y1
+        ABz = z2 - z1
+        
+        # 向量 AP
+        APx = px - x1
+        APy = py - y1
+        APz = pz - z1
+        
+        # AB 的长度平方
+        AB_sq = ABx**2 + ABy**2 + ABz**2
+        
+        # 避免除以0 (如果两端点重合)
+        AB_sq = torch.clamp(AB_sq, min=1e-6)
+        
+        # 投影系数 t = (AP · AB) / |AB|^2
+        t = (APx * ABx + APy * ABy + APz * ABz) / AB_sq
+        
+        # 限制 t 在 [0, 1] 之间（线段约束）
+        t = torch.clamp(t, 0.0, 1.0)
+        
+        # 最近点 (Projection)
+        closestX = x1 + t * ABx
+        closestY = y1 + t * ABy
+        closestZ = z1 + t * ABz
+        
+        # 距离平方
+        dx = px - closestX
+        dy = py - closestY
+        dz = pz - closestZ
+        
+        return dx**2 + dy**2 + dz**2
+
+    def _extract_mesh_projection_based(
+        self, 
+        vtx_result: dict, 
+        edge_result: dict, 
+        resolution: int = 1024,
+        vtx_prob_threshold: float = 0.5,
+        
+        # --- 你的新逻辑参数 ---
+        search_radius: float = 128.0,      # 1. 候选边最大长度
+        project_dist_thresh: float = 1.5,  # 2. 投影距离阈值 (管子半径，单位：voxel)
+        dir_align_threshold: float = 0.6,  # 3. 方向相似度阈值 (cos theta)
+        connect_ratio_threshold: float = 0.4, # 4. 最终连接阈值 (匹配点数 / 理论长度)
+        
+        edge_prob_threshold: float = 0.1,  # 仅仅用于提取"存在的"体素
+    ):
+        t_start = time.perf_counter()
+
+        # ---------------------------------------------------------------------
+        # 1. 准备全局数据：提取所有"活着"的 Edge Voxels (作为点云处理)
+        # ---------------------------------------------------------------------
+        e_probs = torch.sigmoid(edge_result['occ_probs'][:, 0])
+        e_coords = edge_result['coords_4d'][:, 1:].float() # (N, 3)
+        
+        # 获取方向向量
+        if 'predicted_direction_feats' in edge_result:
+            e_dirs = edge_result['predicted_direction_feats'] # (N, 3)
+            # 归一化方向
+            e_dirs = F.normalize(e_dirs, p=2, dim=1)
+        else:
+            print("Warning: No direction features, using dummy.")
+            e_dirs = torch.zeros_like(e_coords)
+
+        # 筛选有效的 Edge Voxels (Global Point Cloud)
+        valid_mask = e_probs > edge_prob_threshold
+        
+        cloud_coords = e_coords[valid_mask] # (M, 3)
+        cloud_dirs = e_dirs[valid_mask]     # (M, 3)
+        
+        num_cloud = cloud_coords.shape[0]
+        print(f"[Projection] Global active edge voxels: {num_cloud}")
+        
+        if num_cloud == 0:
+             return [], []
+
+        # ---------------------------------------------------------------------
+        # 2. 准备顶点和候选边
+        # ---------------------------------------------------------------------
+        v_probs = torch.sigmoid(vtx_result['occ_probs'][:, 0])
+        v_coords = vtx_result['coords_4d'][:, 1:].float()
+        v_mask = v_probs > vtx_prob_threshold
+        valid_v_coords = v_coords[v_mask] # (V, 3)
+
+        if valid_v_coords.shape[0] < 2:
+            return valid_v_coords.cpu().numpy() / resolution, []
+
+        # 生成所有可能的候选边 (基于距离粗筛)
+        dists = torch.cdist(valid_v_coords, valid_v_coords)
+        triu_mask = torch.triu(torch.ones_like(dists), diagonal=1).bool()
+        cand_mask = (dists < search_radius) & triu_mask
+        cand_indices = torch.nonzero(cand_mask, as_tuple=False) # (E_cand, 2)
+        
+        p1s = valid_v_coords[cand_indices[:, 0]] # (E, 3)
+        p2s = valid_v_coords[cand_indices[:, 1]] # (E, 3)
+        
+        num_candidates = p1s.shape[0]
+        print(f"[Projection] Checking {num_candidates} candidate pairs...")
+
+        # ---------------------------------------------------------------------
+        # 3. 循环处理候选边 (使用 Bounding Box 快速裁剪)
+        # ---------------------------------------------------------------------
+        final_edges = []
+        
+        # 预计算所有候选边的方向和长度
+        edge_vecs = p2s - p1s
+        edge_lengths = torch.norm(edge_vecs, dim=1)
+        edge_dirs = F.normalize(edge_vecs, p=2, dim=1)
+
+        # 为了避免显存爆炸，也不要在 Python 里做太慢的循环
+        # 我们对点云进行操作太慢，对每一条边去遍历整个点云也太慢。
+        # 策略：
+        # 我们循环“边”，但在循环内部利用 mask 快速筛选点云。
+        # 由于 Python 循环 10000 次会很慢，我们只处理那些有希望的边。
+        # 这里为了演示逻辑的准确性，我们使用简单的循环，但在 GPU 上做计算。
+        
+        # 将全局点云拆分到各个坐标轴，便于快速 BBox 筛选
+        cx, cy, cz = cloud_coords[:, 0], cloud_coords[:, 1], cloud_coords[:, 2]
+
+        # 优化：如果候选边太多，可以分块。这里假设边在 5万以内，点在 10万以内，可以处理。
+        
+        # 这一步是瓶颈，我们尝试用 Python 循环，但只对局部点计算
+        # 为了加速，我们可以将点云放入 HashGrid 或者只是简单的 BBox Check。
+        
+        # 让我们用简单的逻辑：对于每条边，找出 BBox 内的点，算距离。
+        # 这里的 batch_size 是指一次并行处理多少条边
+        
+        batch_size = 128 # 每次处理 128 条边
+        
+        for i in range(0, num_candidates, batch_size):
+            end = min(i + batch_size, num_candidates)
+            
+            # 当前批次的边数据
+            b_p1 = p1s[i:end] # (B, 3)
+            b_p2 = p2s[i:end] # (B, 3)
+            b_dirs = edge_dirs[i:end] # (B, 3)
+            b_lens = edge_lengths[i:end] # (B,)
+            
+            # --- 步骤 A: 投影 & 距离检查 ---
+            # 这是一个 (B, M) 的大矩阵计算，容易 OOM。
+            # M (点云数) 可能很大。
+            # 解决方法：我们反过来思考。
+            # 不计算矩阵，我们只对单个边进行循环？太慢。
+            
+            # 实用优化：只对 bounding box 内的点进行距离计算。
+            # 由于 GPU 难以动态索引不规则数据，我们还是逐个边循环比较稳妥，
+            # 但为了 Python 速度，必须尽可能向量化。
+            
+            # 这里我采用一种折中方案：逐个处理边，但是利用 torch.where 快速定位。
+            # 实际上，对于 Python 里的 for loop，几千次是可以接受的。
+            
+            current_edges_indices = cand_indices[i:end]
+
+            for j in range(len(b_p1)):
+                # 单条边处理
+                p1 = b_p1[j]
+                p2 = b_p2[j]
+                e_dir = b_dirs[j]
+                e_len = b_lens[j].item()
+                
+                # 1. Bounding Box Filter (快速大幅裁剪)
+                # 找出这条边 BBox 范围内的所有点 (+ padding)
+                padding = project_dist_thresh + 2.0
+                min_xyz = torch.min(p1, p2) - padding
+                max_xyz = torch.max(p1, p2) + padding
+                
+                # 利用 boolean mask 筛选
+                mask_x = (cx >= min_xyz[0]) & (cx <= max_xyz[0])
+                mask_y = (cy >= min_xyz[1]) & (cy <= max_xyz[1])
+                mask_z = (cz >= min_xyz[2]) & (cz <= max_xyz[2])
+                bbox_mask = mask_x & mask_y & mask_z
+                
+                subset_coords = cloud_coords[bbox_mask]
+                subset_dirs = cloud_dirs[bbox_mask]
+                
+                if subset_coords.shape[0] == 0:
+                    continue
+
+                # 2. 精确距离计算 (Projection Distance)
+                # 计算 subset 中每个点到线段 p1-p2 的距离平方
+                dist_sq = self._point_line_segment_distance(
+                    subset_coords[:, 0], subset_coords[:, 1], subset_coords[:, 2],
+                    p1[0], p1[1], p1[2],
+                    p2[0], p2[1], p2[2]
+                )
+                
+                # 3. 距离阈值过滤 (Keep voxels inside the tube)
+                dist_mask = dist_sq < (project_dist_thresh ** 2)
+                
+                # 获取在管子内部的体素
+                tube_dirs = subset_dirs[dist_mask]
+                
+                if tube_dirs.shape[0] == 0:
+                    continue
+                
+                # 4. 方向一致性检查 (Direction Check)
+                # 计算点积 (cos theta)
+                # e_dir 是 (3,), tube_dirs 是 (K, 3)
+                dot_prod = torch.matmul(tube_dirs, e_dir)
+                
+                # 这里使用 abs，因为边可能是无向的，或者网络预测可能反向
+                # 如果你的网络严格预测流向，可以去掉 abs
+                dir_sim = torch.abs(dot_prod)
+                
+                # 统计方向符合要求的体素数量
+                valid_voxel_count = (dir_sim > dir_align_threshold).sum().item()
+                
+                # 5. 比值判决 (Ratio Check)
+                # 量化出的 Voxel 数目 ≈ 边的长度 (e_len)
+                # 如果 e_len 很小(比如<1)，我们设为1防止除以0
+                theoretical_count = max(e_len, 1.0)
+                
+                ratio = valid_voxel_count / theoretical_count
+                
+                if ratio > connect_ratio_threshold:
+                    # 找到了！
+                    global_idx = i + j
+                    edge_tuple = cand_indices[global_idx].cpu().numpy().tolist()
+                    final_edges.append(edge_tuple)
+
+        t_end = time.perf_counter()
+        print(f"[Projection] Logic finished. Accepted {len(final_edges)} edges. Time={t_end - t_start:.4f}s")
+        
+        out_vertices = valid_v_coords.cpu().numpy() / resolution
+        return out_vertices, final_edges
+
+    def _save_voxel_ply(self, coords: torch.Tensor, labels: torch.Tensor, filename: str):
+        if coords.numel() == 0:
+            return
+
+        coords_np = coords.cpu().to(torch.float32).numpy()
+        labels_np = labels.cpu().to(torch.float32).numpy()
+
+        colors = np.zeros((coords_np.shape[0], 3), dtype=np.uint8)
+        colors[labels_np == 0] = [255, 0, 0]
+        colors[labels_np == 1] = [0, 0, 255]
+
+        try:
+            import trimesh
+            point_cloud = trimesh.PointCloud(vertices=coords_np, colors=colors)
+            ply_path = os.path.join(self.output_voxel_dir, f"{filename}.ply")
+            point_cloud.export(ply_path)
+        except ImportError:
+            ply_path = os.path.join(self.output_voxel_dir, f"{filename}.ply")
+            with open(ply_path, 'w') as f:
+                f.write("ply\n")
+                f.write("format ascii 1.0\n")
+                f.write(f"element vertex {coords_np.shape[0]}\n")
+                f.write("property float x\n")
+                f.write("property float y\n")
+                f.write("property float z\n")
+                f.write("property uchar red\n")
+                f.write("property uchar green\n")
+                f.write("property uchar blue\n")
+                f.write("end_header\n")
+                for i in range(coords_np.shape[0]):
+                    f.write(f"{coords_np[i,0]} {coords_np[i,1]} {coords_np[i,2]} {colors[i,0]} {colors[i,1]} {colors[i,2]}\n")
+
+    def _load_config(self, config_path=None):
+        if config_path and os.path.exists(config_path):
+            with open(config_path) as f:
+                return yaml.safe_load(f)
+        
+        ckpt_dir = os.path.dirname(self.ckpt_path)
+        possible_configs = [
+            os.path.join(ckpt_dir, "config.yaml"),
+            os.path.join(os.path.dirname(ckpt_dir), "config.yaml")
+        ]
+        
+        for config_file in possible_configs:
+            if os.path.exists(config_file):
+                with open(config_file) as f:
+                    print(f"Loaded config from: {config_file}")
+                    return yaml.safe_load(f)
+        
+        checkpoint = torch.load(self.ckpt_path, map_location='cpu')
+        if 'config' in checkpoint:
+            print("Loaded config from checkpoint")
+            return checkpoint['config']
+        
+        raise FileNotFoundError("Could not find config_edge.yaml in checkpoint directory or parent, and config not saved in checkpoint.")
+    
+    def _init_models(self):
+        self.voxel_encoder = VoxelFeatureEncoder_active_pointnet(
+            in_channels=15, 
+            hidden_dim=256, 
+            out_channels=1024, 
+            scatter_type='mean', 
+            n_blocks=5, 
+            resolution=128,
+
+        ).to(self.device)
+
+        self.connection_head = ConnectionHead(
+            channels=128 * 2,
+            out_channels=1,
+            mlp_ratio=4,
+        ).to(self.device)
+        
+        self.vae = VoxelVAE( # abalation: VoxelVAE_1volume_dilation
+            in_channels=self.config['model']['in_channels'],
+            latent_dim=self.config['model']['latent_dim'],
+            encoder_blocks=self.config['model']['encoder_blocks'],
+            # decoder_blocks=self.config['model']['decoder_blocks'],
+            decoder_blocks_vtx=self.config['model']['decoder_blocks_vtx'],
+            decoder_blocks_edge=self.config['model']['decoder_blocks_edge'],
+            num_heads=8,
+            num_head_channels=64,
+            mlp_ratio=4.0,
+            attn_mode="swin",
+            window_size=8,
+            pe_mode="ape",
+            use_fp16=False,
+            use_checkpoint=False,
+            qk_rms_norm=False,
+            using_subdivide=True,
+            using_attn=self.config['model']['using_attn'],
+            attn_first=self.config['model'].get('attn_first', True),
+            pred_direction=self.config['model'].get('pred_direction', False),
+        ).to(self.device)
+        
+        load_pretrained_woself(
+            checkpoint_path=self.ckpt_path,
+            voxel_encoder=self.voxel_encoder,
+            connection_head=self.connection_head,
+            vae=self.vae,
+        )
+        # --- 【新增】在这里添加权重检查逻辑 ---
+        print(f"--- 正在检查权重文件中的 NaN/Inf 值... ---")
+        has_nan_inf = False
+        if self._check_weights_for_nan_inf(self.vae, "VoxelVAE"):
+            has_nan_inf = True
+        
+        if self._check_weights_for_nan_inf(self.voxel_encoder, "Vertex Encoder"):
+            has_nan_inf = True
+
+        if self._check_weights_for_nan_inf(self.connection_head, "Connection Head"):
+            has_nan_inf = True
+            
+        if not has_nan_inf:
+            print("--- 权重检查通过。未发现 NaN/Inf 值。 ---")
+        else:
+            # 如果发现坏值，直接抛出异常，因为评估无法继续
+            raise ValueError(f"在检查点 '{self.ckpt_path}' 中发现了 NaN 或 Inf 值。请检查导致训练不稳定的权重文件。")
+        # --- 检查逻辑结束 ---
+        
+        self.vae.eval()
+        self.voxel_encoder.eval()
+        self.connection_head.eval()
+    
+    def _init_dataset(self):
+        self.dataset = VoxelVertexDataset_edge(
+            root_dir=self.dataset_path,
+            base_resolution=self.config['dataset']['base_resolution'],
+            min_resolution=self.config['dataset']['min_resolution'],
+            cache_dir='/home/tiger/yy/src/dataset_cache/test_15c_dora',
+            # cache_dir=self.config['dataset']['cache_dir'],
+            renders_dir=self.config['dataset']['renders_dir'],
+            
+            # filter_active_voxels=self.config['dataset']['filter_active_voxels'],
+            filter_active_voxels=False,
+            cache_filter_path=self.config['dataset']['cache_filter_path'],
+
+            sample_type=self.config['dataset']['sample_type'],
+            active_voxel_res=128,
+            pc_sample_number=819200,
+
+        )
+        
+        self.dataloader = DataLoader(
+            self.dataset,
+            batch_size=1,
+            shuffle=False,
+            collate_fn=partial(collate_fn_pointnet),
+            num_workers=0,
+            pin_memory=True,
+            # prefetch_factor=4,
+        )
+    
+    def _check_weights_for_nan_inf(self, model: torch.nn.Module, model_name: str) -> bool:
+        """
+        检查模型的所有参数中是否存在 NaN 或 Inf 值。
+
+        Args:
+            model (torch.nn.Module): 要检查的模型。
+            model_name (str): 模型的名称，用于打印日志。
+
+        Returns:
+            bool: 如果找到 NaN 或 Inf，则返回 True，否则返回 False。
+        """
+        found_issue = False
+        for name, param in model.named_parameters():
+            if torch.isnan(param.data).any():
+                print(f"[!!!] 严重错误: 在模型 '{model_name}' 的参数 '{name}' 中发现 NaN 值！")
+                found_issue = True
+            if torch.isinf(param.data).any():
+                print(f"[!!!] 严重错误: 在模型 '{model_name}' 的参数 '{name}' 中发现 Inf 值！")
+                found_issue = True
+        return found_issue
+
+
+    def _compute_vertex_metrics(self, pred_coords, gt_coords, threshold=1.0):
+        """
+        修改后的函数，确保一对一匹配，并优先匹配最近的点对。
+        """
+        pred_array = np.unique(pred_coords.round().int().cpu().numpy(), axis=0)
+        gt_array = np.unique(gt_coords.round().int().cpu().numpy(), axis=0)
+
+        pred_total = len(pred_array)
+        gt_total = len(gt_array)
+
+        if pred_total == 0 or gt_total == 0:
+            return {
+                'recall': 0.0,
+                'precision': 0.0,
+                'f1': 0.0,
+                'matches': 0,
+                'pred_count': pred_total,
+                'gt_count': gt_total
+            }
+
+        # 依然在预测点上构建KD-Tree，为每个真实点查找最近的预测点
+        tree = cKDTree(pred_array)
+        dists, pred_idxs = tree.query(gt_array, k=1)
+
+        # --- 核心修改部分 ---
+
+        # 1. 创建一个列表，包含 (距离, 真实点索引, 预测点索引)
+        #    这样我们就可以按距离对所有可能的匹配进行排序
+        possible_matches = []
+        for gt_idx, (dist, pred_idx) in enumerate(zip(dists, pred_idxs)):
+            if dist <= threshold:
+                possible_matches.append((dist, gt_idx, pred_idx))
+        
+        # 2. 按距离从小到大排序（贪心策略）
+        possible_matches.sort(key=lambda x: x[0])
+        
+        matches = 0
+        # 使用集合来跟踪已经使用过的预测点和真实点，确保一对一匹配
+        used_pred_indices = set()
+        used_gt_indices = set() # 虽然当前逻辑下gt不会重复，但加上更严谨
+
+        # 3. 遍历排序后的可能匹配，进行一对一分配
+        for dist, gt_idx, pred_idx in possible_matches:
+            # 如果这个预测点和这个真实点都还没有被使用过
+            if pred_idx not in used_pred_indices and gt_idx not in used_gt_indices:
+                matches += 1
+                used_pred_indices.add(pred_idx)
+                used_gt_indices.add(gt_idx)
+
+        # --- 修改结束 ---
+
+        # matches 现在是真正的 True Positives 数量，它绝不会超过 pred_total 或 gt_total
+        recall = matches / gt_total if gt_total > 0 else 0.0
+        precision = matches / pred_total if pred_total > 0 else 0.0
+        
+        # 计算F1时，使用标准的 Precision 和 Recall 定义
+        if (precision + recall) == 0:
+            f1 = 0.0
+        else:
+            f1 = 2 * (precision * recall) / (precision + recall)
+
+        return {
+            'recall': recall,
+            'precision': precision,
+            'f1': f1,
+            'matches': matches,
+            'pred_count': pred_total,
+            'gt_count': gt_total
+        }
+
+    def _compute_vertex_metrics(self, pred_coords, gt_coords, threshold=1.0):
+        """
+        一个折衷的顶点指标计算方案。
+        它沿用“为每个真实点寻找最近预测点”的逻辑，
+        但通过修正计算方式，确保Precision和F1值不会超过1.0。
+        """
+        # 假设 pred_coords 和 gt_coords 是 PyTorch 张量
+        pred_array = np.unique(pred_coords.round().int().cpu().numpy(), axis=0)
+        gt_array = np.unique(gt_coords.round().int().cpu().numpy(), axis=0)
+
+        pred_total = len(pred_array)
+        gt_total = len(gt_array)
+
+        if pred_total == 0 or gt_total == 0:
+            return {
+                'recall': 0.0,
+                'precision': 0.0,
+                'f1': 0.0,
+                'matches': 0,
+                'pred_count': pred_total,
+                'gt_count': gt_total
+            }
+
+        # 在预测点上构建KD-Tree，为每个真实点查找最近的预测点
+        tree = cKDTree(pred_array)
+        dists, _ = tree.query(gt_array, k=1) # 我们在这里其实不需要 pred 的索引
+
+        # 1. 计算从 gt 角度出发的匹配数 (True Positives for Recall)
+        #    这和您的第一个函数完全一样。
+        #    这个值代表了“有多少个真实点被成功找到了”。
+        matches_from_gt = np.sum(dists <= threshold)
+
+        # 2. 计算 Recall (召回率)
+        #    召回率的分母是真实点的总数，所以这里的计算是合理的。
+        recall = matches_from_gt / gt_total if gt_total > 0 else 0.0
+
+        # 3. 计算 Precision (精确率) - ✅ 这是核心修正点
+        #    精确率的分母是预测点的总数。
+        #    分子（True Positives）不能超过预测点的总数。
+        #    因此，我们取 matches_from_gt 和 pred_total 中的较小值。
+        #    这解决了 Precision > 1 的问题。
+        tp_for_precision = min(matches_from_gt, pred_total)
+        precision = tp_for_precision / pred_total if pred_total > 0 else 0.0
+
+        # 4. 使用标准的F1分数公式
+        #    您原来的 F1 公式 `2 * matches / (pred + gt)` 是 L1-Score，
+        #    更常用的是基于 Precision 和 Recall 的调和平均数。
+        if (precision + recall) == 0:
+            f1 = 0.0
+        else:
+            f1 = 2 * (precision * recall) / (precision + recall)
+
+        return {
+            'recall': recall,
+            'precision': precision,
+            'f1': f1,
+            'matches': matches_from_gt, # 仍然报告原始的匹配数，便于观察
+            'pred_count': pred_total,
+            'gt_count': gt_total
+        }
+
+    def _compute_chamfer_distance(self, p1: torch.Tensor, p2: torch.Tensor, one_sided: bool = False):
+        if len(p1) == 0 or len(p2) == 0:
+            return float('nan')
+        
+        dist_p1_p2 = torch.min(torch.cdist(p1, p2), dim=1)[0].mean()
+        
+        if one_sided:
+            return dist_p1_p2.item()
+        else:
+            dist_p2_p1 = torch.min(torch.cdist(p2, p1), dim=1)[0].mean()
+            return (dist_p1_p2 + dist_p2_p1).item() / 2
+
+    def visualize_latent_space_pca(self, sample_idx: int):
+        """
+        Encodes a sample, performs PCA on its latent features, and saves a 
+        colored PLY file for visualization.
+        
+        The position of each point in the PLY file corresponds to the spatial
+        location in the latent grid.
+        
+        The color of each point represents the first three principal components
+        of its feature vector.
+        """
+        print(f"--- Starting Latent Space PCA Visualization for Sample {sample_idx} ---")
+        self.vae.eval()
+        
+        try:
+            # 1. Get the latent representation for the sample
+            latent = self._get_latent_for_sample(sample_idx)
+        except ValueError as e:
+            print(f"Error: {e}")
+            return
+            
+        latent_coords = latent.coords.detach().cpu().numpy()
+        latent_feats = latent.feats.detach().cpu().numpy()
+
+        if latent_feats.shape[0] < 3:
+            print(f"Warning: Not enough latent points ({latent_feats.shape[0]}) to perform PCA. Skipping.")
+            return
+
+        print(f"--> Performing PCA on {latent_feats.shape[0]} latent vectors of dimension {latent_feats.shape[1]}...")
+        
+        # 2. Perform PCA to reduce feature dimensions to 3
+        pca = PCA(n_components=3)
+        pca_features = pca.fit_transform(latent_feats)
+        
+        print(f"    Explained variance ratio by 3 components: {pca.explained_variance_ratio_}")
+        print(f"    Total explained variance: {np.sum(pca.explained_variance_ratio_):.4f}")
+
+        # 3. Normalize the PCA components to be used as RGB colors [0, 255]
+        # We normalize each component independently to maximize color contrast
+        normalized_colors = np.zeros_like(pca_features)
+        for i in range(3):
+            min_val = pca_features[:, i].min()
+            max_val = pca_features[:, i].max()
+            if max_val - min_val > 1e-6:
+                normalized_colors[:, i] = (pca_features[:, i] - min_val) / (max_val - min_val)
+            else:
+                normalized_colors[:, i] = 0.5 # Handle case of constant value
+
+        colors_uint8 = (normalized_colors * 255).astype(np.uint8)
+
+        # 4. Prepare spatial coordinates for the point cloud
+        # latent_coords is (batch_idx, x, y, z), we want the xyz part
+        spatial_coords = latent_coords[:, 1:]
+
+        # 5. Create and save the colored PLY file
+        try:
+            # Create a Trimesh PointCloud object
+            point_cloud = trimesh.points.PointCloud(vertices=spatial_coords, colors=colors_uint8)
+            
+            # Define the output filename
+            filename = f"sample_{sample_idx}_latent_pca.ply"
+            ply_path = os.path.join(self.output_voxel_dir, filename)
+            
+            # Export the file
+            point_cloud.export(ply_path)
+            print(f"--> Successfully saved PCA visualization to: {ply_path}")
+
+        except Exception as e:
+            print(f"Error during Trimesh export: {e}")
+            print("Please ensure 'trimesh' is installed correctly.")
+
+    def _get_latent_for_sample(self, sample_idx: int) -> SparseTensor:
+        """
+        Encodes a single sample and returns its latent representation.
+        """
+        print(f"--> Encoding sample {sample_idx} to get its latent vector...")
+        # Get data for the specified sample
+        batch_data = self.dataset[sample_idx]
+        if batch_data is None:
+            raise ValueError(f"Sample at index {sample_idx} could not be loaded.")
+        
+        # Use the collate function to form a batch
+        batch_data = collate_fn_pointnet([batch_data])
+
+        with torch.no_grad():
+            active_coords = batch_data['active_voxels_128'].to(self.device)
+            point_cloud = batch_data['point_cloud_128'].to(self.device)
+
+            active_voxel_feats = self.voxel_encoder(
+                p=point_cloud,
+                sparse_coords=active_coords,
+                res=128,
+                bbox_size=(-0.5, 0.5),
+            )
+
+            sparse_input = SparseTensor(
+                feats=active_voxel_feats,
+                coords=active_coords.int()
+            )
+            
+            # 2. Encode to get the latent representation
+            latent_128, posterior = self.vae.encode(sparse_input, sample_posterior=True,)
+            print(f"    Latent for sample {sample_idx} obtained. Shape: {latent_128.feats.shape}")
+            return latent_128
+    
+    
+
+    def evaluate(self, num_samples=None, visualize=False, chamfer_threshold=0.9, threshold=1.):
+        total_samples = len(self.dataset)
+        eval_samples = min(num_samples or total_samples, total_samples)
+        sample_indices = random.sample(range(total_samples), eval_samples) if num_samples else range(total_samples)
+        # sample_indices = range(eval_samples)
+
+        eval_dataset = Subset(self.dataset, sample_indices)
+        eval_loader = DataLoader(
+            eval_dataset,
+            batch_size=1,
+            shuffle=False,
+            collate_fn=partial(collate_fn_pointnet),
+            num_workers=self.config['training']['num_workers'],
+            pin_memory=True,
+        )
+        
+        per_sample_metrics = {
+            'vertex': {res: [] for res in [128, 256, 512]},
+            'edge': {res: [] for res in [128, 256, 512]}, 
+            'sample_names': []
+        }
+        avg_metrics = {
+            'vertex': {res: defaultdict(list) for res in [128, 256, 512]},
+            'edge': {res: defaultdict(list) for res in [128, 256, 512]},
+        }
+
+        self.vae.eval()
+
+        for batch_idx, batch_data in enumerate(tqdm(eval_loader, desc="Evaluating")):
+            if batch_data is None:
+                continue
+            sample_idx = sample_indices[batch_idx]
+            sample_name = f'sample_{sample_idx}'
+            per_sample_metrics['sample_names'].append(sample_name)
+
+            batch_save_path = f"/home/tiger/yy/src/checkpoints/output_slat_flow_matching_active/40w_128to512_head_rope/143500_sample_active_vis_42seed_trellis/gt_data_batch_{batch_idx}.pt"
+            if not os.path.exists(batch_save_path):
+                print(f"Warning: Saved batch file not found: {batch_save_path}")
+                continue
+            batch_data = torch.load(batch_save_path, map_location=self.device)
+            
+            with torch.no_grad():
+                # 1. Get input data
+                combined_voxels_512 = batch_data['combined_voxels_512'].to(self.device)
+                combined_voxel_labels_512 = batch_data['combined_voxel_labels_512'].to(self.device)
+                gt_combined_endpoints_512 = batch_data['gt_combined_endpoints_512'].to(self.device)
+                gt_combined_errors_512 = batch_data['gt_combined_errors_512'].to(self.device)
+
+                edge_mask = (combined_voxel_labels_512 == 1)
+
+                gt_edge_endpoints_512 = gt_combined_endpoints_512[edge_mask].to(self.device)
+                
+                gt_edge_voxels_512 = combined_voxels_512[edge_mask].to(self.device)
+
+                p1 = gt_edge_endpoints_512[:, 1:4].float()
+                p2 = gt_edge_endpoints_512[:, 4:7].float()
+
+                mask = ( (p1[:,0] < p2[:,0]) |
+                        ((p1[:,0] == p2[:,0]) & (p1[:,1] < p2[:,1])) |
+                        ((p1[:,0] == p2[:,0]) & (p1[:,1] == p2[:,1]) & (p1[:,2] <= p2[:,2])) )
+
+                pA = torch.where(mask[:, None], p1, p2)  # smaller one
+                pB = torch.where(mask[:, None], p2, p1)  # larger one
+
+                d = pB - pA
+                dir_gt = F.normalize(d, dim=-1, eps=1e-6)
+
+                gt_vertex_voxels_512 = batch_data['gt_vertex_voxels_512'].to(self.device).int()
+
+                vtx_128 = downsample_voxels(gt_vertex_voxels_512, input_resolution=512, output_resolution=128)
+                vtx_256 = downsample_voxels(gt_vertex_voxels_512, input_resolution=512, output_resolution=256)
+                
+                edge_128 = downsample_voxels(combined_voxels_512, input_resolution=512, output_resolution=128)
+                edge_256 = downsample_voxels(combined_voxels_512, input_resolution=512, output_resolution=256)
+                edge_512 = combined_voxels_512
+
+
+                gt_edge_voxels_list = [
+                    edge_128,
+                    edge_256,
+                    edge_512,
+                ]
+
+                active_coords = batch_data['active_voxels_128'].to(self.device)
+                point_cloud = batch_data['point_cloud_128'].to(self.device)
+
+                
+                active_voxel_feats = self.voxel_encoder(
+                    p=point_cloud,
+                    sparse_coords=active_coords,
+                    res=128,
+                    bbox_size=(-0.5, 0.5),
+                )
+
+                sparse_input = SparseTensor(
+                    feats=active_voxel_feats,
+                    coords=active_coords.int()
+                )
+                
+                latent_128, posterior = self.vae.encode(sparse_input)
+
+                
+                load_path = f'/home/tiger/yy/src/checkpoints/output_slat_flow_matching_active/40w_128to512_head_rope/143500_sample_active_vis_42seed_trellis/sample_latent_{batch_idx}.pt'
+                latent_128 = torch.load(load_path, map_location=self.device)
+
+                print('latent_128.feats.mean()', latent_128.feats.mean(), 'latent_128.feats.std()', latent_128.feats.std())
+                print('posterior.mean', posterior.mean.mean(), 'posterior.std', posterior.std.mean(), 'posterior.var', posterior.var.mean())
+                print('latent_128.coords.shape', latent_128.coords.shape)
+                
+    
+                latent_128 = SparseTensor(
+                    coords=latent_128.coords,
+                    feats=latent_128.feats + 0. * torch.randn_like(latent_128.feats),
+                )
+
+                self.output_voxel_dir = os.path.dirname(load_path)
+                self.output_obj_dir = os.path.dirname(load_path)
+                
+                # 7. Decoding with separate vertex and edge processing
+                decoded_results = self.vae.decode(
+                    latent_128, 
+                    gt_vertex_voxels_list=[],
+                    gt_edge_voxels_list=[],
+                    training=False,
+
+                    inference_threshold=0.5,
+                    vis_last_layer=False,
+                )
+            
+                error = 0 # decoded_results[-1]['edge']['predicted_offset_feats']
+                
+                if self.config['model'].get('pred_direction', False):
+                    pred_dir = decoded_results[-1]['edge']['predicted_direction_feats']
+                    zero_mask = (pred_dir == 0).all(dim=1)  # [N]，True 表示这一行全为0
+                    num_zeros = zero_mask.sum().item()
+                    print("Number of zero vectors:", num_zeros)
+
+                    pred_edge_coords_3d = decoded_results[-1]['edge']['coords']
+                    print('pred_edge_coords_3d.shape', pred_edge_coords_3d.shape)
+                    print('pred_dir.shape', pred_dir.shape)
+                    if pred_edge_coords_3d.shape[-1] == 4:
+                        pred_edge_coords_3d = pred_edge_coords_3d[:, 1:]
+                        
+                    save_pth = os.path.join(self.output_voxel_dir, f"{sample_name}_direction.ply")
+                    visualize_colored_points_ply(pred_edge_coords_3d, pred_dir, save_pth)
+
+                    save_pth = os.path.join(self.output_voxel_dir, f"{sample_name}_direction_gt.ply")
+                    visualize_colored_points_ply((gt_edge_voxels_512[:, 1:]), dir_gt, save_pth)
+
+
+                pred_vtx_coords_3d = decoded_results[-1]['vertex']['coords']
+                pred_edge_coords_3d = decoded_results[-1]['edge']['coords']
+
+
+                gt_vertex_voxels_512 = batch_data['gt_vertex_voxels_512'][:, 1:].to(self.device)
+                gt_edge_voxels_512 = batch_data['gt_edge_voxels_512'][:, 1:].to(self.device)
+                
+
+                # Calculate metrics and save results
+                matches, match_rate, pred_total, gt_total = compute_vertex_matching(pred_vtx_coords_3d, gt_vertex_voxels_512, threshold=threshold,)
+                print(f"\n----- Resolution {512} vtx -----")
+                print(f"Pred Vertices: {pred_total} | GT Vertices: {gt_total}")
+                print(f"Matched Vertices: {matches} | Match Rate: {match_rate:.2%}")
+
+                self._save_voxel_ply(pred_vtx_coords_3d / 512., torch.zeros(len(pred_vtx_coords_3d)), f"{sample_name}_pred_vtx")
+                self._save_voxel_ply((pred_edge_coords_3d) / 512, torch.zeros(len(pred_edge_coords_3d)), f"{sample_name}_pred_edge")
+                
+                self._save_voxel_ply(gt_vertex_voxels_512 / 512, torch.zeros(len(gt_vertex_voxels_512)), f"{sample_name}_gt_vertex")
+                self._save_voxel_ply((combined_voxels_512[:, 1:]) / 512., torch.zeros(len(gt_combined_errors_512)), f"{sample_name}_gt_edge")
+
+
+                # Calculate vertex-specific metrics
+                matches, match_rate, pred_total, gt_total = compute_vertex_matching(pred_edge_coords_3d, combined_voxels_512[:, 1:], threshold=threshold,)
+                print(f"\n----- Resolution {512} edge -----")
+                print('pred_edge_coords_3d.shape', pred_edge_coords_3d.shape)
+                print('gt_edge_voxels_512.shape', gt_edge_voxels_512.shape)
+
+                print(f"Pred Vertices: {pred_total} | GT Vertices: {gt_total}")
+                print(f"Matched Vertices: {matches} | Match Rate: {match_rate:.2%}")
+
+                pred_vertex_coords_np = np.round(pred_vtx_coords_3d.cpu().numpy()).astype(int)
+                pred_edges = []
+                gt_vertex_coords_np = np.round(gt_vertex_voxels_512.cpu().numpy()).astype(int)
+                if visualize:
+                    if pred_vtx_coords_3d.shape[-1] == 4:
+                        pred_vtx_coords_float = pred_vtx_coords_3d[:, 1:].float()
+                    else:
+                        pred_vtx_coords_float = pred_vtx_coords_3d.float()
+                        
+                    pred_vtx_feats = decoded_results[-1]['vertex']['feats']
+
+                # ==========================================
+                # Link Prediction & Mesh Generation
+                # ==========================================
+                print("Predicting connectivity...")
+                
+                # 1. 预测边
+                # 注意：K_neighbors 的设置。如果是物体，64 足够了。
+                # 如果点非常稀疏，可能需要更大。
+                pred_edges = predict_mesh_connectivity(
+                    connection_head=self.connection_head, # 或者是 self.connection_head，取决于你在哪里定义的
+                    vtx_feats=pred_vtx_feats,
+                    vtx_coords=pred_vtx_coords_float,
+                    batch_size=4096,
+                    threshold=0.5,
+                    k_neighbors=None, 
+                    device=self.device
+                )
+                print(f"Predicted {len(pred_edges)} edges.")
+
+                # 2. 构建三角形
+                num_verts = pred_vtx_coords_float.shape[0]
+                pred_faces = build_triangles_from_edges(pred_edges, num_verts)
+                print(f"Constructed {len(pred_faces)} triangles.")
+
+                # 3. 保存 OBJ
+                import trimesh
+                
+                # 坐标归一化/还原 (根据你的需求，这里假设你是 0-512 的体素坐标)
+                # 如果想保存为归一化坐标：
+                mesh_verts = pred_vtx_coords_float.cpu().numpy() / 512.0
+                
+                # 如果有 error offset，记得加上！
+                # 你之前的代码好像没有对 vertex 加 offset，只对 edge 加了
+                # 如果 vertex 也有 offset (如 dual contouring)，在这里加上
+                
+                # 移动到中心 (可选)
+                mesh_verts = mesh_verts - 0.5 
+                
+                mesh = trimesh.Trimesh(vertices=mesh_verts, faces=pred_faces)
+                
+                trimesh.repair.fix_normals(mesh)
+                
+                output_obj_path = os.path.join(self.output_voxel_dir, f"{sample_name}_recon.obj")
+                mesh.export(output_obj_path)
+                print(f"Saved mesh to {output_obj_path}")
+                
+                # 保存边线 (用于 Debug)
+                # 有时候三角形很难形成，只看边也很有用
+                edges_path = os.path.join(self.output_voxel_dir, f"{sample_name}_edges.ply")
+                # self._visualize_vertices(pred_edge_coords_np, gt_edge_coords_np, f"{sample_name}_edge_comparison")
+                
+                
+                # Process results at different resolutions
+                for i, res in enumerate([128, 256, 512]):
+                    if i >= len(decoded_results):
+                        continue
+
+                    gt_key = f'gt_vertex_voxels_{res}'
+                    if gt_key not in batch_data:
+                        continue
+                    if i == 0:
+                        pred_coords_res = decoded_results[i]['vtx_sp'].coords[:, 1:].float()
+                        gt_coords_res = batch_data[gt_key][:, 1:].float().to(self.device)
+                    else:
+                        pred_coords_res = decoded_results[i]['vertex']['coords'].float()
+                        gt_coords_res = batch_data[gt_key][:, 1:].float().to(self.device)
+                        
+                    
+                    v_metrics = self._compute_vertex_metrics(pred_coords_res, gt_coords_res, threshold=threshold)
+                    
+                    per_sample_metrics['vertex'][res].append({
+                        'recall': v_metrics['recall'],
+                        'precision': v_metrics['precision'],
+                        'f1': v_metrics['f1'],
+                        'num_pred': len(pred_coords_res),
+                        'num_gt': len(gt_coords_res)
+                    })
+                    
+                    avg_metrics['vertex'][res]['recall'].append(v_metrics['recall'])
+                    avg_metrics['vertex'][res]['precision'].append(v_metrics['precision'])
+                    avg_metrics['vertex'][res]['f1'].append(v_metrics['f1'])
+
+                    gt_edge_key = f'gt_edge_voxels_{res}'
+                    if gt_edge_key not in batch_data:
+                        continue
+                        
+                    if i == 0:
+                        pred_edge_coords_res = decoded_results[i]['edge_sp'].coords[:, 1:].float()
+                        # gt_edge_coords_res = batch_data[gt_edge_key][:, 1:].float().to(self.device)
+                        idx = i
+                        gt_edge_coords_res = gt_edge_voxels_list[idx][:, 1:].float().to(self.device)
+                    elif i == 1:
+                        idx = i
+                        #################################
+                        # pred_edge_coords_res = decoded_results[i]['edge']['coords'].float() - error / 2. + 0.5
+                        # # gt_edge_coords_res = batch_data[gt_edge_key][:, 1:].float().to(self.device)
+                        # gt_edge_coords_res = gt_edge_voxels_list[idx][:, 1:].float().to(self.device) - gt_combined_errors_512[:, 1:].to(self.device) + 0.5
+
+
+                        pred_edge_coords_res = decoded_results[i]['edge']['coords'].float()
+                        gt_edge_coords_res = gt_edge_voxels_list[idx][:, 1:].float().to(self.device)
+
+                        # self._save_voxel_ply(gt_edge_voxels_list[idx][:, 1:].float().to(self.device) / (128*2**i), torch.zeros(len(gt_edge_coords_res)), f"{sample_name}_gt_edge_{128*2**i}res_wooffset")
+                        # self._save_voxel_ply(decoded_results[i]['edge']['coords'].float() / (128*2**i), torch.zeros(len(pred_edge_coords_res)), f"{sample_name}_pred_edge_{128*2**i}res_wooffset")
+
+                    else:
+                        idx = i
+                        pred_edge_coords_res = decoded_results[i]['edge']['coords'].float()
+                        # gt_edge_coords_res = batch_data[gt_edge_key][:, 1:].float().to(self.device)
+                        gt_edge_coords_res = gt_edge_voxels_list[idx][:, 1:].float().to(self.device)
+
+                    # self._save_voxel_ply(gt_edge_coords_res / (128*2**i), torch.zeros(len(gt_edge_coords_res)), f"{sample_name}_gt_edge_{128*2**i}res")
+                    # self._save_voxel_ply(pred_edge_coords_res / (128*2**i), torch.zeros(len(pred_edge_coords_res)), f"{sample_name}_pred_edge_{128*2**i}res")
+
+                    e_metrics = self._compute_vertex_metrics(pred_edge_coords_res, gt_edge_coords_res, threshold=threshold)
+                    
+                    per_sample_metrics['edge'][res].append({
+                        'recall': e_metrics['recall'],
+                        'precision': e_metrics['precision'],
+                        'f1': e_metrics['f1'],
+                        'num_pred': len(pred_edge_coords_res),
+                        'num_gt': len(gt_edge_coords_res)
+                    })
+                    
+                    avg_metrics['edge'][res]['recall'].append(e_metrics['recall'])
+                    avg_metrics['edge'][res]['precision'].append(e_metrics['precision'])
+                    avg_metrics['edge'][res]['f1'].append(e_metrics['f1'])
+                    
+        avg_metrics_processed = {}
+        for category, res_dict in avg_metrics.items():
+            avg_metrics_processed[category] = {}
+            for resolution, metric_dict in res_dict.items():
+                avg_metrics_processed[category][resolution] = {
+                    metric_name: np.mean(values) if values else float('nan')
+                    for metric_name, values in metric_dict.items()
+                }
+
+        result_data = {
+            'config': self.config,
+            'checkpoint': self.ckpt_path,
+            'dataset': self.dataset_path,
+            'num_samples': eval_samples,
+            'per_sample_metrics': per_sample_metrics,
+            'avg_metrics': avg_metrics_processed
+        }
+
+        results_file_path = os.path.join(self.result_dir, f"evaluation_results_epoch{self.epoch}.yaml")
+        with open(results_file_path, 'w') as f:
+            yaml.dump(result_data, f, default_flow_style=False)
+
+        return result_data
+    
+    def _generate_line_voxels(
+        self,
+        p1: torch.Tensor, 
+        p2: torch.Tensor
+    ) -> Tuple[
+        List[Tuple[int, int, int]], 
+        List[Tuple[torch.Tensor, torch.Tensor]], 
+        List[np.ndarray]
+    ]:
+        """
+        Improved version using better sampling strategy
+        """
+        p1_np = p1 #.cpu().numpy()
+        p2_np = p2 #.cpu().numpy()
+        voxel_dict = OrderedDict()
+        
+        # Use proper 3D line voxelization algorithm
+        def bresenham_3d(p1, p2):
+            """3D Bresenham's line algorithm"""
+            x1, y1, z1 = np.round(p1).astype(int)
+            x2, y2, z2 = np.round(p2).astype(int)
+            
+            points = []
+            dx = abs(x2 - x1)
+            dy = abs(y2 - y1)
+            dz = abs(z2 - z1)
+            
+            xs = 1 if x2 > x1 else -1
+            ys = 1 if y2 > y1 else -1
+            zs = 1 if z2 > z1 else -1
+            
+            # Driving axis is X
+            if dx >= dy and dx >= dz:
+                err_1 = 2 * dy - dx
+                err_2 = 2 * dz - dx
+                for i in range(dx + 1):
+                    points.append((x1, y1, z1))
+                    if err_1 > 0:
+                        y1 += ys
+                        err_1 -= 2 * dx
+                    if err_2 > 0:
+                        z1 += zs
+                        err_2 -= 2 * dx
+                    err_1 += 2 * dy
+                    err_2 += 2 * dz
+                    x1 += xs
+            
+            # Driving axis is Y
+            elif dy >= dx and dy >= dz:
+                err_1 = 2 * dx - dy
+                err_2 = 2 * dz - dy
+                for i in range(dy + 1):
+                    points.append((x1, y1, z1))
+                    if err_1 > 0:
+                        x1 += xs
+                        err_1 -= 2 * dy
+                    if err_2 > 0:
+                        z1 += zs
+                        err_2 -= 2 * dy
+                    err_1 += 2 * dx
+                    err_2 += 2 * dz
+                    y1 += ys
+            
+            # Driving axis is Z
+            else:
+                err_1 = 2 * dx - dz
+                err_2 = 2 * dy - dz
+                for i in range(dz + 1):
+                    points.append((x1, y1, z1))
+                    if err_1 > 0:
+                        x1 += xs
+                        err_1 -= 2 * dz
+                    if err_2 > 0:
+                        y1 += ys
+                        err_2 -= 2 * dz
+                    err_1 += 2 * dx
+                    err_2 += 2 * dy
+                    z1 += zs
+            
+            return points
+        
+        # Get all voxels using Bresenham algorithm
+        voxel_coords = bresenham_3d(p1_np, p2_np)
+        
+        # Add all voxels to dictionary
+        for coord in voxel_coords:
+            voxel_dict[tuple(coord)] = (p1, p2)
+        
+        voxel_coords = list(voxel_dict.keys())
+        endpoint_pairs = list(voxel_dict.values())
+
+        # --- compute error vectors ---
+        error_vectors = []
+        diff = p2_np - p1_np
+        d_norm_sq = np.dot(diff, diff)
+        
+        for v in voxel_coords:
+            v_center = np.array(v, dtype=float) + 0.5
+            if d_norm_sq == 0:  # degenerate line
+                closest = p1_np
+            else:
+                t = np.dot(v_center - p1_np, diff) / d_norm_sq
+                t = np.clip(t, 0.0, 1.0)
+                closest = p1_np + t * diff
+            error_vectors.append(v_center - closest)
+        
+        return voxel_coords, endpoint_pairs, error_vectors
+
+
+# 使用示例
+def set_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+def evaluate_checkpoint(ckpt_path, dataset_path, eval_dir):
+    set_seed(42)
+    tester = Tester(ckpt_path=ckpt_path, dataset_path=dataset_path)
+    result_data = tester.evaluate(num_samples=NUM_SAMPLES, visualize=VISUALIZE, chamfer_threshold=CHAMFER_EDGE_THRESHOLD, threshold=THRESHOLD)
+    
+    # 生成文件名
+    epoch_str = os.path.basename(ckpt_path).split('_')[1].split('.')[0]
+    dataset_name = os.path.basename(os.path.normpath(dataset_path))
+    
+    # 保存简版报告(TXT)
+    summary_path = os.path.join(eval_dir, f"epoch{epoch_str}_{dataset_name}_summary_threshold{THRESHOLD}_one2one.txt")
+    with open(summary_path, 'w') as f:
+        # 头部信息
+        f.write(f"Checkpoint: {os.path.basename(ckpt_path)}\n")
+        f.write(f"Dataset: {dataset_name}\n")
+        f.write(f"Evaluation Samples: {result_data['num_samples']}\n\n")
+        
+        # 平均指标
+        f.write("=== Average Metrics ===\n")
+        for category, data in result_data['avg_metrics'].items():
+            if isinstance(data, dict):  # 处理多分辨率情况
+                f.write(f"\n{category.upper()}:\n")
+                for res, metrics in data.items():
+                    f.write(f"  Resolution {res}:\n")
+                    for k, v in metrics.items():
+                        # 确保值是数字类型后再格式化
+                        if isinstance(v, (int, float)):
+                            f.write(f"    {str(k).ljust(15)}: {v:.4f}\n")
+                        else:
+                            f.write(f"    {str(k).ljust(15)}: {str(v)}\n")
+            else:  # 处理非多分辨率情况
+                f.write(f"\n{category.upper()}:\n")
+                for k, v in data.items():
+                    if isinstance(v, (int, float)):
+                        f.write(f"  {str(k).ljust(15)}: {v:.4f}\n")
+                    else:
+                        f.write(f"  {str(k).ljust(15)}: {str(v)}\n")
+        
+        # 样本级详细统计
+        f.write("\n\n=== Detailed Per-Sample Metrics ===\n")
+        for name, vertex_metrics, edge_metrics in zip(
+            result_data['per_sample_metrics']['sample_names'],
+            zip(*[result_data['per_sample_metrics']['vertex'][res] for res in [128, 256, 512]]),
+            zip(*[result_data['per_sample_metrics']['edge'][res] for res in [128, 256, 512]])
+        ):
+            # 样本标题
+            f.write(f"\n◆ Sample: {name}\n")
+            
+            # 顶点指标
+            f.write(f"[Vertex Prediction]\n")
+            f.write(f"  {'Resolution'.ljust(10)} {'Recall'.ljust(8)} {'Precision'.ljust(8)} {'F1'.ljust(8)} {'Pred/Gt'.ljust(10)}\n")
+            for res, metrics in zip([128, 256, 512], vertex_metrics):
+                f.write(f"  {str(res).ljust(10)} "
+                      f"{metrics['recall']:.4f}    "
+                      f"{metrics['precision']:.4f}    "
+                      f"{metrics['f1']:.4f}    "
+                      f"{metrics['num_pred']}/{metrics['num_gt']}\n")
+            
+            # Edge指标
+            f.write(f"[Edge Prediction]\n")
+            f.write(f"  {'Resolution'.ljust(10)} {'Recall'.ljust(8)} {'Precision'.ljust(8)} {'F1'.ljust(8)} {'Pred/Gt'.ljust(10)}\n")
+            for res, metrics in zip([128, 256, 512], edge_metrics):
+                f.write(f"  {str(res).ljust(10)} "
+                      f"{metrics['recall']:.4f}    "
+                      f"{metrics['precision']:.4f}    "
+                      f"{metrics['f1']:.4f}    "
+                      f"{metrics['num_pred']}/{metrics['num_gt']}\n")
+            
+            f.write("-"*60 + "\n")
+
+    print(f"Saved summary to: {summary_path}")
+    return result_data
+
+
+if __name__ == '__main__':
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+        evaluate_all_checkpoints = False  # 设置为 True 启用范围过滤
+        EPOCH_START = 1
+        EPOCH_END = 12
+        CHAMFER_EDGE_THRESHOLD=0.5
+        NUM_SAMPLES=50
+        VISUALIZE=True
+        THRESHOLD=1.5
+        VISUAL_FIELD=False
+        
+        ckpt_path = '/home/tiger/yy/src/checkpoints/vae/unique_files_glb_under6000face_2degree_30ratio_0.01/shapenet_bs2_128to512_wolabel_dir_sorted_dora_small_lowlr/checkpoint_epoch1_batch12000_loss0.1140.pt'
+        dataset_path = '/home/tiger/yy/src/trellis_clean_mesh/mesh_data'
+        
+        if dataset_path == '/HOME/paratera_xy/pxy1054/HDD_POOL/Trisf/data/mesh/objaverse_200_2000':
+            RENDERS_DIR = '/HOME/paratera_xy/pxy1054/HDD_POOL/Trisf/data/mesh_render_img/objaverse_200_2000/renders_cond'
+        else:
+            RENDERS_DIR = ''
+
+    
+        ckpt_dir = os.path.dirname(ckpt_path)
+        eval_dir = os.path.join(ckpt_dir, "evaluate")
+        os.makedirs(eval_dir, exist_ok=True)
+
+        if False:
+            for i in range(NUM_SAMPLES):
+                print("--- Starting Latent Space PCA Visualization ---")
+                tester = Tester(ckpt_path=ckpt_path, dataset_path=dataset_path)
+                tester.visualize_latent_space_pca(sample_idx=i)
+                print("--- PCA Visualization Finished ---")
+
+        if not evaluate_all_checkpoints:
+            evaluate_checkpoint(ckpt_path, dataset_path, eval_dir)
+        else:
+            pt_files = sorted([f for f in os.listdir(ckpt_dir) if f.endswith('.pt')])
+
+            filtered_pt_files = []
+            for f in pt_files:
+                try:
+                    parts = f.split('_')
+                    epoch_str = parts[1].replace('epoch', '')
+                    epoch = int(epoch_str)
+                    if EPOCH_START <= epoch <= EPOCH_END:
+                        filtered_pt_files.append(f)
+                except Exception as e:
+                    print(f"Warning: Could not parse epoch from {f}: {e}")
+                    continue
+
+            for pt_file in filtered_pt_files:
+                full_ckpt_path = os.path.join(ckpt_dir, pt_file)
+                evaluate_checkpoint(full_ckpt_path, dataset_path, eval_dir)

test_slat_vae_128to512_pointnet_vae_head.py

@@ -744,7 +744,7 @@ class Tester:
             root_dir=self.dataset_path,
             base_resolution=self.config['dataset']['base_resolution'],
             min_resolution=self.config['dataset']['min_resolution'],
-            cache_dir='/gemini/user/private/zhaotianhao/dataset_cache/test_15c_dora',
+            cache_dir='/home/tiger/yy/src/dataset_cache',
             # cache_dir=self.config['dataset']['cache_dir'],
             renders_dir=self.config['dataset']['renders_dir'],
             
@@ -1262,15 +1262,12 @@ class Tester:
                 # 如果想保存为归一化坐标：
                 mesh_verts = pred_vtx_coords_float.cpu().numpy() / 512.0
                 
-                # 如果有 error offset，记得加上！
-                # 你之前的代码好像没有对 vertex 加 offset，只对 edge 加了
-                # 如果 vertex 也有 offset (如 dual contouring)，在这里加上
-                
                 # 移动到中心 (可选)
                 mesh_verts = mesh_verts - 0.5 
                 
                 mesh = trimesh.Trimesh(vertices=mesh_verts, faces=pred_faces)
                 
+                trimesh.repair.fix_normals(mesh)
                 # 过滤孤立点 (可选)
                 # mesh.remove_unreferenced_vertices()
                 
@@ -1581,22 +1578,23 @@ def evaluate_checkpoint(ckpt_path, dataset_path, eval_dir):
 if __name__ == '__main__':
     with torch.cuda.amp.autocast(dtype=torch.bfloat16):
         evaluate_all_checkpoints = True  # 设置为 True 启用范围过滤
-        EPOCH_START = 0
+        EPOCH_START = 1
         EPOCH_END = 12
         CHAMFER_EDGE_THRESHOLD=0.5
-        NUM_SAMPLES=20
+        NUM_SAMPLES=50
         VISUALIZE=True
         THRESHOLD=1.5
         VISUAL_FIELD=False
         
-        ckpt_path = '/gemini/user/private/zhaotianhao/checkpoints/vae/train_9w_200_2000face/shapenet_bs2_128to512_wolabel_dir_sorted_dora_small/checkpoint_epoch0_batch10433_loss1.2657.pt'
-        ckpt_path = '/gemini/user/private/zhaotianhao/checkpoints/vae/unique_files_glb_under6000face_2degree_30ratio_0.01/shapenet_bs2_128to512_wolabel_dir_sorted_dora_small/checkpoint_epoch0_batch2000_loss0.3315.pt'
-        
-        dataset_path = '/gemini/user/private/zhaotianhao/data/MERGED_DATASET_count_200_2000_100000/test'
-        dataset_path = '/gemini/user/private/zhaotianhao/data/why_filter_unquantized'
-        # dataset_path = '/gemini/user/private/zhaotianhao/data/trellis500k_compress_glb'
-        dataset_path = '/gemini/user/private/zhaotianhao/data/unique_files_glb_under6000face_2degree_30ratio_0.01'
+        # ckpt_path = '/gemini/user/private/zhaotianhao/checkpoints/vae/train_9w_200_2000face/shapenet_bs2_128to512_wolabel_dir_sorted_dora_small/checkpoint_epoch0_batch10433_loss1.2657.pt'
+        ckpt_path = '/home/tiger/yy/src/checkpoints/vae/unique_files_glb_under6000face_2degree_30ratio_0.01/shapenet_bs2_128to512_wolabel_dir_sorted_dora_small_lowlr/checkpoint_epoch0_batch6000_loss0.1150.pt'
         
+        dataset_path = '/home/tiger/yy/src/unique_files_glb_under6000face_2degree_30ratio_0.01'
+        # dataset_path = '/gemini/user/private/zhaotianhao/data/why_filter_unquantized'
+        # # dataset_path = '/gemini/user/private/zhaotianhao/data/trellis500k_compress_glb'
+        # dataset_path = '/gemini/user/private/zhaotianhao/data/unique_files_glb_under6000face_2degree_30ratio_0.01'
+        dataset_path = '/home/tiger/yy/src/trellis_clean_mesh/mesh_data'
+
         if dataset_path == '/HOME/paratera_xy/pxy1054/HDD_POOL/Trisf/data/mesh/objaverse_200_2000':
             RENDERS_DIR = '/HOME/paratera_xy/pxy1054/HDD_POOL/Trisf/data/mesh_render_img/objaverse_200_2000/renders_cond'
         else:

train_slat_flow_128to512_pointnet_head.py

@@ -0,0 +1,507 @@
+import os
+# os.environ['ATTN_BACKEND'] = 'xformers' # xformers is generally compatible with DDP
+# os.environ["OMP_NUM_THREADS"] = "1"
+# os.environ["MKL_NUM_THREADS"] = "1"
+import torch
+import numpy as np
+import yaml
+from torch.utils.data import DataLoader, DistributedSampler
+from functools import partial
+import torch.nn.functional as F
+from torch.optim import AdamW
+from torch.amp import GradScaler, autocast
+from typing import *
+from transformers import CLIPTextModel, AutoTokenizer, CLIPTextConfig
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+# --- Updated Imports based on VAE script ---
+from dataset_triposf_head import VoxelVertexDataset_edge, collate_fn_pointnet
+from triposf.models.triposf_vae.VoxelFeatureVAE_edge_woself_128to1024_decoder_head import VoxelVAE
+from vertex_encoder import VoxelFeatureEncoder_active_pointnet, ConnectionHead
+from triposf.modules.sparse.basic import SparseTensor
+
+from trellis.models.structured_latent_flow import SLatFlowModel
+from trellis.trainers.flow_matching.sparse_flow_matching_alone import SparseFlowMatchingTrainer
+from safetensors.torch import load_file
+import torch.multiprocessing as mp
+from PIL import Image
+import torch.nn as nn
+
+from triposf.modules.utils import DiagonalGaussianDistribution
+import torchvision.transforms as transforms
+import re
+import contextlib
+
+# --- Distributed Setup Functions ---
+def setup_distributed(backend="nccl"):
+    """Initializes the distributed environment."""
+    if not dist.is_initialized():
+        rank = int(os.environ["RANK"])
+        world_size = int(os.environ["WORLD_SIZE"])
+        local_rank = int(os.environ["LOCAL_RANK"])
+        
+        torch.cuda.set_device(local_rank)
+        dist.init_process_group(backend=backend)
+    
+    return int(os.environ["RANK"]), int(os.environ["LOCAL_RANK"]), int(os.environ["WORLD_SIZE"])
+
+def cleanup_distributed():
+    dist.destroy_process_group()
+
+# --- Modified Trainer Class ---
+class SLatFlowMatchingTrainer(SparseFlowMatchingTrainer):
+    def __init__(self, *args, rank: int, local_rank: int, world_size: int, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.cfg = kwargs.pop('cfg', None)
+        if self.cfg is None:
+            raise ValueError("Configuration dictionary 'cfg' must be provided.")
+            
+        # --- Distributed-related attributes ---
+        self.rank = rank
+        self.local_rank = local_rank
+        self.world_size = world_size
+        self.device = torch.device(f"cuda:{self.local_rank}")
+        self.is_master = (self.rank == 0)
+        self.gradient_accumulation_steps = 8
+
+        self.i_save = self.cfg['training']['save_every']
+        self.save_dir = self.cfg['training']['output_dir']
+
+        self.resolution = 128
+        self.condition_type = 'image'
+        self.is_cond = False
+        self.img_res = 518
+
+        if self.is_master:
+            os.makedirs(self.save_dir, exist_ok=True)
+            print(f"Checkpoints and logs will be saved to: {self.save_dir}")
+
+        # Initialize components and set up for DDP
+        self._init_components(
+            clip_model_path=self.cfg['training'].get('clip_model_path', None),
+            dinov2_model_path=self.cfg['training'].get('dinov2_model_path', None),
+            vae_path=self.cfg['training']['vae_path'],
+        )
+
+        self._setup_ddp()
+
+        self.denoiser_checkpoint_path = self.cfg['training'].get('denoiser_checkpoint_path', None)
+
+        trainable_params = list(self.denoiser.parameters())         
+        self.optimizer = AdamW(trainable_params, lr=self.cfg['training'].get('lr', 0.0001), weight_decay=0.0)
+        
+        self.scaler = GradScaler()
+
+        if self.is_master:
+            print("Using Automatic Mixed Precision (AMP) with GradScaler.")
+            
+    def _init_components(self,
+            clip_model_path=None,
+            dinov2_model_path=None,
+            vae_path=None,
+        ):
+        """
+        Initializes VAE, VoxelEncoder (PointNet), and condition models.
+        """
+        # Use the Dataset from the VAE script
+        self.dataset = VoxelVertexDataset_edge(
+            root_dir=self.cfg['dataset']['path'],
+            base_resolution=self.cfg['dataset']['base_resolution'],
+            min_resolution=self.cfg['dataset']['min_resolution'],
+            cache_dir=self.cfg['dataset']['cache_dir'],
+            renders_dir=self.cfg['dataset']['renders_dir'],
+            
+            filter_active_voxels=self.cfg['dataset']['filter_active_voxels'],
+            cache_filter_path=self.cfg['dataset']['cache_filter_path'],
+
+            active_voxel_res=128,
+            pc_sample_number=819200,
+            
+            sample_type=self.cfg['dataset']['sample_type'],
+
+        )
+        
+        self.sampler = DistributedSampler(
+            self.dataset,
+            num_replicas=self.world_size,
+            rank=self.rank,
+            shuffle=True 
+        )
+        
+        # Use collate_fn_pointnet
+        self.dataloader = DataLoader(
+            self.dataset,
+            batch_size=self.cfg['training']['batch_size'],
+            shuffle=False,
+            collate_fn=partial(collate_fn_pointnet,),
+            num_workers=self.cfg['training']['num_workers'],
+            pin_memory=True,
+            sampler=self.sampler,
+            prefetch_factor=4,
+            persistent_workers=True,
+            drop_last=True,
+        )
+
+        def load_file_func(path, device='cpu'):
+            return torch.load(path, map_location=device)
+
+        def _load_and_broadcast(model, load_fn=None, path=None, strict=True):
+            if self.is_master:
+                try:
+                    state = load_fn(path) if load_fn else model.state_dict()
+                except Exception as e:
+                    raise RuntimeError(f"Failed to load weights from {path}: {e}")
+            else:
+                state = None
+
+            dist.barrier()
+            state_b = [state] if self.is_master else [None]
+            dist.broadcast_object_list(state_b, src=0)
+
+            try:
+                # Handle potential key mismatches (e.g. 'module.' prefix)
+                model.load_state_dict(state_b[0], strict=strict)
+            except Exception as e:
+                if self.is_master: print(f"Strict loading failed for {model.__class__.__name__}, trying non-strict: {e}")
+                model.load_state_dict(state_b[0], strict=False)
+
+        # ------------------------- Voxel Encoder (PointNet) -------------------------
+        # Matching the VAE script configuration
+        self.voxel_encoder = VoxelFeatureEncoder_active_pointnet(
+            in_channels=15, 
+            hidden_dim=256, 
+            out_channels=1024, 
+            scatter_type='mean', 
+            n_blocks=5, 
+            resolution=128,
+            add_label=False,
+        ).to(self.device)
+
+        # ------------------------- VAE -------------------------
+        self.vae = VoxelVAE(
+            in_channels=self.cfg['model']['in_channels'],
+            latent_dim=self.cfg['model']['latent_dim'],
+            encoder_blocks=self.cfg['model']['encoder_blocks'],
+            decoder_blocks_vtx=self.cfg['model']['decoder_blocks_vtx'],
+            decoder_blocks_edge=self.cfg['model']['decoder_blocks_edge'],
+            num_heads=8,
+            num_head_channels=64,
+            mlp_ratio=4.0,
+            attn_mode="swin",
+            window_size=8,
+            pe_mode="ape",
+            use_fp16=False,
+            use_checkpoint=True,
+            qk_rms_norm=False,
+            using_subdivide=True,
+            using_attn=self.cfg['model']['using_attn'],
+            attn_first=self.cfg['model'].get('attn_first', True),
+            pred_direction=self.cfg['model'].get('pred_direction', False),
+        ).to(self.device)
+
+
+        # ------------------------- Conditioning -------------------------
+        if self.condition_type == 'text':
+            self.tokenizer = AutoTokenizer.from_pretrained(clip_model_path)
+            if self.is_master:
+                self.condition_model = CLIPTextModel.from_pretrained(clip_model_path)
+            else:
+                config = CLIPTextConfig.from_pretrained(clip_model_path)
+                self.condition_model = CLIPTextModel(config)
+            _load_and_broadcast(self.condition_model)
+
+        elif self.condition_type == 'image':
+            if self.is_master:
+                print("Initializing for IMAGE conditioning (DINOv2).")
+
+            # Update paths as per your environment
+            local_repo_path = "/home/tiger/yy/src/gemini/user/private/zhaotianhao/dinov2_resources/dinov2-main"
+            weights_path = "/home/tiger/yy/src/gemini/user/private/zhaotianhao/dinov2_resources/dinov2_vitl14_reg4_pretrain.pth"
+
+            dinov2_model = torch.hub.load(
+                repo_or_dir=local_repo_path,
+                model='dinov2_vitl14_reg',
+                source='local',
+                pretrained=False 
+            )
+            self.condition_model = dinov2_model
+
+            _load_and_broadcast(self.condition_model, load_fn=torch.load, path=weights_path)
+            
+            self.image_cond_model_transform = transforms.Compose([
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            ])
+        else:
+            raise ValueError(f"Unsupported condition type: {self.condition_type}")
+
+        self.condition_model.to(self.device).eval()
+        for p in self.condition_model.parameters(): p.requires_grad = False
+
+        # ------------------------- Load VAE/Encoder Weights -------------------------
+        # Load weights corresponding to the logic in VAE script's `load_pretrained_woself`
+        # Assuming checkpoint contains 'vae' and 'voxel_encoder' keys
+        _load_and_broadcast(self.vae,
+                            load_fn=lambda p: load_file_func(p)['vae'],
+                            path=vae_path)
+
+        _load_and_broadcast(self.voxel_encoder,
+                            load_fn=lambda p: load_file_func(p)['voxel_encoder'],
+                            path=vae_path)
+
+        self.vae.eval()
+        self.voxel_encoder.eval()
+        for p in self.vae.parameters(): p.requires_grad = False
+        for p in self.voxel_encoder.parameters(): p.requires_grad = False
+
+    def _load_denoiser(self):
+        """Loads a checkpoint for the denoiser."""
+        path = self.denoiser_checkpoint_path
+        if not path or not os.path.isfile(path):
+            if self.is_master:
+                print("No valid checkpoint path provided for denoiser. Starting from scratch.")
+            return
+
+        if self.is_master:
+            print(f"Loading denoiser checkpoint from: {path}")
+            checkpoint = torch.load(path, map_location=self.device)
+        else:
+            checkpoint = None
+
+        dist.barrier()
+        dist_list = [checkpoint] if self.is_master else [None]
+        dist.broadcast_object_list(dist_list, src=0)
+        checkpoint = dist_list[0]
+        
+        try:
+            self.denoiser.module.load_state_dict(checkpoint['denoiser'])
+            if self.is_master: print("Denoiser weights loaded successfully.")
+        except Exception as e:
+            if self.is_master: print(f"[ERROR] Failed to load denoiser state_dict: {e}")
+             
+        if 'step' in checkpoint and self.is_master:
+            print(f"Checkpoint from step {checkpoint['step']}.")
+        
+        dist.barrier()
+
+    def _setup_ddp(self):
+        """Sets up DDP and DataLoaders."""
+        self.denoiser = self.denoiser.to(self.device)
+        self.denoiser = DDP(self.denoiser, device_ids=[self.local_rank])
+        
+        for param in self.denoiser.parameters():
+            param.requires_grad = True
+
+    @torch.no_grad()
+    def encode_image(self, images) -> torch.Tensor:
+        if isinstance(images, torch.Tensor):
+            batch_tensor = images.to(self.device)
+        elif isinstance(images, list):
+            assert all(isinstance(i, Image.Image) for i in images), "Image list should be list of PIL images"
+            image = [i.resize((518, 518), Image.LANCZOS) for i in images]
+            image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
+            image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
+            batch_tensor = torch.stack(image).to(self.device)
+        else:
+            raise ValueError(f"Unsupported type of image: {type(image)}")
+
+        if batch_tensor.shape[-2:] != (518, 518):
+             batch_tensor = F.interpolate(batch_tensor, (518, 518), mode='bicubic', align_corners=False)
+
+        features = self.condition_model(batch_tensor, is_training=True)['x_prenorm']
+        patchtokens = F.layer_norm(features, features.shape[-1:])
+        return patchtokens
+        
+    def process_batch(self, batch):
+        preprocessed_images = batch['image'] 
+        cond_ = self.encode_image(preprocessed_images)
+        return cond_
+        
+    def train(self, num_epochs=1000):
+        # 1. 无条件生成的准备工作 (和之前一样)
+        if self.is_cond == False:
+            if self.condition_type == 'text':
+                txt = ['']
+                encoding = self.tokenizer(txt, max_length=77, padding='max_length', truncation=True, return_tensors='pt')
+                tokens = encoding['input_ids'].to(self.device)
+                with torch.no_grad():
+                    cond_ = self.condition_model(input_ids=tokens).last_hidden_state
+            else:
+                blank_img = Image.fromarray(np.zeros((self.img_res, self.img_res, 3), dtype=np.uint8))
+                with torch.no_grad():
+                    dummy_cond = self.encode_image([blank_img])
+                    cond_ = torch.zeros_like(dummy_cond)
+                if self.is_master: print(f"Generated unconditional image prompt with shape: {cond_.shape}")
+
+        self._load_denoiser()
+        self.denoiser.train()
+        
+        # 获取全局步数
+        global_step = 0
+        if self.denoiser_checkpoint_path:
+            match = re.search(r'step(\d+)', self.denoiser_checkpoint_path)
+            if match:
+                global_step = int(match.group(1))
+
+        accum_steps = self.gradient_accumulation_steps
+        if self.is_master:
+            print(f"Training with Gradient Accumulation Steps: {accum_steps}")
+
+        # 确保循环开始前梯度清零
+        self.optimizer.zero_grad(set_to_none=True)
+
+        for epoch in range(num_epochs):
+            self.dataloader.sampler.set_epoch(epoch)
+            epoch_losses = []
+            
+            # 遍历数据
+            for i, batch in enumerate(self.dataloader):
+                
+                # --- A. 数据准备 ---
+                if self.is_cond and self.condition_type == 'image':
+                    cond_ = self.process_batch(batch)
+                
+                point_cloud = batch['point_cloud_128'].to(self.device)
+                active_coords = batch['active_voxels_128'].to(self.device)
+                
+                batch_size = int(active_coords[:, 0].max().item() + 1)
+                if cond_.shape[0] != batch_size:
+                    cond_ = cond_.expand(batch_size, -1, -1).contiguous().to(self.device)
+                else:
+                    cond_ = cond_.to(self.device)
+
+                # --- B. 前向传播 & Loss计算 ---
+                with autocast(device_type='cuda', dtype=torch.bfloat16):
+                    with torch.no_grad():
+                        active_voxel_feats = self.voxel_encoder(
+                            p=point_cloud,
+                            sparse_coords=active_coords,
+                            res=128,
+                            bbox_size=(-0.5, 0.5),
+                        )
+                        sparse_input = SparseTensor(
+                            feats=active_voxel_feats,
+                            coords=active_coords.int()
+                        )
+                        latent_128, posterior = self.vae.encode(sparse_input)
+
+                    terms, _ = self.training_losses(x_0=latent_128, cond=cond_)
+                    loss = terms['loss']
+                    
+                    # [重点] Loss 除以累积步数
+                    loss = loss / accum_steps
+
+                # --- C. 反向传播 ---
+                # 注意：这里没有 no_sync，每次都会同步梯度
+                self.scaler.scale(loss).backward()
+                
+                # 记录还原后的 Loss 用于显示
+                current_real_loss = loss.item() * accum_steps
+                epoch_losses.append(current_real_loss)
+
+                # --- D. 梯度累积判断与更新 ---
+                if (i + 1) % accum_steps == 0:
+                    # 如果需要 clip_grad_norm，可以在这里加:
+                    # self.scaler.unscale_(self.optimizer)
+                    # torch.nn.utils.clip_grad_norm_(self.denoiser.parameters(), max_norm=1.0)
+                    
+                    self.scaler.step(self.optimizer)
+                    self.scaler.update()
+                    self.optimizer.zero_grad(set_to_none=True)
+                    
+                    global_step += 1
+                    
+                    # --- Logging (只在更新步进行) ---
+                    if self.is_master:
+                        if global_step % 10 == 0:
+                            print(f"Epoch {epoch+1} Step {global_step}: "
+                                  f"Batch_Loss = {current_real_loss:.4f}, "
+                                  f"Epoch_Mean = {np.mean(epoch_losses):.4f}")
+                        
+                        if global_step % self.i_save == 0:
+                            checkpoint = {
+                                'denoiser': self.denoiser.module.state_dict(),
+                                'step': global_step
+                            }
+                            loss_str = f"{np.mean(epoch_losses):.6f}".replace('.', '_')
+                            save_path = os.path.join(self.save_dir, f"checkpoint_step{global_step}_loss{loss_str}.pt")
+                            torch.save(checkpoint, save_path)
+                            print(f"Saved checkpoint to {save_path}")
+
+            # --- E. 处理 Epoch 结束时的残留 Batch (Leftovers) ---
+            # 如果 dataloader 长度不能被 accum_steps 整除，且 drop_last=False，
+            # 这里需要把最后累积的一点梯度更新掉。
+            # (如果你的 dataloader 设置了 drop_last=True 且总数够整除，这里不会触发，但写上比较保险)
+            if (i + 1) % accum_steps != 0:
+                self.scaler.step(self.optimizer)
+                self.scaler.update()
+                self.optimizer.zero_grad(set_to_none=True)
+                # 注意：这里通常不增加 global_step 或者看你习惯，
+                # 因为这是一个“不完整”的 step，通常梯度也是不对等的（因为除数还是accum_steps）
+                # 所以很多实现为了稳定，直接设置 drop_last=True 避开这种情况。
+
+            if self.is_master:
+                avg_loss = np.mean(epoch_losses) if epoch_losses else 0
+                log_path = os.path.join(self.save_dir, "loss_log.txt")
+                with open(log_path, "a") as f:
+                    f.write(f"Epoch {epoch+1}, Step {global_step}, AvgLoss {avg_loss:.6f}\n")
+            
+            dist.barrier()
+
+        if self.is_master:
+            print("Training complete.")
+
+def main():
+    # if mp.get_start_method(allow_none=True) != 'spawn':
+    #     mp.set_start_method('spawn', force=True)
+
+    # if mp.get_start_method(allow_none=True) != 'forkserver':
+    #     mp.set_start_method('forkserver', force=True)
+
+    rank, local_rank, world_size = setup_distributed()
+    torch.manual_seed(42+rank) 
+    np.random.seed(42+rank)
+    
+    # Path to your config
+    config_path = "/home/tiger/yy/src/Michelangelo-master/config_slat_flow_128to512_pointnet_head.yaml"
+    with open(config_path) as f:
+        cfg = yaml.safe_load(f)
+
+    # Initialize Flow Model (on CPU first)
+    diffusion_model = SLatFlowModel(
+        resolution=cfg['flow']['resolution'],
+        in_channels=cfg['flow']['in_channels'],
+        out_channels=cfg['flow']['out_channels'],
+        model_channels=cfg['flow']['model_channels'],
+        cond_channels=cfg['flow']['cond_channels'],
+        num_blocks=cfg['flow']['num_blocks'],
+        num_heads=cfg['flow']['num_heads'],
+        mlp_ratio=cfg['flow']['mlp_ratio'],
+        patch_size=cfg['flow']['patch_size'],
+        num_io_res_blocks=cfg['flow']['num_io_res_blocks'],
+        io_block_channels=cfg['flow']['io_block_channels'],
+        pe_mode=cfg['flow']['pe_mode'],
+        qk_rms_norm=cfg['flow']['qk_rms_norm'],
+        qk_rms_norm_cross=cfg['flow']['qk_rms_norm_cross'],
+        use_fp16=cfg['flow'].get('use_fp16', False),
+    )
+
+    torch.manual_seed(42 + rank)
+    np.random.seed(42 + rank)
+    
+    trainer = SLatFlowMatchingTrainer(
+        denoiser=diffusion_model,
+        t_schedule=cfg['t_schedule'],
+        sigma_min=cfg['sigma_min'],
+        cfg=cfg,
+        rank=rank,
+        local_rank=local_rank,
+        world_size=world_size,
+    )
+
+    trainer.train()
+    cleanup_distributed()
+
+if __name__ == '__main__':
+    main()

trellis/__init__.py

@@ -2,5 +2,5 @@ from . import models
 from . import modules
 from . import pipelines
 from . import renderers
-from . import representations
+# from . import representations
 from . import utils