display_model.py

import numpy as np
from plyfile import PlyData, PlyElement

def ply_to_glb(ply_file, glb_file):
    print("Converting PLY to GLB...")
    # Import trimesh here to ensure it's only required when this function is called
    import trimesh
    # Load the PLY file with trimesh
    try:
        mesh = trimesh.load(ply_file)
        # Export to GLB format
        mesh.export(glb_file, file_type='glb')
        print("Conversion finished.")
        return "PLY to GLB conversion complete."
    except Exception as e:
        # In case of any issue, print the error
        print(f"Error during conversion: {e}")
        return "Conversion failed."


def merge_box_to_ply(ply_file, box_ply_file):
    pass

def add_1box_to_ply(box, ply_file, new_ply_file, line_width=0.05, obj_id=1):
    print("adding 1 box to ply...")
    print("ply_file:",ply_file)
    print("new_ply_file:",new_ply_file)
    # box format: [xmin, ymin, zmin, xmax, ymax, zmax]
    xmin, ymin, zmin, xmax, ymax, zmax = box
    box_coords = np.array(
            [[xmin, ymin, zmin], #0
             [xmin-line_width, ymin-line_width, zmin], #1
             [xmax, ymin, zmin], #2
             [xmax+line_width, ymin-line_width, zmin], #3
             [xmax, ymax, zmin], #4
             [xmax+line_width, ymax+line_width, zmin], #5
             [xmin, ymax, zmin], #6
             [xmin-line_width, ymax+line_width, zmin], #7
             [xmin, ymin, zmax], #8
             [xmin-line_width, ymin-line_width, zmax], #9
             [xmax, ymin, zmax], #10
             [xmax+line_width, ymin-line_width, zmax], #11
             [xmax, ymax, zmax], #12
             [xmax+line_width, ymax+line_width, zmax], #13
             [xmin, ymax, zmax], #14
             [xmin-line_width, ymax+line_width, zmax] #15
            ])

    # read in ply
    with open(ply_file, 'rb') as f:
        ply_data = PlyData.read(f)
        vertices = ply_data['vertex'].data

    # handle vertices and update
    color = [255, 0, 0]
    box_vertices = np.zeros(len(box_coords), dtype=vertices.dtype)
    box_vertices['x'] = [coord[0] for coord in box_coords]
    box_vertices['y'] = [coord[1] for coord in box_coords]
    box_vertices['z'] = [coord[2] for coord in box_coords]
    box_vertices['red'] = [color[0]] * 16
    box_vertices['green'] = [color[1]] * 16
    box_vertices['blue'] = [color[2]] * 16
    box_vertices['alpha'] = [obj_id] * 16

    # 将新的顶点数据添加到原始顶点数据后面
    updated_vertices = np.concatenate((vertices, box_vertices))

    # 创建包含新顶点的PlyElement对象
    updated_vertex_element = PlyElement.describe(updated_vertices, 'vertex')

    # 将更新后的PlyElement对象替换原始的顶点数据
    # ply_data['vertex'] = updated_vertex_element

    # get the number of original vertices:
    num_origin_vertices = len(vertices)

    # define connections of new faces
    
    box_connections=[
        [num_origin_vertices+0,  num_origin_vertices+1,  num_origin_vertices+3 ],
        [num_origin_vertices+0,  num_origin_vertices+3,  num_origin_vertices+2 ],
        [num_origin_vertices+2,  num_origin_vertices+3,  num_origin_vertices+5 ],
        [num_origin_vertices+2,  num_origin_vertices+5,  num_origin_vertices+4 ],
        [num_origin_vertices+4,  num_origin_vertices+5,  num_origin_vertices+7 ],
        [num_origin_vertices+4,  num_origin_vertices+7,  num_origin_vertices+6 ],
        [num_origin_vertices+0,  num_origin_vertices+1,  num_origin_vertices+7 ],
        [num_origin_vertices+0,  num_origin_vertices+7,  num_origin_vertices+6 ],
        [num_origin_vertices+0,  num_origin_vertices+1,  num_origin_vertices+9 ],
        [num_origin_vertices+0,  num_origin_vertices+9,  num_origin_vertices+8 ],
        [num_origin_vertices+2,  num_origin_vertices+3,  num_origin_vertices+11],
        [num_origin_vertices+2,  num_origin_vertices+11, num_origin_vertices+10],
        [num_origin_vertices+4,  num_origin_vertices+5,  num_origin_vertices+13],
        [num_origin_vertices+4,  num_origin_vertices+13, num_origin_vertices+12],
        [num_origin_vertices+6,  num_origin_vertices+7,  num_origin_vertices+15],
        [num_origin_vertices+6,  num_origin_vertices+15, num_origin_vertices+14],
        [num_origin_vertices+8,  num_origin_vertices+9,  num_origin_vertices+11],
        [num_origin_vertices+8,  num_origin_vertices+11, num_origin_vertices+10],
        [num_origin_vertices+10, num_origin_vertices+11, num_origin_vertices+13],
        [num_origin_vertices+10, num_origin_vertices+13, num_origin_vertices+12],
        [num_origin_vertices+12, num_origin_vertices+13, num_origin_vertices+15],
        [num_origin_vertices+12, num_origin_vertices+15, num_origin_vertices+14],
        [num_origin_vertices+8,  num_origin_vertices+9,  num_origin_vertices+15],
        [num_origin_vertices+8,  num_origin_vertices+15, num_origin_vertices+14]
    ]

    # handle faces and update
    faces = ply_data['face'].data
    box_faces = np.zeros(len(box_connections), dtype=faces.dtype)
    box_faces['vertex_indices'] = box_connections

    # 将新的face数据添加到原始顶点数据后面
    updated_faces = np.concatenate((faces, box_faces))

    # 创建包含新顶点的PlyElement对象
    updated_face_element = PlyElement.describe(updated_faces, 'face')

    # 将更新后的PlyElement对象替换原始的顶点数据
    # ply_data['face'] = updated_face_element

    new_ply_data = PlyData([updated_vertex_element, updated_face_element])

    # 将更新后的PlyData对象写回Ply文件
    with open(new_ply_file, 'wb') as f:
        new_ply_data.write(f)

    print("add 1 box to ply finished.")


