Iterate over triangles, get their neighbors, and show as movie

mesh_near_neighborgs.py

import open3d as o3d
import numpy as np
from PIL import Image
from tqdm import tqdm
from collections import defaultdict
import cv2
import copy


def capture_frame(vis):
    img = o3d.visualization.gui.Application.instance.render_to_image(vis.scene, 1024, 768)
    img = np.asarray(img)
    return img


def visualize_colored_mesh(mesh, triangle_index, adjacency_list, vis):
    # Paint the current triangle red
    triangle = mesh.triangles[triangle_index]
    for vertex_index in triangle:
        mesh.vertex_colors[vertex_index] = [1, 0, 0]  # Red

    # Paint neighboring triangles blue
    neighbors = adjacency_list[triangle_index]
    # print(len(neighbors))
    for neighbor_index in neighbors:
        neighbor_triangle = mesh.triangles[neighbor_index]
        for vertex_index in neighbor_triangle:
            mesh.vertex_colors[vertex_index] = [0, 0, 1]  # Blue

    mesh.compute_vertex_normals()

    vis.remove_geometry("mesh")
    vis.add_geometry("mesh", mesh)
    #vis.poll_events()
    #vis.update_renderer()
    vis.post_redraw()
    image = capture_frame(vis)

    return image


def visualize_white_mesh(mesh, triangle_index, adjacency_list, vis):
    mesh.paint_uniform_color([1, 1, 1])
    return visualize_colored_mesh(mesh, triangle_index, adjacency_list, vis)


def compute_adjacency_list(mesh):
    adjacency_list = defaultdict(list)
    for triangle_index, triangle in enumerate(mesh.triangles):
        for i in range(3):
            edge = tuple(sorted((triangle[i], triangle[(i + 1) % 3])))
            adjacency_list[edge].append(triangle_index)

    triangle_adjacency = defaultdict(list)
    for edge, triangle_indices in adjacency_list.items():
        if len(triangle_indices) == 2:  # Shared edge
            t1, t2 = triangle_indices
            triangle_adjacency[t1].append(t2)
            triangle_adjacency[t2].append(t1)

    return triangle_adjacency


def compute_uv_coordinates(vertices):
    # Assuming a planar mesh lying on the XY plane
    # UV coordinates are mapped based on the XY positions of vertices
    # Normalize vertex positions to range [0, 1]
    min_x = np.min(vertices[:, 0])
    max_x = np.max(vertices[:, 0])
    min_y = np.min(vertices[:, 1])
    max_y = np.max(vertices[:, 1])

    normalized_vertices = (vertices - [min_x, min_y, 0]) / [max_x - min_x, max_y - min_y, 1]

    # UV coordinates are simply XY positions scaled to range [0, 1]
    uv_coordinates = normalized_vertices[:, :2]
    
    return uv_coordinates


def read_color_mesh(obj_path):
    # Read the mesh
    mesh = o3d.io.read_triangle_mesh(obj_path, True, print_progress=True)
    texture_image = np.asarray(mesh.textures).squeeze()

    # Ensure the mesh has vertex normals
    mesh.compute_vertex_normals()

    uv_coords = np.asarray(mesh.triangle_uvs)
    uv_coords = uv_coords - np.floor(uv_coords)

    vertex_colors = np.sort(np.zeros((np.asarray(mesh.vertices).shape[0], 3)))

    triangles = np.asarray(mesh.triangles)
    for i, tri in enumerate(triangles):
        for j in range(3):
            # Get the UV coordinates of the j-th vertex of the i-th triangle
            uv = uv_coords[i * 3 + j]
            
            # Map UV coordinates to texture image coordinates
            tex_x = int(uv[0] * (texture_image.shape[1] - 1))
            #tex_y = int((1 - uv[1]) * (texture_image.shape[0] - 1))  # Flip UV vertically
            tex_y = int(uv[1] * (texture_image.shape[0] - 1))  # Flip UV vertically
            
            # Get the texture color
            color = texture_image[tex_y, tex_x, :3] / 255.0  # Normalize color
            
            # Assign the texture color to the corresponding vertex
            vertex_colors[tri[j], :] = color

    # Apply colors to mesh vertices

    mesh.vertex_colors = o3d.utility.Vector3dVector(vertex_colors)
    return mesh

# Load a sample mesh with color or create one
obj_path = o3d.data.MonkeyModel().path
#obj_path = o3d.data.DamagedHelmetModel().path
#obj_path = "/home/pedro/Downloads/winter_girl/winter_girl_C0-L5_deq_tri.obj"


mesh = read_color_mesh(obj_path)
#mesh = o3d.io.read_triangle_mesh(obj_path, False, print_progress=True)

print("meshes=", dir(mesh))
print("mesh.has_vertex_colors=", mesh.has_vertex_colors())
print("mesh.vertex_colors=", np.asarray(mesh.vertex_colors))
print("mesh.triangles=", mesh.triangles)
print("mesh.textures=", mesh.textures)


adjacency_list = compute_adjacency_list(mesh)

#o3d.visualization.draw([mesh])

o3d.visualization.gui.Application.instance.initialize()

# Create a visualizer
vis = o3d.visualization.O3DVisualizer("Open3D", 1024, 768)

vis.add_geometry("mesh", mesh)

vis.set_background((1.0, 1.0, 1.0, 1.0), None)
vis.reset_camera_to_default()


white_mesh = copy.copy(mesh)
white_mesh.paint_uniform_color([1, 1, 1])
white_mesh.textures[0] = np.ones(np.asarray(white_mesh.textures[0]).shape)


out = None

# Iterate over each triangle and capture the coloring
for triangle_index in tqdm(range(len(mesh.triangles))):
    color = visualize_colored_mesh(mesh, triangle_index, adjacency_list, vis)
    white = visualize_white_mesh(white_mesh, triangle_index, adjacency_list, vis)
    frame = np.hstack((color, white))
    if out == None:
        frame_height, frame_width, _ = frame.shape
        out = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 1, (frame_width, frame_height))
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    out.write(frame)
    if triangle_index == 50:
        break

vis.close()
out.release()