#3D视觉基础：点云处理、深度估计、立体视觉详解

import numpy as np
import open3d as o3d

class PointCloudProcessor:
    """简洁的Open3D点云处理器封装"""
    def create_from_numpy(self, points: np.ndarray) -> o3d.geometry.PointCloud:
        """从(N,3)的numpy数组创建点云"""
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)
        return pcd

    def preprocess(self, pcd: o3d.geometry.PointCloud, 
                   voxel_size=0.05, nb_neighbors=20, std_ratio=2.0) -> o3d.geometry.PointCloud:
        """
        一站式预处理：下采样 → 法向量估计 → 离群点去除
        """
        # 体素下采样（解决密度不均，保留几何特征）
        down_pcd = pcd.voxel_down_sample(voxel_size=voxel_size)
        # 法向量估计（为后续分割、重建做准备）
        down_pcd.estimate_normals()
        # 统计离群点去除（过滤激光雷达/扫描带来的噪点）
        filtered_pcd, _ = down_pcd.remove_statistical_outlier(
            nb_neighbors=nb_neighbors, std_ratio=std_ratio
        )
        return filtered_pcd

# 快速测试
if __name__ == "__main__":
    # 生成10000个高斯分布的随机点云
    noisy_points = np.random.randn(10000, 3) * 0.5
    pcd_raw = PointCloudProcessor().create_from_numpy(noisy_points)
    pcd_filtered = PointCloudProcessor().preprocess(pcd_raw)
    # 对比可视化（本地运行时打开draw_geometries）
    # o3d.visualization.draw_geometries([pcd_raw, pcd_filtered], window_name="Raw vs Filtered")
    print(f"原始点数: {len(pcd_raw.points)}, 预处理后点数: {len(pcd_filtered.points)}")

import torch
import torch.nn as nn

class SimplePointNet(nn.Module):
    """简化版PointNet（不含T-net空间变换模块）"""
    def __init__(self, num_classes=10):
        super().__init__()
        # 共享MLP：逐点提取特征 (B,N,3) → (B,N,1024)
        self.point_mlp = nn.Sequential(
            nn.Conv1d(3, 64, kernel_size=1),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Conv1d(64, 128, kernel_size=1),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Conv1d(128, 1024, kernel_size=1),
            nn.BatchNorm1d(1024)
        )
        # 全局特征聚合
        self.global_pool = nn.MaxPool1d(kernel_size=1024)
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(1024, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, num_classes)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (batch_size, num_points, 3)
        Returns:
            logits: (batch_size, num_classes)
        """
        # 调整维度以适配Conv1d (B,N,3) → (B,3,N)
        x = x.transpose(2, 1)
        # 逐点特征提取
        x = self.point_mlp(x)
        # 全局特征聚合 (B,1024,N) → (B,1024,1) → (B,1024)
        x = self.global_pool(x).squeeze(-1)
        # 分类
        return self.classifier(x)

# 快速测试维度
if __name__ == "__main__":
    test_input = torch.randn(8, 1024, 3)  # batch=8, 点数=1024
    model = SimplePointNet(num_classes=40)  # ModelNet40数据集的40类
    output = model(test_input)
    print(f"输入维度: {test_input.shape}, 输出维度: {output.shape}")

import cv2
import numpy as np
import matplotlib.pyplot as plt

def sgbm_stereo_depth(left_path: str, right_path: str, 
                       baseline=0.2, focal_length=720):
    """
    用SGBM计算双目图像的视差图和深度图
    Args:
        left_path/right_path: 左右图像路径
        baseline: 相机基线距离（米）
        focal_length: 相机焦距（像素）
    """
    # 1. 读取并预处理灰度图
    left = cv2.imread(left_path, cv2.IMREAD_GRAYSCALE)
    right = cv2.imread(right_path, cv2.IMREAD_GRAYSCALE)
    h, w = left.shape[:2]

    # 2. 创建SGBM匹配器
    stereo = cv2.StereoSGBM_create(
        minDisparity=0,          # 最小视差
        numDisparities=16*5,     # 视差范围（必须是16的倍数）
        blockSize=5,             # 匹配块大小（奇数）
        P1=8*3*5**2,             # 小视差变化的惩罚项
        P2=32*3*5**2,            # 大视差变化的惩罚项
        disp12MaxDiff=1,         # 左右一致性检查的最大差异
        uniquenessRatio=15,      # 唯一匹配的阈值
        preFilterCap=63          # 预处理的截断值
    )

    # 3. 计算视差图（归一化到[0,1]）
    disparity = stereo.compute(left, right).astype(np.float32) / 16.0
    disparity_normalized = (disparity - disparity.min()) / (disparity.max() - disparity.min() + 1e-5)

    # 4. 视差转深度（避免除零）
    depth = (baseline * focal_length) / (disparity + 1e-5)
    depth_normalized = (depth - depth.min()) / (depth.max() - depth.min() + 1e-5)

    # 5. 可视化
    plt.figure(figsize=(15, 5))
    plt.subplot(131), plt.imshow(left, cmap='gray'), plt.title('Left Image')
    plt.subplot(132), plt.imshow(disparity_normalized, cmap='jet'), plt.title('Normalized Disparity')
    plt.subplot(133), plt.imshow(depth_normalized, cmap='jet'), plt.title('Normalized Depth')
    plt.tight_layout()
    plt.show()

    return disparity, depth