import open3d as o3d
import numpy as np
from matplotlib import pyplot as plt
import cv2

color_raw = o3d.io.read_image("aucpuncture2point/configs/using_img/color.png")
depth_raw = o3d.io.read_image("aucpuncture2point/configs/using_img/depth.png")

# 模拟空洞：将指定区域的深度值设置为无效值（0或者其他很大的值）
def create_occlusion(depth_raw, occlusion_rect=(290, 160, 370, 270)):
    """
    Create an occlusion (hole) in the depth image by setting a specific rectangular region to zero (no depth).
    The rect is defined by (x_min, y_min, x_max, y_max).
    """
    depth_image = np.asarray(depth_raw)
    x_min, y_min, x_max, y_max = occlusion_rect
    
    # 将指定区域的深度设置为无效（可以选择设置为0或其他特定的无效值）
    depth_image[y_min:y_max, x_min:x_max] = 0  # 模拟遮挡，深度值为0
    
    return o3d.geometry.Image(depth_image)

# 在深度图像中创建遮挡
depth_with_occlusion = create_occlusion(depth_raw)

# 将深度图像转换为 NumPy 数组
depth_image = np.asarray(depth_with_occlusion)

# 创建掩膜：深度值为0的区域
mask = (depth_image == 0).astype(np.uint8)

# 使用 OpenCV 的 inpaint 函数进行深度修复
depth_image_inpainted = cv2.inpaint(depth_image, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)

# 将修复后的深度图像转换回 Open3D 图像
depth_fixed = o3d.geometry.Image(depth_image_inpainted)


rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
    color_raw, depth_fixed, convert_rgb_to_intensity=False)

intrinsics = o3d.camera.PinholeCameraIntrinsic()
intrinsics.set_intrinsics(640, 480, 453.17746, 453.17746, 325.943024, 243.559982)

# 从RGBD生成点云
pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, intrinsics)
pcd.transform([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
pcd = pcd.voxel_down_sample(voxel_size=0.01)
# 估计点云的法向量
pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))

# 计算几何中心
centroid = np.mean(np.asarray(pcd.points), axis=0)

# 调整法向量方向，使其指向人体内侧
normals = np.asarray(pcd.normals)
for i, normal in enumerate(normals):
    point = np.asarray(pcd.points)[i]
    if np.dot(normal, centroid - point) < 0:  # 法向量与到中心向量的点积如果小于0，则需要翻转
        normals[i] *= -1

pcd.normals = o3d.utility.Vector3dVector(normals)
# 创建一个小球表示几何中心
sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.02)  # 半径为0.02的小球
sphere.paint_uniform_color([1, 0, 0])  # 颜色为红色
sphere.translate(centroid)  # 将小球平移到几何中心

# 可视化点云和几何中心
o3d.visualization.draw_geometries([pcd, sphere], point_show_normal=True)