src/mesh/uvmapper.py

import numpy as np
import trimesh

DEFAULT_UV_FALLBACK = (0.5, 0.5)


def apply_atlas_to_mesh(mesh, atlas_texture, color_uv_mapping, palette_mapper, face_colors=None):
    from ..extraction.reader import get_face_colors

    if face_colors is None:
        original_face_colors = get_face_colors(mesh)
    else:
        original_face_colors = face_colors
    palette_matched_colors = palette_mapper.map_colors(original_face_colors)

    uv_coordinates_per_face = generate_face_uv_coordinates(
        palette_matched_colors, color_uv_mapping, len(mesh.faces)
    )

    duplicated_mesh = create_mesh_with_per_face_vertices(mesh, uv_coordinates_per_face)
    apply_texture_to_mesh(duplicated_mesh, atlas_texture, uv_coordinates_per_face)

    return duplicated_mesh


def generate_face_uv_coordinates(face_colors, color_to_uv_map, total_faces):
    uv_coordinates_array = np.zeros((total_faces, 3, 2), dtype=np.float32)
    faces_without_mapping = 0

    for face_index, face_color in enumerate(face_colors):
        rgb_color_tuple = convert_color_to_tuple(face_color)
        uv_position = color_to_uv_map.get(rgb_color_tuple, DEFAULT_UV_FALLBACK)

        if (
            uv_position == DEFAULT_UV_FALLBACK
            and rgb_color_tuple not in color_to_uv_map
        ):
            faces_without_mapping += 1

        assign_uniform_uv_to_face(uv_coordinates_array, face_index, uv_position)

    if faces_without_mapping > 0:
        print(f"Warning: {faces_without_mapping} faces had unmapped colors")

    return uv_coordinates_array


def convert_color_to_tuple(color):
    return tuple(int(channel) for channel in color)


def assign_uniform_uv_to_face(uv_array, face_index, uv_coordinate):
    uv_array[face_index] = [uv_coordinate, uv_coordinate, uv_coordinate]


def create_mesh_with_per_face_vertices(original_mesh, face_uv_coordinates):
    face_vertex_positions = original_mesh.vertices[original_mesh.faces]
    flattened_vertices = face_vertex_positions.reshape(-1, 3)
    sequential_face_indices = np.arange(len(flattened_vertices)).reshape(-1, 3)

    duplicated_mesh = trimesh.Trimesh(
        vertices=flattened_vertices,
        faces=sequential_face_indices,
        vertex_normals=None,
        process=False,
    )

    return duplicated_mesh


def apply_texture_to_mesh(mesh, atlas_image, face_uv_array):
    flattened_uv_coordinates = face_uv_array.reshape(-1, 2)

    pbr_material = trimesh.visual.material.PBRMaterial(
        baseColorTexture=atlas_image,
        baseColorFactor=[1.0, 1.0, 1.0, 1.0],
        roughnessFactor=1.0,
        metallicFactor=0.0,
    )

    texture_visual = trimesh.visual.texture.TextureVisuals(
        uv=flattened_uv_coordinates, image=atlas_image, material=pbr_material
    )

    mesh.visual = texture_visual