def ply_to_obj(ply_file, obj_file, mtl_file):
    # 读取PLY文件
    with open(ply_file, 'rb') as f:
        plydata = PlyData.read(f)

        # 获取顶点和面数据
        vertices = np.vstack([plydata['vertex'][prop] for prop in ['x', 'y', 'z']]).T
        colors = np.vstack([plydata['vertex'][prop] for prop in ['red', 'green', 'blue', 'alpha']]).T/255.0
        faces = plydata['face']['vertex_indices']

    # 写入OBJ文件
    with open(obj_file, 'w') as f:
        # 写入依赖的mtl文件（颜色）
        f.write("mtllib %s\n"%mtl_file.split('/')[-1])

        # 写入顶点信息
        for vertex in vertices:
            f.write(f"v {' '.join(map(str, vertex))}\n")

        # 写入颜色信息
        for idx in range(len(vertices)):
            f.write("usemtl mat%d\n"%(idx+1))

        # 写入面信息
        for face in faces:
            f.write("f")
            for vertex_index in face:
                f.write(f" {vertex_index + 1}")  # OBJ文件索引从1开始
            f.write("\n")

    # 写入mtl文件
    with open(mtl_file, 'w') as f:
        for idx, color in enumerate(colors):
            f.write("newmtl mat%d\n" % (idx+1))
            f.write("Kd %f %f %f\n\n" % (color[0], color[1],color[2]))

if __name__ == "__main__":
    # ply_to_obj("./scenes/scene0132_00_vh_clean_2_aligned.ply", "./scenes/scene0132_00_vh_clean_2_aligned.obj", "./scenes/scene0132_00_vh_clean_2_aligned_colors.mtl")
    add_1box_to_ply([0,0,0,1,1,1],"scenes\scene0132_00_vh_clean_2_aligned.ply","scenes\scene0132_00_add1box.ply")