tran.py

from trimesh import load, Trimesh, Scene
from trimesh.visual import TextureVisuals
from trimesh.visual.material import SimpleMaterial, color_image
from math import sqrt
from scipy.spatial import cKDTree
from numpy import sin, cos, dot, pi, eye
from json import dump
from os.path import basename

WOOL_PALETTE = {
    "minecraft:white_wool": (234, 236, 237),
    "minecraft:orange_wool": (241, 118, 20),
    "minecraft:magenta_wool": (191, 75, 201),
    "minecraft:light_blue_wool": (58, 175, 217),
    "minecraft:yellow_wool": (249, 198, 39),
    "minecraft:lime_wool": (112, 185, 25),
    "minecraft:pink_wool": (237, 141, 172),
    "minecraft:gray_wool": (62, 68, 71),
    "minecraft:light_gray_wool": (142, 142, 134),
    "minecraft:cyan_wool": (21, 137, 145),
    "minecraft:purple_wool": (121, 42, 172),
    "minecraft:blue_wool": (53, 57, 157),
    "minecraft:brown_wool": (114, 71, 40),
    "minecraft:green_wool": (84, 109, 27),
    "minecraft:red_wool": (161, 39, 34),
    "minecraft:black_wool": (20, 21, 25),
}
CONCRETE_PALETTE = {
    "minecraft:white_concrete": (207, 213, 214),
    "minecraft:orange_concrete": (224, 97, 0),
    "minecraft:magenta_concrete": (170, 47, 156),
    "minecraft:light_blue_concrete": (36, 137, 199),
    "minecraft:yellow_concrete": (241, 175, 21),
    "minecraft:lime_concrete": (94, 168, 24),
    "minecraft:pink_concrete": (214, 101, 143),
    "minecraft:gray_concrete": (54, 57, 61),
    "minecraft:light_gray_concrete": (125, 125, 115),
    "minecraft:cyan_concrete": (21, 119, 136),
    "minecraft:purple_concrete": (100, 32, 156),
    "minecraft:blue_concrete": (44, 46, 143),
    "minecraft:brown_concrete": (96, 59, 31),
    "minecraft:green_concrete": (73, 91, 36),
    "minecraft:red_concrete": (142, 32, 32),
    "minecraft:black_concrete": (8, 10, 15),
}
TERRACOTTA_PALETTE = {
    "minecraft:white_terracotta": (209, 178, 161),
    "minecraft:orange_terracotta": (161, 83, 37),
    "minecraft:magenta_terracotta": (150, 88, 109),
    "minecraft:light_blue_terracotta": (113, 108, 137),
    "minecraft:yellow_terracotta": (186, 133, 35),
    "minecraft:lime_terracotta": (103, 117, 53),
    "minecraft:pink_terracotta": (160, 77, 78),
    "minecraft:gray_terracotta": (57, 42, 35),
    "minecraft:light_gray_terracotta": (135, 107, 98),
    "minecraft:cyan_terracotta": (86, 91, 91),
    "minecraft:purple_terracotta": (118, 70, 86),
    "minecraft:blue_terracotta": (74, 59, 91),
    "minecraft:brown_terracotta": (77, 51, 36),
    "minecraft:green_terracotta": (76, 82, 42),
    "minecraft:red_terracotta": (143, 61, 46),
    "minecraft:black_terracotta": (37, 23, 16),
}
GLASS_PALETTE = {
    "minecraft:white_stained_glass": (255, 255, 255),
    "minecraft:orange_stained_glass": (216, 127, 51),
    "minecraft:magenta_stained_glass": (178, 76, 216),
    "minecraft:light_blue_stained_glass": (102, 153, 216),
    "minecraft:yellow_stained_glass": (229, 229, 51),
    "minecraft:lime_stained_glass": (127, 204, 25),
    "minecraft:pink_stained_glass": (242, 127, 165),
    "minecraft:gray_stained_glass": (76, 76, 76),
    "minecraft:light_gray_stained_glass": (153, 153, 153),
    "minecraft:cyan_stained_glass": (76, 127, 153),
    "minecraft:purple_stained_glass": (127, 63, 178),
    "minecraft:blue_stained_glass": (51, 76, 178),
    "minecraft:brown_stained_glass": (102, 76, 51),
    "minecraft:green_stained_glass": (102, 127, 51),
    "minecraft:red_stained_glass": (153, 51, 51),
    "minecraft:black_stained_glass": (25, 25, 25),
}

