print_setting/sui_01.py

import open3d as o3d
import os
import numpy as np
from scipy.spatial.transform import Rotation
import sys

import argparse
# import cv2
import matplotlib.pyplot as plt
import numpy as np
from numba import njit, prange
import time


# 核心计算函数（支持Numba加速）
@njit(fastmath=True, cache=True)
def calculate_rotation_z(angle_x, angle_y, angle_z, points, cos_cache, sin_cache, angle_step):
    """计算单个旋转组合后的重心Z坐标（无显式平移）"""
    # 获取预计算的三角函数值
    idx_x = angle_x // angle_step
    idx_y = angle_y // angle_step
    idx_z = angle_z // angle_step

    cos_x = cos_cache[idx_x]
    sin_x = sin_cache[idx_x]
    cos_y = cos_cache[idx_y]
    sin_y = sin_cache[idx_y]
    cos_z = cos_cache[idx_z]
    sin_z = sin_cache[idx_z]

    # 构造旋转矩阵（展开矩阵乘法优化）
    # R = Rz @ Ry @ Rx
    # 计算矩阵元素（手动展开矩阵乘法）
    m00 = cos_z * cos_y
    m01 = cos_z * sin_y * sin_x - sin_z * cos_x
    m02 = cos_z * sin_y * cos_x + sin_z * sin_x

    m10 = sin_z * cos_y
    m11 = sin_z * sin_y * sin_x + cos_z * cos_x
    m12 = sin_z * sin_y * cos_x - cos_z * sin_x

    m20 = -sin_y
    m21 = cos_y * sin_x
    m22 = cos_y * cos_x

    # 计算所有点的Z坐标
    z_values = np.empty(points.shape[0], dtype=np.float64)
    for i in prange(points.shape[0]):
        x, y, z = points[i, 0], points[i, 1], points[i, 2]
        # 应用旋转矩阵
        rotated_z = m20 * x + m21 * y + m22 * z
        z_values[i] = rotated_z

    # 计算重心Z（等效于平移后的重心）
    min_z = np.min(z_values)
    avg_z = np.mean(z_values)
    return avg_z - min_z  # 等效于平移后的重心Z坐标


