allenai / Holodeck

CVPR 2024: Language Guided Generation of 3D Embodied AI Environments.
https://yueyang1996.github.io/holodeck
Apache License 2.0
304 stars 25 forks source link

Converting .pkl.gz to .glb files, incorrect material #34

Closed jean-yoo closed 2 months ago

jean-yoo commented 3 months ago

Hi,

I'm trying to convert the .pkl.gz object models into a glb file that can be visualized with textures in Blender. I used the script in #9 to convert it to a .obj file, and used the code in #14 to map the textures back using bpy. But when I import the model in Blender, the materials don't seem to have textures mapped correctly; the whole armchair appears to be green, including the pillow and chair legs, which are both white in the Unity rendering (see images below). Exported gltf vs. Unity render:

Screenshot 2024-04-23 at 11 42 38 PM Screenshot 2024-04-23 at 11 50 13 PM

(obj uid 0e8e009f398249b9bc13e8ff7078530a)

Could you please verify whether the script provided in #14 is correct, or whether there's an error somewhere in my code? Thanks.

def create_and_assign_textures(material, albedo_path, emission_path, normal_path):
    # Enable 'Use Nodes'
    material.use_nodes = True
    nodes = material.node_tree.nodes

    # Find or create Principled BSDF node
    bsdf = next((node for node in nodes if node.type == 'BSDF_PRINCIPLED'), None)
    if not bsdf:
        bsdf = nodes.new(type='ShaderNodeBsdfPrincipled')
        bsdf.location = (0, 0)

    # Function to create or get an image texture node
    def get_image_node(img_path, location):
        for node in nodes:
            if node.type == 'TEX_IMAGE' and node.image and node.image.filepath == img_path:
                return node
        new_node = nodes.new('ShaderNodeTexImage')
        new_node.image = bpy.data.images.load(img_path)
        new_node.location = location
        return new_node

    # Create or get Albedo Texture node
    albedo_tex = get_image_node(albedo_path, (-300, 100))
    material.node_tree.links.new(bsdf.inputs['Base Color'], albedo_tex.outputs['Color'])

    # Create or get Emission Texture node
    emission_tex = get_image_node(emission_path, (-300, -100))
    material.node_tree.links.new(bsdf.inputs['Emission'], emission_tex.outputs['Color'])

    # Create or get Normal Map node
    normal_map = next((node for node in nodes if node.type == 'NORMAL_MAP'), None)
    if not normal_map:
        normal_map = nodes.new('ShaderNodeNormalMap')
        normal_map.location = (-300, -300)
    normal_tex = get_image_node(normal_path, (-500, -300))
    material.node_tree.links.new(normal_map.inputs['Color'], normal_tex.outputs['Color'])
    material.node_tree.links.new(bsdf.inputs['Normal'], normal_map.outputs['Normal'])

# Path to the OBJ file
obj_path = output_file_path

# Import OBJ file
bpy.ops.import_scene.obj(filepath=obj_path)

# Get the imported object
obj = bpy.context.selected_objects[0]

# Create a new material
material = bpy.data.materials.new(name="MyMaterial")

path = # path to holodeck objaverse here
uid = # obj uid 

# Assign textures
create_and_assign_textures(material, albedo_path=f'{path}/{uid}/albedo.jpg', emission_path=f'{path}/{uid}/emission.jpg', normal_path=f'{path}/{uid}/normal.jpg')
obj.material_slots[0].material = material
# Export as glTF or GLB
export_path = # export path here 
bpy.ops.export_scene.gltf(filepath=export_path, use_selection = True, export_format='GLB', export_materials = 'EXPORT')
YueYANG1996 commented 3 months ago

@sunfanyunn

sunfanyunn commented 3 months ago

Not entirely sure what is happening here -- but you can try to export the .pkl.gz file into Blender directly with the following script to see if the texture still appears to be incorrect:

def load_pickled_3d_asset(file_path):
    import gzip
    import pickle
    # Open the compressed pickled file
    with gzip.open(file_path, 'rb') as f:
        # Load the pickled object
        loaded_object_data = pickle.load(f)

    # Create a new mesh object in Blender
    mesh = bpy.data.meshes.new(name='LoadedMesh')
    obj = bpy.data.objects.new('LoadedObject', mesh)

    # Link the object to the scene
    bpy.context.scene.collection.objects.link(obj)

    # Set the mesh data for the object
    obj.data = mesh

    # Update the mesh with the loaded data
    # print(loaded_object_data.keys())
    # print(loaded_object_data['triangles'])
    # triangles = [vertex_index for face in loaded_object_data['triangles'] for vertex_index in face]
    triangles = np.array(loaded_object_data['triangles']).reshape(-1,3)
    vertices = []

    for v in loaded_object_data['vertices']:
        vertices.append([v['x'],v['z'],v['y']])

    mesh.from_pydata(vertices, [], triangles)

    uvs = []
    for uv in loaded_object_data['uvs']:
        uvs.append([uv['x'],uv['y']]) 

    mesh.update()

    # Ensure UV coordinates are stored
    if not mesh.uv_layers:
        mesh.uv_layers.new(name="UVMap")

    uv_layer = mesh.uv_layers["UVMap"]
    for poly in mesh.polygons:
        for loop_index in poly.loop_indices:
            vertex_index = mesh.loops[loop_index].vertex_index
            uv_layer.data[loop_index].uv = uvs[vertex_index]

    material = bpy.data.materials.new(name="AlbedoMaterial")
    obj.data.materials.append(material)

    # Assign albedo color to the material
    material.use_nodes = True
    nodes = material.node_tree.nodes
    principled_bsdf = nodes.get("Principled BSDF")

    texture_node = nodes.new(type='ShaderNodeTexImage')

    image_path = f"{'/'.join(file_path.split('/')[:-1])}/albedo.jpg"  # Replace with your image file path

    image = bpy.data.images.load(image_path)

    # Assign the image to the texture node
    texture_node.image = image

    # Connect the texture node to the albedo color
    material.node_tree.links.new(
        texture_node.outputs["Color"],
        principled_bsdf.inputs["Base Color"]
    )

    # normal
    image_path = f"{'/'.join(file_path.split('/')[:-1])}/normal.jpg"
    img_normal = bpy.data.images.load(image_path)
    image_texture_node_normal = material.node_tree.nodes.new(type='ShaderNodeTexImage')
    image_texture_node_normal.image = img_normal    
    image_texture_node_normal.image.colorspace_settings.name = 'Non-Color'

    normal_map_node = material.node_tree.nodes.new(type='ShaderNodeNormalMap')

    material.node_tree.links.new(normal_map_node.outputs["Normal"], principled_bsdf.inputs["Normal"])
    material.node_tree.links.new(image_texture_node_normal.outputs["Color"], normal_map_node.inputs["Color"])

    # Assign the material to the object
    obj.data.materials[0] = material    

    # Update mesh to apply UV changes
    mesh.update()

    return obj
jean-yoo commented 2 months ago

Yep, resolved, thanks -- the helper function to create textures also needed to load in the UV-coordinate information from the pkl.gz files and connect it to the normal map node using bpy.