PITCH = 1.0  # Voxel size, smaller is finer (note MC block size limit)
ROTATE_ANGLE = (0, 0, 0)  # Rotation angle (x, y, z), in degrees

def call_back_null(stage_index, stage_num, current_step, stage_steps):
    pass

# Find the closest block in the palette based on color
def find_closest_block_with_glass(color, palette, glass_palette=GLASS_PALETTE):
    r, g, b, a = color
    min_distance = float('inf')
    closest_block = None
    if a < 200:
        for block_name, (br, bg, bb) in glass_palette.items():
            distance = sqrt((r-br)**2 + (g-bg)**2 + (b-bb)**2)
            if distance < min_distance:
                min_distance = distance
                closest_block = block_name
    else:
        for block_name, (br, bg, bb) in palette.items():
            distance = sqrt((r-br)**2 + (g-bg)**2 + (b-bb)**2)
            if distance < min_distance:
                min_distance = distance
                closest_block = block_name
    return closest_block
def find_closest_block_without_glass(color, palette):
    r, g, b, a = color
    min_distance = float('inf')
    closest_block = None
    for block_name, (br, bg, bb) in palette.items():
        distance = sqrt((r-br)**2 + (g-bg)**2 + (b-bb)**2)
        if distance < min_distance:
            min_distance = distance
            closest_block = block_name
    return closest_block
def find_closest_block_only_glass(color, palette, glass_palette=GLASS_PALETTE):
    r, g, b, a = color
    min_distance = float('inf')
    closest_block = None
    for block_name, (br, bg, bb) in glass_palette.items():
        distance = sqrt((r-br)**2 + (g-bg)**2 + (b-bb)**2)
        if distance < min_distance:
            min_distance = distance
            closest_block = block_name
    return closest_block

# Get the color of a voxel
def get_voxel_color(point, colors, tree):
    _, index = tree.query(point)
    voxel_color = colors[index]
    return voxel_color

# Get the palette based on the selected blocks
def get_palette(wool=True, concrete=True, terracotta=True):
    palette = {}
    if wool:
        palette.update(WOOL_PALETTE)
    if concrete:
        palette.update(CONCRETE_PALETTE)
    if terracotta:
        palette.update(TERRACOTTA_PALETTE)
    return palette

# Get the block function based on the selected blocks
def get_block_function(wool=True, concrete=True, terracotta=True, glass=True):
    if wool or concrete or terracotta:
        if glass:
            return find_closest_block_with_glass
        else:
            return find_closest_block_without_glass
    else:
        return find_closest_block_only_glass

# Load the model
def load_model(obj_file):
    print(f"Loading model from {obj_file}...")
    mesh = load(obj_file)
    print(f'Type of mesh: {type(mesh)}')
    if type(mesh) == Scene:
        meshes = list(mesh.geometry.values())
    elif type(mesh) == Trimesh:
        meshes = [mesh]
    else:
        raise TypeError(f"Unsupported mesh type: {type(mesh)}")
    print(f"Model loaded, {len(meshes)} meshes in total")
    return meshes

# Voxelize the model
def voxelize_model(mesh, pitch=PITCH):
    print("Voxelizing model...")
    print(f'Type of visual: {type(mesh.visual)}')
    if type(mesh.visual) == TextureVisuals:
        # If the material is a SimpleMaterial and the image is None, set the image to a color image
        if type(mesh.visual.material) == SimpleMaterial and mesh.visual.material.image is None:
            mesh.visual.material.image = color_image(mesh.visual.material.main_color)
        mesh.visual = mesh.visual.to_color() # Convert to ColorVisuals
        #mesh.visual = trimesh.visual.ColorVisuals(mesh=mesh) # Convert to ColorVisuals
    colors = mesh.visual.vertex_colors.astype(int)  # Get vertex colors
    vertices = mesh.vertices
    tree = cKDTree(vertices)  # Build a KDTree for fast color lookup
    voxels = mesh.voxelized(pitch=pitch)
    points = voxels.points  # Voxel coordinates
    print(f"Voxelization completed, {len(points)} blocks in total")
    return points, colors, tree

# Insert blocks
def insert_blocks(block_data, points, colors, tree, palette, pitch=PITCH, call_back=call_back_null, stage_index=0, stage_num=1):
    for i, point in enumerate(points):
        # Calculate real world coordinates
        x, y, z = point / pitch
        x, y, z = map(round, (x, y, z))

        if (x == 0 and y == 0 and z == 0):
            continue

        # Get voxel color
        voxel_color = get_voxel_color(point, colors, tree)
        # Find closest block
        closest_block = find_closest_block(voxel_color, palette)

        if closest_block not in block_data:
            block_data[closest_block] = []
        block_data[closest_block].append([x, y, z])

        call_back(stage_index, stage_num, len(points)+i, len(points)*2)

    print(f"Successfully set {len(points)} blocks")
    return block_data

# Calculate rotation matrix
def calculate_rotation_matrix(rotate_angle=ROTATE_ANGLE):
    # Convert degrees to radians for rotation
    rx, ry, rz = map(lambda x: x * 3.1415926 / 180, rotate_angle)
    # Rotation matrix for x-axis
    rx_matrix = [
        [1, 0, 0],
        [0, cos(rx), -sin(rx)],
        [0, sin(rx), cos(rx)]
    ]
    # Rotation matrix for y-axis
    ry_matrix = [
        [cos(ry), 0, sin(ry)],
        [0, 1, 0],
        [-sin(ry), 0, cos(ry)]
    ]
    # Rotation matrix for z-axis
    rz_matrix = [
        [cos(rz), -sin(rz), 0],
        [sin(rz), cos(rz), 0],
        [0, 0, 1]
    ]
    # Combine rotation matrices
    return dot(rz_matrix, dot(ry_matrix, rx_matrix))

# Main function
def model_to_minecraft(obj_file, save_path, pitch=PITCH, rotate_angle = ROTATE_ANGLE, wool=True, concrete=True, terracotta=True, glass=True, call_back=call_back_null):
    global find_closest_block
    find_closest_block = get_block_function(wool, concrete, terracotta, glass)
    palette = get_palette(wool, concrete, terracotta)

    meshes = load_model(obj_file)
    call_back(len(meshes), len(meshes)*2, 0, 1)

    # Calculate rotation matrix
    if rotate_angle != (0, 0, 0):
        temp_matrix = calculate_rotation_matrix(rotate_angle)
        rotation_matrix = eye(4)
        rotation_matrix[:3, :3] = temp_matrix
        for mesh in meshes:
            mesh.apply_transform(rotation_matrix)

    # Find the minimum x, y, z coordinates of all meshes
    min_x = float('inf')
    min_y = float('inf')
    min_z = float('inf')
    for mesh in meshes:
        corner = mesh.bounds[0]
        min_x = min(min_x, corner[0])
        min_y = min(min_y, corner[1])
        min_z = min(min_z, corner[2])

    block_data = {}
    for i, mesh in enumerate(meshes):
        # Translate the mesh to the origin
        if min_x < 0 or min_y < 0 or min_z < 0:
            mesh.apply_translation((-min_x, -min_y, -min_z))
        points, colors, tree = voxelize_model(mesh, pitch)
        call_back(len(meshes)+i+1, len(meshes)*2, len(points)-1, len(points)*2)
        block_data = insert_blocks(block_data, points, colors, tree, palette, pitch, call_back, len(meshes)+i+1, len(meshes)*2+1)
    
    # Save block data to JSON file
    filename = save_path + '/' + basename(obj_file).split('.')[0] + '.json'
    with open(filename, 'w') as f:
        dump(block_data, f, indent=2)
    print("Block data saved to", filename)
    print('-'*20)