# 并行优化主函数
def parallel_rotation2(points, angle_step=3):
    """
    参数：
    points : numpy.ndarray (N,3) - 三维点云
    angle_step : int - 角度搜索步长（度数）

    返回：
    (best_angle_x, best_angle_y, best_angle_z, min_z)
    """
    points = np.ascontiguousarray(points.astype(np.float64))

    # 生成所有可能角度
    angles = np.arange(0, 360, angle_step)
    n_angles = len(angles)

    # 预计算三角函数值（大幅减少重复计算）
    rads = np.radians(angles)
    cos_cache = np.cos(rads).astype(np.float64)
    sin_cache = np.sin(rads).astype(np.float64)

    # 生成所有角度组合（内存优化版）
    total_combinations = n_angles ** 3
    print(f"Total combinations: {total_combinations:,}")

    # 分块处理以避免内存溢出
    best_z = np.inf
    best_angles = (0, 0, 0)
    batch_size = 10 ** 6  # 根据可用内存调整

    for x_chunk in range(0, n_angles, max(1, n_angles // 4)):
        angles_x = angles[x_chunk:x_chunk + max(1, n_angles // 4)]
        for y_chunk in range(0, n_angles, max(1, n_angles // 4)):
            angles_y = angles[y_chunk:y_chunk + max(1, n_angles // 4)]

            # 生成当前分块的所有组合
            xx, yy, zz = np.meshgrid(angles_x, angles_y, angles)
            current_batch = np.stack([xx.ravel(), yy.ravel(), zz.ravel()], axis=1)

            # 处理子批次
            for i in range(0, len(current_batch), batch_size):
                batch = current_batch[i:i + batch_size]
                results = np.zeros(len(batch), dtype=np.float64)
                _process_batch(batch, points, cos_cache, sin_cache, angle_step, results)

                # 更新最佳结果
                min_idx = np.argmin(results)
                if results[min_idx] < best_z:
                    best_z = results[min_idx]
                    best_angles = tuple(batch[min_idx])
                    print(f"New best: {best_angles} -> Z={best_z:.4f}")

    return (*best_angles, best_z)


@njit(parallel=True, fastmath=True)
def _process_batch(batch, points, cos_cache, sin_cache, angle_step, results):
    for i in prange(len(batch)):
        ax, ay, az = batch[i]
        results[i] = calculate_rotation_z(
            ax, ay, az, points,
            cos_cache, sin_cache, angle_step
        )


class ModelProcessor:
    def __init__(self):

        # argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
        parser = argparse.ArgumentParser()

        parser.add_argument(
            "--id",
            required=False,
        )

        args = parser.parse_args()

        self.id = args.id

        self.mesh = None
        self.asset_dir = f"/home/algo/Documents/datasets/{self.id}"

    def load_model(self):
        """加载并初始化3D模型"""
        # model_path = f"{self.asset_dir}/baked/{self.id}.obj"
        # model_path = f"{self.asset_dir}/repair_{self.id}_mesh.ply"
        model_path = "/data/datasets_20t/8/88884_253283_P65951_6cm_x1.obj"
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model file not found: {model_path}")

        print(model_path)

        mesh_native = o3d.io.read_triangle_mesh(model_path, enable_post_processing=False)
        # self.mesh = o3d.io.read_triangle_mesh(model_path, enable_post_processing=False)

        print("Open3D去重前顶点数:", len(mesh_native.vertices))
        self.mesh = mesh_native.merge_close_vertices(eps=1e-6)

        vertices2 = np.asarray(self.mesh.vertices)
        print("Open3D去重后顶点数:", len(vertices2))
        vertices2_sorted = sorted(
            vertices2.tolist(),
            key=lambda x: (x[0], x[1], x[2])
        )

        if not self.mesh.has_vertex_colors():
            num_vertices = len(self.mesh.vertices)
            self.mesh.vertex_colors = o3d.utility.Vector3dVector(
                np.ones((num_vertices, 3))
            )

        self.uv_array = np.asarray(self.mesh.triangle_uvs)
        # print(f"UV 坐标形状：{self.uv_array.shape}, {self.uv_array[0][1]}")

    def calculate_rotation_and_center_of_mass(self, angle_x, angle_y, angle_z, points):
        """计算某一组旋转角度后的重心"""
        # 计算绕X轴、Y轴和Z轴的旋转矩阵
        R_x = np.array([
            [1, 0, 0],
            [0, np.cos(np.radians(angle_x)), -np.sin(np.radians(angle_x))],
            [0, np.sin(np.radians(angle_x)), np.cos(np.radians(angle_x))]
        ])

        R_y = np.array([
            [np.cos(np.radians(angle_y)), 0, np.sin(np.radians(angle_y))],
            [0, 1, 0],
            [-np.sin(np.radians(angle_y)), 0, np.cos(np.radians(angle_y))]
        ])

        R_z = np.array([
            [np.cos(np.radians(angle_z)), -np.sin(np.radians(angle_z)), 0],
            [np.sin(np.radians(angle_z)), np.cos(np.radians(angle_z)), 0],
            [0, 0, 1]
        ])

        # 综合旋转矩阵
        R = R_z @ R_y @ R_x

        # 执行旋转
        rotated_points = points @ R.T

        # 计算最小z值
        min_z = np.min(rotated_points[:, 2])

        # 计算平移向量，将最小Z值平移到0
        translation_vector = np.array([0, 0, -min_z])
        rotated_points += translation_vector

        # 计算重心
        center_of_mass = np.mean(rotated_points, axis=0)

        return center_of_mass[2], angle_x, angle_y, angle_z

    def parallel_rotation(self, angle_step=3):
        """顺序计算最优旋转角度（单线程）"""
        min_center_of_mass_y = float('inf')
        best_angle_x, best_angle_y, best_angle_z = 0, 0, 0

        # 遍历所有角度组合
        for angle_x in range(0, 360, angle_step):
            for angle_y in range(0, 360, angle_step):
                for angle_z in range(0, 360, angle_step):
                    center_of_mass_z, ax, ay, az = self.calculate_rotation_and_center_of_mass(
                        angle_x, angle_y, angle_z, self.mesh.vertices
                    )
                    if center_of_mass_z < min_center_of_mass_y:
                        min_center_of_mass_y = center_of_mass_z
                        best_angle_x, best_angle_y, best_angle_z = ax, ay, az

        return best_angle_x, best_angle_y, best_angle_z, min_center_of_mass_y

    def process(self):
        """执行完整处理流程"""
        self.load_model()

        try:
            start = time.time()

            # mesh = o3d.geometry.TriangleMesh()
            # mesh.vertices = o3d.utility.Vector3dVector(np.random.rand(100, 3))
            # points = np.asarray(mesh.vertices)

            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(self.mesh.vertices)

            # 自动计算合理体素大小
            raw_points = np.asarray(pcd.points)
            bounds = np.ptp(raw_points, axis=0)
            voxel_size = np.max(bounds) / 50  # 默认取最大边长的2%

            # 执行下采样并验证
            pcd_downsampled = pcd.voxel_down_sample(voxel_size)
            if len(pcd_downsampled.points) < 10:  # 最少保留10个点
                raise RuntimeError(f"下采样失败：voxel_size={voxel_size:.3f}过大")

            print(f"下采样后点数: {len(pcd_downsampled.points)} (voxel_size={voxel_size:.3f})")

            # pcd.paint_uniform_color([1,0,0])  # 原始红色
            # pcd_downsampled.paint_uniform_color([0,0,1])  # 采样后蓝色
            # o3d.visualization.draw_geometries([pcd, pcd_downsampled])

            # 继续后续处理
            points = np.asarray(pcd_downsampled.points)

            best_angle_x, best_angle_y, best_angle_z, min_z = parallel_rotation2(points, angle_step=5)
            print("best=", best_angle_x, best_angle_y, best_angle_z, min_z)
            print(time.time() - start)

        except Exception as e:
            print(f"Error during processing: {str(e)}")
            raise


if __name__ == "__main__":
    ModelProcessor().process()