Loads 3D model with UV coordinates, assign random colors to texture charts, render 3D model with randomly colored charts, render 2D texture atlas with randomly colored charts

show_3d_model_charts.py

# pip install trimesh pyvista numpy scipy matplotlib pillow pyglet<2
import trimesh
import numpy as np
from PIL import ImageDraw
from scipy.spatial import KDTree
from scipy.sparse.csgraph import connected_components
from scipy.sparse import coo_matrix
import matplotlib.pyplot as plt
from matplotlib.colors import hsv_to_rgb
import pyvista as pv
from collections import defaultdict
from PIL import Image

# Загрузите вашу 3D модель (замените 'your_model.obj' на путь к вашей модели)
mesh = trimesh.load('/.../model.obj')

faces_by_uv_vertex = defaultdict(lambda: [])
for fi0, f0 in enumerate(mesh.faces):
    for vi in f0:
        faces_by_uv_vertex[vi].append(fi0)

# Найти пары соседних треугольников (по UV-вершинам)
pairs = set()
for fi0, f0 in enumerate(mesh.faces):
    for vi in f0:
        for fi1 in faces_by_uv_vertex[vi]:
            pairs.add((min(fi0, fi1), max(fi0, fi1)))

# Создаем граф соседних треугольников
rows, cols = zip(*pairs)
data = np.ones(len(pairs))
num_faces = len(mesh.faces)
adj_matrix = coo_matrix((data, (rows, cols)), shape=(num_faces, num_faces))

# Находим компоненты связности в графе
num_components, labels = connected_components(csgraph=adj_matrix, directed=False)

# labels - это и есть chart_ids
chart_ids = labels

# Генерируем случайные цвета для каждого чарта
unique_charts = np.unique(chart_ids)
colors = np.array([hsv_to_rgb(((i * 239.17) % 1.0, 1, 1)) for i in range(len(unique_charts))])

# Присваиваем цвета треугольникам в зависимости от их чарта
face_colors = colors[chart_ids]

# Создаем PolyData для визуализации в PyVista
points = mesh.vertices
faces = np.hstack([np.full((mesh.faces.shape[0], 1), 3), mesh.faces]).flatten()  # PyVista требует формат [3, i0, i1, i2, ...]

# Создаем объект PolyData
polydata = pv.PolyData(points, faces)

# Добавляем цвета к PolyData
polydata["face_colors"] = (face_colors * 255).astype(np.uint8)

# Визуализация
plotter = pv.Plotter()
plotter.add_mesh(polydata, scalars="face_colors", rgb=True)
# plotter.show()

# Размер текстурного атласа (предполагаем квадратный атлас)
atlas_size = 2048
atlas_image = Image.new('RGB', (atlas_size, atlas_size), color=(255, 255, 255))
draw = ImageDraw.Draw(atlas_image)

# Получаем UV координаты для каждого треугольника
uvs = mesh.visual.uv

# Преобразование UV координат к размеру текстурного атласа
uvs_scaled = uvs * atlas_size

# Присваиваем цвета чартам на текстурном атласе
for face_id, chart_id in enumerate(chart_ids):
    color = tuple((colors[chart_id] * 255).astype(int))
    face = mesh.faces[face_id]
    polygon = [(uvs_scaled[vertex_id][0], atlas_size - uvs_scaled[vertex_id][1]) for vertex_id in face]
    draw.polygon(polygon, fill=color)

# Сохраняем
#atlas_image.show() # показываем текстурный атлас
atlas_image.save('texture_atlas.png')

plotter.show()