Last active
June 9, 2024 16:52
-
-
Save pedrogarciafreitas/7ccffa40e5e1445811b64a837b92d367 to your computer and use it in GitHub Desktop.
Iterate over triangles, get their neighbors, and show as movie
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment