chore: import upstream snapshot with attribution

data_toolkit/blender_script/dump_mesh.py

@@ -0,0 +1,242 @@
+import argparse, sys, os, math, io
+from typing import *
+import bpy
+import bmesh
+from mathutils import Vector, Matrix
+import numpy as np
+import pickle
+
+
+"""=============== BLENDER ==============="""
+
+IMPORT_FUNCTIONS: Dict[str, Callable] = {
+    "obj": bpy.ops.import_scene.obj if bpy.app.version[0] < 4 else bpy.ops.wm.obj_import,
+    "glb": bpy.ops.import_scene.gltf,
+    "gltf": bpy.ops.import_scene.gltf,
+    "usd": bpy.ops.import_scene.usd,
+    "fbx": bpy.ops.import_scene.fbx,
+    "stl": bpy.ops.import_mesh.stl if bpy.app.version[0] < 4 else bpy.ops.wm.stl_import,
+    "usda": bpy.ops.import_scene.usda,
+    "dae": bpy.ops.wm.collada_import,
+    "ply": bpy.ops.import_mesh.ply if bpy.app.version[0] < 4 else bpy.ops.wm.ply_import,
+    "abc": bpy.ops.wm.alembic_import,
+    "blend": bpy.ops.wm.append,
+}
+
+
+def init_scene() -> None:
+    """Resets the scene to a clean state.
+
+    Returns:
+        None
+    """
+    # delete everything
+    for obj in bpy.data.objects:
+        bpy.data.objects.remove(obj, do_unlink=True)
+
+    # delete all the materials
+    for material in bpy.data.materials:
+        bpy.data.materials.remove(material, do_unlink=True)
+
+    # delete all the textures
+    for texture in bpy.data.textures:
+        bpy.data.textures.remove(texture, do_unlink=True)
+
+    # delete all the images
+    for image in bpy.data.images:
+        bpy.data.images.remove(image, do_unlink=True)
+
+
+def load_object(object_path: str) -> None:
+    """Loads a model with a supported file extension into the scene.
+
+    Args:
+        object_path (str): Path to the model file.
+
+    Raises:
+        ValueError: If the file extension is not supported.
+
+    Returns:
+        None
+    """
+    file_extension = object_path.split(".")[-1].lower()
+    if file_extension is None:
+        raise ValueError(f"Unsupported file type: {object_path}")
+
+    if file_extension == "usdz":
+        # install usdz io package
+        dirname = os.path.dirname(os.path.realpath(__file__))
+        usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
+        bpy.ops.preferences.addon_install(filepath=usdz_package)
+        # enable it
+        addon_name = "io_scene_usdz"
+        bpy.ops.preferences.addon_enable(module=addon_name)
+        # import the usdz
+        from io_scene_usdz.import_usdz import import_usdz
+
+        import_usdz(context, filepath=object_path, materials=True, animations=True)
+        return None
+
+    # load from existing import functions
+    import_function = IMPORT_FUNCTIONS[file_extension]
+
+    print(f"Loading object from {object_path}")
+    if file_extension == "blend":
+        import_function(directory=object_path, link=False)
+    elif file_extension in {"glb", "gltf"}:
+        import_function(filepath=object_path, merge_vertices=True, import_shading='NORMALS', bone_heuristic='TEMPERANCE')
+    else:
+        import_function(filepath=object_path)
+        
+        
+def delete_invisible_objects() -> None:
+    """Deletes all invisible objects in the scene.
+
+    Returns:
+        None
+    """
+    to_remove = []
+    for obj in bpy.context.scene.objects:
+        if obj.hide_viewport or obj.hide_render:
+            obj.hide_viewport = False
+            obj.hide_render = False
+            obj.hide_select = False
+            to_remove.append(obj)
+    for obj in to_remove:
+        bpy.data.objects.remove(obj, do_unlink=True)
+
+    # Delete invisible collections
+    invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
+    for col in invisible_collections:
+        bpy.data.collections.remove(col)
+      
+
+def scene_bbox() -> Tuple[Vector, Vector]:
+    """Returns the bounding box of the scene.
+
+    Taken from Shap-E rendering script
+    (https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)
+
+    Returns:
+        Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
+    """
+    bbox_min = (math.inf,) * 3
+    bbox_max = (-math.inf,) * 3
+    found = False
+    scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
+    for obj in scene_meshes:
+        found = True
+        for coord in obj.bound_box:
+            coord = Vector(coord)
+            coord = obj.matrix_world @ coord
+            bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
+            bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
+    if not found:
+        raise RuntimeError("no objects in scene to compute bounding box for")
+    return Vector(bbox_min), Vector(bbox_max)
+
+
+def normalize_scene() -> Tuple[float, Vector]:
+    """Normalizes the scene by scaling and translating it to fit in a unit cube centered
+    at the origin.
+
+    Mostly taken from the Point-E / Shap-E rendering script
+    (https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
+    but fix for multiple root objects: (see bug report here:
+    https://github.com/openai/shap-e/pull/60).
+
+    Returns:
+        Tuple[float, Vector]: The scale factor and the offset applied to the scene.
+    """
+    scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
+    if len(scene_root_objects) > 1:
+        # create an empty object to be used as a parent for all root objects
+        scene = bpy.data.objects.new("ParentEmpty", None)
+        bpy.context.scene.collection.objects.link(scene)
+
+        # parent all root objects to the empty object
+        for obj in scene_root_objects:
+            obj.parent = scene
+    else:
+        scene = scene_root_objects[0]
+
+    bbox_min, bbox_max = scene_bbox()
+    scale = 1 / max(bbox_max - bbox_min)
+    scene.scale = scene.scale * scale
+
+    # Apply scale to matrix_world.
+    bpy.context.view_layer.update()
+    bbox_min, bbox_max = scene_bbox()
+    offset = -(bbox_min + bbox_max) / 2
+    scene.matrix_world.translation += offset
+    
+    return scale, offset
+
+
+def main(arg):    
+    # Initialize context
+    if arg.object.endswith(".blend"):
+        delete_invisible_objects()
+    else:
+        init_scene()
+        load_object(arg.object)
+    print('[INFO] Scene initialized.')
+    
+    # Normalize scene
+    scale, offset = normalize_scene()
+    print('[INFO] Scene normalized.')
+    
+    # Start dumping
+    depsgraph = bpy.context.evaluated_depsgraph_get()
+    scene = bpy.context.scene
+    output = {
+        'objects': [],
+    }
+
+    # Dumping meshes
+    for obj in scene.objects:
+        if obj.type != 'MESH':
+            continue
+        
+        pack = {
+            "vertices": None,
+            "faces": None,
+        }
+        
+        eval_obj = obj.evaluated_get(depsgraph)
+        eval_mesh = eval_obj.to_mesh()
+        
+        bm = bmesh.new()
+        bm.from_mesh(eval_mesh)
+        bm.transform(obj.matrix_world)
+        bmesh.ops.triangulate(bm, faces=bm.faces)
+        bm.to_mesh(eval_mesh)
+        bm.free()
+                
+        pack["vertices"] = np.array([
+            v.co[:] for v in eval_mesh.vertices
+        ], dtype=np.float32)   # (N, 3)
+        
+        pack["faces"] = np.array([
+            [eval_mesh.loops[i].vertex_index for i in poly.loop_indices]
+            for poly in eval_mesh.polygons
+        ], dtype=np.int32)   # (F, 3)
+
+        output['objects'].append(pack)
+
+    # Save output
+    os.makedirs(os.path.dirname(arg.output_path), exist_ok=True)
+    with open(arg.output_path, 'wb') as f:
+        pickle.dump(output, f)
+    print('[INFO] Output saved to {}.'.format(arg.output_path))
+
+        
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
+    parser.add_argument('--object', type=str, help='Path to the 3D model file to be rendered.')
+    parser.add_argument('--output_path', type=str, default='/tmp', help='The path the output will be dumped to.')
+    argv = sys.argv[sys.argv.index("--") + 1:]
+    args = parser.parse_args(argv)
+
+    main(args)
+    

data_toolkit/blender_script/dump_pbr.py

@@ -0,0 +1,485 @@
+import argparse, sys, os, math, io
+from typing import *
+import bpy
+import bmesh
+from mathutils import Vector, Matrix
+import numpy as np
+from PIL import Image
+import pickle
+
+
+"""=============== BLENDER ==============="""
+
+IMPORT_FUNCTIONS: Dict[str, Callable] = {
+    "obj": bpy.ops.import_scene.obj if bpy.app.version[0] < 4 else bpy.ops.wm.obj_import,
+    "glb": bpy.ops.import_scene.gltf,
+    "gltf": bpy.ops.import_scene.gltf,
+    "usd": bpy.ops.import_scene.usd,
+    "fbx": bpy.ops.import_scene.fbx,
+    "stl": bpy.ops.import_mesh.stl if bpy.app.version[0] < 4 else bpy.ops.wm.stl_import,
+    "usda": bpy.ops.import_scene.usda,
+    "dae": bpy.ops.wm.collada_import,
+    "ply": bpy.ops.import_mesh.ply if bpy.app.version[0] < 4 else bpy.ops.wm.ply_import,
+    "abc": bpy.ops.wm.alembic_import,
+    "blend": bpy.ops.wm.append,
+}
+
+
+def init_scene() -> None:
+    """Resets the scene to a clean state.
+
+    Returns:
+        None
+    """
+    # delete everything
+    for obj in bpy.data.objects:
+        bpy.data.objects.remove(obj, do_unlink=True)
+
+    # delete all the materials
+    for material in bpy.data.materials:
+        bpy.data.materials.remove(material, do_unlink=True)
+
+    # delete all the textures
+    for texture in bpy.data.textures:
+        bpy.data.textures.remove(texture, do_unlink=True)
+
+    # delete all the images
+    for image in bpy.data.images:
+        bpy.data.images.remove(image, do_unlink=True)
+
+
+def load_object(object_path: str) -> None:
+    """Loads a model with a supported file extension into the scene.
+
+    Args:
+        object_path (str): Path to the model file.
+
+    Raises:
+        ValueError: If the file extension is not supported.
+
+    Returns:
+        None
+    """
+    file_extension = object_path.split(".")[-1].lower()
+    if file_extension is None:
+        raise ValueError(f"Unsupported file type: {object_path}")
+
+    if file_extension == "usdz":
+        # install usdz io package
+        dirname = os.path.dirname(os.path.realpath(__file__))
+        usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
+        bpy.ops.preferences.addon_install(filepath=usdz_package)
+        # enable it
+        addon_name = "io_scene_usdz"
+        bpy.ops.preferences.addon_enable(module=addon_name)
+        # import the usdz
+        from io_scene_usdz.import_usdz import import_usdz
+
+        import_usdz(context, filepath=object_path, materials=True, animations=True)
+        return None
+
+    # load from existing import functions
+    import_function = IMPORT_FUNCTIONS[file_extension]
+
+    print(f"Loading object from {object_path}")
+    if file_extension == "blend":
+        import_function(directory=object_path, link=False)
+    elif file_extension in {"glb", "gltf"}:
+        import_function(filepath=object_path, merge_vertices=True, import_shading='NORMALS', bone_heuristic='TEMPERANCE')
+    else:
+        import_function(filepath=object_path)
+        
+        
+def delete_invisible_objects() -> None:
+    """Deletes all invisible objects in the scene.
+
+    Returns:
+        None
+    """
+    to_remove = []
+    for obj in bpy.context.scene.objects:
+        if obj.hide_viewport or obj.hide_render:
+            obj.hide_viewport = False
+            obj.hide_render = False
+            obj.hide_select = False
+            to_remove.append(obj)
+    for obj in to_remove:
+        bpy.data.objects.remove(obj, do_unlink=True)
+
+    # Delete invisible collections
+    invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
+    for col in invisible_collections:
+        bpy.data.collections.remove(col)
+      
+
+def scene_bbox() -> Tuple[Vector, Vector]:
+    """Returns the bounding box of the scene.
+
+    Taken from Shap-E rendering script
+    (https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)
+
+    Returns:
+        Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
+    """
+    bbox_min = (math.inf,) * 3
+    bbox_max = (-math.inf,) * 3
+    found = False
+    scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
+    for obj in scene_meshes:
+        found = True
+        for coord in obj.bound_box:
+            coord = Vector(coord)
+            coord = obj.matrix_world @ coord
+            bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
+            bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
+    if not found:
+        raise RuntimeError("no objects in scene to compute bounding box for")
+    return Vector(bbox_min), Vector(bbox_max)
+
+
+def normalize_scene() -> Tuple[float, Vector]:
+    """Normalizes the scene by scaling and translating it to fit in a unit cube centered
+    at the origin.
+
+    Mostly taken from the Point-E / Shap-E rendering script
+    (https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
+    but fix for multiple root objects: (see bug report here:
+    https://github.com/openai/shap-e/pull/60).
+
+    Returns:
+        Tuple[float, Vector]: The scale factor and the offset applied to the scene.
+    """
+    scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
+    if len(scene_root_objects) > 1:
+        # create an empty object to be used as a parent for all root objects
+        scene = bpy.data.objects.new("ParentEmpty", None)
+        bpy.context.scene.collection.objects.link(scene)
+
+        # parent all root objects to the empty object
+        for obj in scene_root_objects:
+            obj.parent = scene
+    else:
+        scene = scene_root_objects[0]
+
+    bbox_min, bbox_max = scene_bbox()
+    scale = 1 / max(bbox_max - bbox_min)
+    scene.scale = scene.scale * scale
+
+    # Apply scale to matrix_world.
+    bpy.context.view_layer.update()
+    bbox_min, bbox_max = scene_bbox()
+    offset = -(bbox_min + bbox_max) / 2
+    scene.matrix_world.translation += offset
+    
+    return scale, offset
+
+
+# =============== NODE TREE PARSING ===============
+
+def extract_image(tex_node, channels):
+        image = tex_node.image
+        pixels = np.array(image.pixels[:])
+        data = pixels.reshape(image.size[1], image.size[0], -1)
+        data = data[..., channels]
+
+        if data.dtype != np.uint8:
+            data = np.clip(data, 0.0, 1.0)
+            data = (data * 255).astype(np.uint8)
+
+        if len(data.shape) == 2:  # Single channel
+            pil_image = Image.fromarray(data, mode='L')
+        elif data.shape[2] == 3:
+            pil_image = Image.fromarray(data, mode='RGB')
+        elif data.shape[2] == 4:
+            pil_image = Image.fromarray(data, mode='RGBA')
+        else:
+            raise ValueError("Unsupported channel shape for image")
+
+        buffer = io.BytesIO()
+        pil_image.save(buffer, format='PNG')
+        png_bytes = buffer.getvalue()
+
+        return {
+            'image': png_bytes,
+            'interpolation': tex_node.interpolation,
+            'extension': tex_node.extension,
+        }
+
+
+def try_extract_image(link, expected_channel='RGB'):
+    """
+    Tries to extract an image from a texture node link.
+    Supported sub tree modes:
+      - RGB:
+        TEX_IMAGE ->
+      - R, G, B:
+        TEX_IMAGE -> SEPARATE_COLOR ->
+      - A:
+        TEX_IMAGE ->
+    """
+    assert expected_channel in ['RGB', 'R', 'G', 'B', 'A'], "Unsupported channel"
+
+    if expected_channel == 'RGB':
+        assert link.from_node.type == 'TEX_IMAGE', "Material is not supported"
+        assert link.from_socket.name == 'Color', "Material is not supported"
+        tex_node = link.from_node
+        return extract_image(tex_node, [0, 1, 2])
+
+    if expected_channel in ['R', 'G', 'B']:
+        socket_name = {
+            'R': 'Red',
+            'G': 'Green',
+            'B': 'Blue',
+        }[expected_channel]
+        assert link.from_node.type == 'SEPARATE_COLOR' and link.from_node.mode == 'RGB', \
+            f"Material is not supported, {link.from_node.type}, {link.from_node.mode}"
+        assert link.from_socket.name == socket_name, "Material is not supported"
+        sep_node = link.from_node
+        assert sep_node.inputs[0].is_linked and sep_node.inputs[0].links[0].from_node.type == 'TEX_IMAGE', \
+            "Material is not supported"
+        assert sep_node.inputs[0].links[0].from_socket.name == 'Color', "Material is not supported"
+        tex_node = sep_node.inputs[0].links[0].from_node
+        channel_index = {
+            'R': 0,
+            'G': 1,
+            'B': 2,
+        }[expected_channel]
+        return extract_image(tex_node, channel_index)
+
+    if expected_channel == 'A':
+        assert link.from_node.type == 'TEX_IMAGE', "Material is not supported"
+        assert link.from_socket.name == 'Alpha', "Material is not supported"
+        tex_node = link.from_node
+        return extract_image(tex_node, 3)
+
+
+def try_extract_factor(link, mode='color'):
+    """
+    Tries to extract a factor from a math node link.
+    Supported sub tree modes:
+      - color:
+       ANY -> MIX(MULTIPLY) ->
+      - scalar:
+       ANY -> MATH(MULTIPLY) ->
+    """
+    assert mode in ['color','scalar'], "Unsupported mode"
+
+    if mode == 'color':
+        if link.from_node.type == 'MIX':
+            mix_node = link.from_node
+            assert mix_node.data_type == 'RGBA' and mix_node.blend_type == 'MULTIPLY', f"Material is not supported, {mix_node.data_type}, {mix_node.blend_type}"
+            assert not mix_node.inputs['Factor'].is_linked and mix_node.inputs['Factor'].default_value == 1.0, \
+                "Material is not supported"
+            if mix_node.inputs['A'].is_linked:
+                assert not mix_node.inputs['B'].is_linked, "Material is not supported"
+                return (list(mix_node.inputs['B'].default_value)[:3], mix_node.inputs['A'].links[0])
+            else:
+                assert not mix_node.inputs['A'].is_linked, "Material is not supported"
+                assert mix_node.inputs['B'].is_linked, "Material is not supported"
+                return (list(mix_node.inputs['A'].default_value)[:3], mix_node.inputs['B'].links[0])
+        return ([1.0, 1.0, 1.0], link)
+
+    if mode =='scalar':
+        if link.from_node.type == 'MATH':
+            math_node = link.from_node
+            assert math_node.operation == 'MULTIPLY', "Material is not supported"
+            assert math_node.inputs[0].is_linked, "Material is not supported"
+            assert not math_node.inputs[1].is_linked, "Material is not supported"
+            return (math_node.inputs[1].default_value, math_node.inputs[0].links[0])
+        return (1.0, link)
+
+
+def try_extract_image_with_factor(link, expected_channel='RGB'):
+    """
+    Tries to extract an image and a factor from a texture node link.
+    """
+    factor, link = try_extract_factor(link, 'color' if expected_channel in ['RGB'] else 'scalar')
+    image = try_extract_image(link, expected_channel)
+    return (factor, image)
+
+
+def main(arg):    
+    # Initialize context
+    if arg.object.endswith(".blend"):
+        delete_invisible_objects()
+    else:
+        init_scene()
+        load_object(arg.object)
+    print('[INFO] Scene initialized.')
+    
+    # Normalize scene
+    scale, offset = normalize_scene()
+    print('[INFO] Scene normalized.')
+    
+    # Start dumping
+    depsgraph = bpy.context.evaluated_depsgraph_get()
+    scene = bpy.context.scene
+    output = {
+        'materials': [],
+        'objects': [],
+    }
+
+    # Dumping materials
+    for mat in bpy.data.materials:
+        assert mat.use_nodes == True, "Material is not supported"
+
+        pack = {
+            "baseColorFactor": [1.0, 1.0, 1.0],
+            "alphaFactor": 1.0,
+            "metallicFactor": 1.0,
+            "roughnessFactor": 1.0,
+            "alphaMode": "OPAQUE",
+            "alphaCutoff": 0.5,
+            "baseColorTexture": None,
+            "alphaTexture": None,
+            "metallicTexture": None,
+            "roughnessTexture": None,
+        }
+
+        try:
+            principled_node = mat.node_tree.nodes.get('Principled BSDF')
+            assert principled_node is not None, "Material is not supported"
+
+            # Handle base color
+            if not principled_node.inputs['Base Color'].is_linked:
+                pack["baseColorFactor"] = list(principled_node.inputs['Base Color'].default_value)
+            else:
+                link = principled_node.inputs['Base Color'].links[0]
+                if link.from_node.type == 'RGB':
+                    pack["baseColorFactor"] = list(link.from_node.outputs[0].default_value)
+                else:
+                    factor, image = try_extract_image_with_factor(link, 'RGB')
+                    pack["baseColorFactor"] = factor
+                    pack["baseColorTexture"] = image
+
+            # Handle alpha
+            if not principled_node.inputs['Alpha'].is_linked:
+                pack["alphaFactor"] = principled_node.inputs['Alpha'].default_value
+                if pack["alphaFactor"] < 1.0:
+                    pack["alphaMode"] = "BLEND"
+            else:
+                link = principled_node.inputs['Alpha'].links[0]
+                node = link.from_node
+                if node.type == 'VALUE':
+                    pack["alphaFactor"] = node.outputs[0].default_value
+                    if pack["alphaFactor"] < 1.0:
+                        pack["alphaMode"] = "BLEND"
+                else:
+                    pack["alphaMode"] = "BLEND"
+                    if node.type == 'MATH':
+                        if node.operation == 'ROUND':
+                            assert node.inputs[0].is_linked, "Material is not supported"
+                            pack["alphaMode"] = "MASK"
+                            link = node.inputs[0].links[0]
+                        elif node.operation == 'SUBTRACT':
+                            assert node.inputs[0].default_value == 1.0 and \
+                                node.inputs[1].is_linked and \
+                                node.inputs[1].links[0].from_node.type == 'MATH' and \
+                                node.inputs[1].links[0].from_node.operation == 'LESS_THAN', \
+                                "Material is not supported"
+                            assert node.inputs[1].links[0].from_node.inputs[0].is_linked, "Material is not supported"
+                            pack["alphaMode"] = "MASK"
+                            pack["alphaCutoff"] = node.inputs[1].links[0].from_node.inputs[1].default_value
+                            link = node.inputs[1].links[0].from_node.inputs[0].links[0]
+                    factor, image = try_extract_image_with_factor(link, 'A')
+                    pack["alphaFactor"] = factor
+                    pack["alphaTexture"] = image
+
+            # Handle metallic
+            if not principled_node.inputs['Metallic'].is_linked:
+                pack["metallicFactor"] = principled_node.inputs['Metallic'].default_value
+            else:
+                link = principled_node.inputs['Metallic'].links[0]
+                node = link.from_node
+                if node.type == 'VALUE':
+                    pack["metallicFactor"] = node.outputs[0].default_value
+                else:
+                    factor, image = try_extract_image_with_factor(link, 'B')
+                    pack["metallicFactor"] = factor
+                    pack["metallicTexture"] = image
+
+            # Handle roughness
+            if not principled_node.inputs['Roughness'].is_linked:
+                pack["roughnessFactor"] = principled_node.inputs['Roughness'].default_value
+            else:
+                link = principled_node.inputs['Roughness'].links[0]
+                node = link.from_node
+                if node.type == 'VALUE':
+                    pack["roughnessFactor"] = node.outputs[0].default_value
+                else:
+                    factor, image = try_extract_image_with_factor(link, 'G')
+                    pack["roughnessFactor"] = factor
+                    pack["roughnessTexture"] = image
+
+            output['materials'].append(pack)
+        except:
+            with open(arg.output_path + '_error.txt', 'w') as f:
+                f.write(str([[n.name] for n in mat.node_tree.nodes]))
+            raise RuntimeError("Material is not supported")
+
+    # Dumping meshes
+    for obj in scene.objects:
+        if obj.type != 'MESH':
+            continue
+        
+        pack = {
+            "vertices": None,
+            "faces": None,
+            "uvs": None,
+            "matIDs": None,
+        }
+        
+        eval_obj = obj.evaluated_get(depsgraph)
+        eval_mesh = eval_obj.to_mesh()
+        
+        bm = bmesh.new()
+        bm.from_mesh(eval_mesh)
+        bm.transform(obj.matrix_world)
+        bmesh.ops.triangulate(bm, faces=bm.faces)
+        bm.to_mesh(eval_mesh)
+        bm.free()
+                
+        pack["vertices"] = np.array([
+            v.co[:] for v in eval_mesh.vertices
+        ], dtype=np.float32)   # (N, 3)
+        
+        pack["faces"] = np.array([
+            [eval_mesh.loops[i].vertex_index for i in poly.loop_indices]
+            for poly in eval_mesh.polygons
+        ], dtype=np.int32)   # (F, 3)
+        
+        pack["normals"] = np.array([
+            [eval_mesh.loops[i].normal for i in poly.loop_indices]
+            for poly in eval_mesh.polygons
+        ], dtype=np.float32)  # (F, 3, 3)
+        
+        if eval_mesh.uv_layers.active is not None:
+            pack["uvs"] = np.array([
+                [eval_mesh.uv_layers.active.data[i].uv for i in poly.loop_indices]
+                for poly in eval_mesh.polygons
+            ], dtype=np.float32)  # (F, 3, 2)
+
+        pack["mat_ids"] = np.array([
+            bpy.data.materials.find(obj.material_slots[poly.material_index].name)
+            if len(obj.material_slots) > 0 and obj.material_slots[poly.material_index].material is not None else -1
+            for poly in eval_mesh.polygons
+        ], dtype=np.int32)
+
+        output['objects'].append(pack)
+
+    # Save output
+    os.makedirs(os.path.dirname(arg.output_path), exist_ok=True)
+    with open(arg.output_path, 'wb') as f:
+        pickle.dump(output, f)
+    print('[INFO] Output saved to {}.'.format(arg.output_path))
+
+        
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
+    parser.add_argument('--object', type=str, help='Path to the 3D model file to be rendered.')
+    parser.add_argument('--output_path', type=str, default='/tmp', help='The path the output will be dumped to.')
+    argv = sys.argv[sys.argv.index("--") + 1:]
+    args = parser.parse_args(argv)
+
+    main(args)
+    

data_toolkit/blender_script/install_pillow.py

@@ -0,0 +1,6 @@
+import subprocess
+import sys
+import ensurepip
+
+ensurepip.bootstrap()
+subprocess.check_call([sys.executable, "-m", "pip", "install", "Pillow"])

data_toolkit/blender_script/render_cond.py

@@ -0,0 +1,529 @@
+import argparse, sys, os, math, re, glob
+from typing import *
+import bpy
+from mathutils import Vector, Matrix
+import numpy as np
+import json
+import glob
+from PIL import Image
+
+
+"""=============== BLENDER ==============="""
+
+IMPORT_FUNCTIONS: Dict[str, Callable] = {
+    "obj": bpy.ops.import_scene.obj,
+    "glb": bpy.ops.import_scene.gltf,
+    "gltf": bpy.ops.import_scene.gltf,
+    "usd": bpy.ops.import_scene.usd,
+    "fbx": bpy.ops.import_scene.fbx,
+    "stl": bpy.ops.import_mesh.stl,
+    "usda": bpy.ops.import_scene.usda,
+    "dae": bpy.ops.wm.collada_import,
+    "ply": bpy.ops.import_mesh.ply,
+    "abc": bpy.ops.wm.alembic_import,
+    "blend": bpy.ops.wm.append,
+}
+
+EXT = {
+    'PNG': 'png',
+    'JPEG': 'jpg',
+    'OPEN_EXR': 'exr',
+    'TIFF': 'tiff',
+    'BMP': 'bmp',
+    'HDR': 'hdr',
+    'TARGA': 'tga'
+}
+
+
+def init_render(engine='CYCLES', resolution=512):
+    bpy.context.scene.render.engine = engine
+    bpy.context.scene.render.resolution_x = resolution
+    bpy.context.scene.render.resolution_y = resolution
+    bpy.context.scene.render.resolution_percentage = 100
+    bpy.context.scene.render.image_settings.file_format = 'PNG'
+    bpy.context.scene.render.image_settings.color_mode = 'RGBA'
+    bpy.context.scene.render.film_transparent = True
+    
+    bpy.context.scene.cycles.device = 'GPU'
+    bpy.context.scene.cycles.samples = 32
+    bpy.context.scene.cycles.filter_type = 'BOX'
+    bpy.context.scene.cycles.filter_width = 1
+    bpy.context.scene.cycles.diffuse_bounces = 1
+    bpy.context.scene.cycles.glossy_bounces = 1
+    bpy.context.scene.cycles.transparent_max_bounces = 3
+    bpy.context.scene.cycles.transmission_bounces = 3
+    bpy.context.scene.cycles.use_denoising = True
+        
+    bpy.context.preferences.addons['cycles'].preferences.get_devices()
+    bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
+    
+
+def init_scene() -> None:
+    """Resets the scene to a clean state.
+
+    Returns:
+        None
+    """
+    # delete everything
+    for obj in bpy.data.objects:
+        bpy.data.objects.remove(obj, do_unlink=True)
+
+    # delete all the materials
+    for material in bpy.data.materials:
+        bpy.data.materials.remove(material, do_unlink=True)
+
+    # delete all the textures
+    for texture in bpy.data.textures:
+        bpy.data.textures.remove(texture, do_unlink=True)
+
+    # delete all the images
+    for image in bpy.data.images:
+        bpy.data.images.remove(image, do_unlink=True)
+        
+
+def init_camera():
+    cam = bpy.data.objects.new('Camera', bpy.data.cameras.new('Camera'))
+    bpy.context.collection.objects.link(cam)
+    bpy.context.scene.camera = cam
+    cam.data.sensor_height = cam.data.sensor_width = 32
+    cam_constraint = cam.constraints.new(type='TRACK_TO')
+    cam_constraint.track_axis = 'TRACK_NEGATIVE_Z'
+    cam_constraint.up_axis = 'UP_Y'
+    cam_empty = bpy.data.objects.new("Empty", None)
+    cam_empty.location = (0, 0, 0)
+    bpy.context.scene.collection.objects.link(cam_empty)
+    cam_constraint.target = cam_empty
+    return cam
+
+
+def init_uniform_lighting():
+    # Clear existing lights
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.ops.object.select_by_type(type="LIGHT")
+    bpy.ops.object.delete()
+    
+    # Create environment light
+    if bpy.context.scene.world is None:
+        world = bpy.data.worlds.new("World")
+        bpy.context.scene.world = world
+    else:
+        world = bpy.context.scene.world
+        
+    # Enabling nodes
+    world.use_nodes = True
+    node_tree = world.node_tree
+    nodes = node_tree.nodes
+    links = node_tree.links
+    
+    # Remove default nodes
+    for node in nodes:
+        nodes.remove(node)
+
+    # Create background node
+    bg_node = nodes.new(type="ShaderNodeBackground")
+    bg_node.inputs["Color"].default_value = (1.0, 1.0, 1.0, 1.0)
+    bg_node.inputs["Strength"].default_value = 1.0
+    output_node = nodes.new(type="ShaderNodeOutputWorld")
+    links.new(bg_node.outputs["Background"], output_node.inputs["Surface"])
+
+
+def init_random_lighting(camera_dir: np.ndarray) -> None:
+    # Clear existing lights
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.ops.object.select_by_type(type="LIGHT")
+    bpy.ops.object.delete()
+    
+    # Create environment light
+    if bpy.context.scene.world is None:
+        world = bpy.data.worlds.new("World")
+        bpy.context.scene.world = world
+    else:
+        world = bpy.context.scene.world
+        
+    # Enabling nodes
+    world.use_nodes = True
+    node_tree = world.node_tree
+    nodes = node_tree.nodes
+    links = node_tree.links
+    
+    # Remove default nodes
+    for node in nodes:
+        nodes.remove(node)
+    
+    # Random place lights
+    num_lights = np.random.randint(1, 4)
+    total_strength = 1.5
+    for i in range(num_lights):
+        new_light = bpy.data.objects.new(f"Light_{i}", bpy.data.lights.new(f"Light_{i}", type="POINT"))
+        bpy.context.collection.objects.link(new_light)
+        
+        new_light_distance = 1 / np.random.uniform(1/100, 1/10)
+        new_light_dir = np.random.randn(3)
+        new_light_dir[2] += 0.6
+        new_light_dir = new_light_dir / np.linalg.norm(new_light_dir)
+        new_light_location = new_light_dir * new_light_distance
+        new_light_camera_strength_ratio = max(np.sum(camera_dir * new_light_dir) * 0.5 + 0.5, 0)
+        new_light_max_energy = total_strength / (np.sum(camera_dir * new_light_dir) * 0.45 + 0.55)
+        new_light_strength = np.sqrt(np.random.uniform(0.01, 1)) * new_light_max_energy
+        new_light_camera_strength = new_light_camera_strength_ratio * new_light_strength
+        total_strength -= new_light_camera_strength
+        
+        new_light.location = (new_light_location[0], new_light_location[1], new_light_location[2])
+        new_light.data.color = (1.0, 1.0, 1.0)
+        new_light.data.energy = new_light_strength * new_light_distance**2 * 31.4
+        new_light.data.shadow_soft_size = np.random.uniform(0.1, 0.1 * new_light_distance)
+        
+    # Create background node
+    bg_node = nodes.new(type="ShaderNodeBackground")
+    bg_node.inputs["Color"].default_value = (1.0, 1.0, 1.0, 1.0)
+    bg_node.inputs["Strength"].default_value = total_strength
+    output_node = nodes.new(type="ShaderNodeOutputWorld")
+    links.new(bg_node.outputs["Background"], output_node.inputs["Surface"])
+
+
+def load_object(object_path: str) -> None:
+    """Loads a model with a supported file extension into the scene.
+
+    Args:
+        object_path (str): Path to the model file.
+
+    Raises:
+        ValueError: If the file extension is not supported.
+
+    Returns:
+        None
+    """
+    file_extension = object_path.split(".")[-1].lower()
+    if file_extension is None:
+        raise ValueError(f"Unsupported file type: {object_path}")
+
+    if file_extension == "usdz":
+        # install usdz io package
+        dirname = os.path.dirname(os.path.realpath(__file__))
+        usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
+        bpy.ops.preferences.addon_install(filepath=usdz_package)
+        # enable it
+        addon_name = "io_scene_usdz"
+        bpy.ops.preferences.addon_enable(module=addon_name)
+        # import the usdz
+        from io_scene_usdz.import_usdz import import_usdz
+
+        import_usdz(context, filepath=object_path, materials=True, animations=True)
+        return None
+
+    # load from existing import functions
+    import_function = IMPORT_FUNCTIONS[file_extension]
+
+    print(f"Loading object from {object_path}")
+    if file_extension == "blend":
+        import_function(directory=object_path, link=False)
+    elif file_extension in {"glb", "gltf"}:
+        import_function(filepath=object_path, merge_vertices=True, import_shading='NORMALS')
+    else:
+        import_function(filepath=object_path)
+        
+
+def delete_invisible_objects() -> None:
+    """Deletes all invisible objects in the scene.
+
+    Returns:
+        None
+    """
+    # bpy.ops.object.mode_set(mode="OBJECT")
+    bpy.ops.object.select_all(action="DESELECT")
+    for obj in bpy.context.scene.objects:
+        if obj.hide_viewport or obj.hide_render:
+            obj.hide_viewport = False
+            obj.hide_render = False
+            obj.hide_select = False
+            obj.select_set(True)
+    bpy.ops.object.delete()
+
+    # Delete invisible collections
+    invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
+    for col in invisible_collections:
+        bpy.data.collections.remove(col)
+
+
+def unhide_all_objects() -> None:
+    """Unhides all objects in the scene.
+
+    Returns:
+        None
+    """
+    for obj in bpy.context.scene.objects:
+        obj.hide_set(False)
+        
+        
+def convert_to_meshes() -> None:
+    """Converts all objects in the scene to meshes.
+
+    Returns:
+        None
+    """
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.context.view_layer.objects.active = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"][0]
+    for obj in bpy.context.scene.objects:
+        obj.select_set(True)
+    bpy.ops.object.convert(target="MESH")
+    
+        
+def triangulate_meshes() -> None:
+    """Triangulates all meshes in the scene.
+
+    Returns:
+        None
+    """
+    bpy.ops.object.select_all(action="DESELECT")
+    objs = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
+    bpy.context.view_layer.objects.active = objs[0]
+    for obj in objs:
+        obj.select_set(True)
+    bpy.ops.object.mode_set(mode="EDIT")
+    bpy.ops.mesh.reveal()
+    bpy.ops.mesh.select_all(action="SELECT")
+    bpy.ops.mesh.quads_convert_to_tris(quad_method="BEAUTY", ngon_method="BEAUTY")
+    bpy.ops.object.mode_set(mode="OBJECT")
+    bpy.ops.object.select_all(action="DESELECT")
+    
+
+def scene_bbox() -> Tuple[Vector, Vector]:
+    """Returns the bounding box of the scene.
+
+    Taken from Shap-E rendering script
+    (https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)
+
+    Returns:
+        Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
+    """
+    bbox_min = (math.inf,) * 3
+    bbox_max = (-math.inf,) * 3
+    found = False
+    scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
+    for obj in scene_meshes:
+        found = True
+        for coord in obj.bound_box:
+            coord = Vector(coord)
+            coord = obj.matrix_world @ coord
+            bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
+            bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
+    if not found:
+        raise RuntimeError("no objects in scene to compute bounding box for")
+    return Vector(bbox_min), Vector(bbox_max)
+
+
+def normalize_scene() -> Tuple[float, Vector]:
+    """Normalizes the scene by scaling and translating it to fit in a unit cube centered
+    at the origin.
+
+    Mostly taken from the Point-E / Shap-E rendering script
+    (https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
+    but fix for multiple root objects: (see bug report here:
+    https://github.com/openai/shap-e/pull/60).
+
+    Returns:
+        Tuple[float, Vector]: The scale factor and the offset applied to the scene.
+    """
+    scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
+    if len(scene_root_objects) > 1:
+        # create an empty object to be used as a parent for all root objects
+        scene = bpy.data.objects.new("ParentEmpty", None)
+        bpy.context.scene.collection.objects.link(scene)
+
+        # parent all root objects to the empty object
+        for obj in scene_root_objects:
+            obj.parent = scene
+    else:
+        scene = scene_root_objects[0]
+
+    bbox_min, bbox_max = scene_bbox()
+    scale = 1 / max(bbox_max - bbox_min)
+    scene.scale = scene.scale * scale
+
+    # Apply scale to matrix_world.
+    bpy.context.view_layer.update()
+    bbox_min, bbox_max = scene_bbox()
+    offset = -(bbox_min + bbox_max) / 2
+    scene.matrix_world.translation += offset
+    bpy.ops.object.select_all(action="DESELECT")
+    
+    return scale, offset
+
+
+def get_transform_matrix(obj: bpy.types.Object) -> list:
+    pos, rt, _ = obj.matrix_world.decompose()
+    rt = rt.to_matrix()
+    matrix = []
+    for ii in range(3):
+        a = []
+        for jj in range(3):
+            a.append(rt[ii][jj])
+        a.append(pos[ii])
+        matrix.append(a)
+    matrix.append([0, 0, 0, 1])
+    return matrix
+
+
+def check_mask_boundary_distance(image_path: str, threshold: int = 0) -> Tuple[bool, bool, int]:
+    """Check the rendered object's mask distance to image boundary.
+    
+    Args:
+        image_path: Path to the rendered PNG image with alpha channel.
+        threshold: Alpha value threshold to consider as valid mask (0-255).
+        
+    Returns:
+        Tuple of (touches_boundary, too_far_from_boundary, min_distance):
+        - touches_boundary: True if the mask touches any boundary
+        - too_far_from_boundary: True if the mask is too far from all boundaries (>80 pixels)
+        - min_distance: Minimum distance from mask to any boundary
+    """
+    img = Image.open(image_path)
+    if img.mode != 'RGBA':
+        return False, False, 0
+    
+    # Get alpha channel
+    alpha = np.array(img)[:, :, 3]
+    h, w = alpha.shape
+    
+    # Find all pixels with alpha > threshold (mask pixels)
+    mask_pixels = np.where(alpha > threshold)
+    
+    if len(mask_pixels[0]) == 0:
+        # No valid mask pixels
+        return False, True, max(h, w)
+    
+    # Calculate distances to each boundary
+    min_row = np.min(mask_pixels[0])  # Distance to top edge
+    max_row = np.max(mask_pixels[0])  # Distance to bottom edge
+    min_col = np.min(mask_pixels[1])  # Distance to left edge
+    max_col = np.max(mask_pixels[1])  # Distance to right edge
+    
+    dist_top = min_row
+    dist_bottom = (h - 1) - max_row
+    dist_left = min_col
+    dist_right = (w - 1) - max_col
+    
+    min_distance = min(dist_top, dist_bottom, dist_left, dist_right)
+    
+    # Check if touches boundary (distance <= 0)
+    touches_boundary = min_distance <= 0
+    
+    # Check if too far from boundary (distance > 130 pixels)
+    too_far = min_distance > 130
+    
+    return touches_boundary, too_far, min_distance
+
+
+def main(arg):
+    if arg.object.endswith(".blend"):
+        delete_invisible_objects()
+    else:
+        init_scene()
+        load_object(arg.object)
+    print('[INFO] Scene initialized.')
+    
+    # normalize scene
+    scale, offset = normalize_scene()
+    print('[INFO] Scene normalized.')
+    
+    # Initialize camera and lighting
+    cam = init_camera()
+    init_uniform_lighting()
+    print('[INFO] Camera and lighting initialized.')
+        
+    # ============= Render conditional views =============
+    init_render(engine=arg.engine, resolution=arg.cond_resolution)
+    # Create a list of views
+    to_export = {
+        "aabb": [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        "scale": scale,
+        "offset": [offset.x, offset.y, offset.z],
+        "frames": []
+    }
+    views = json.loads(arg.cond_views)
+    
+    # Parameters for boundary check and radius adjustment
+    max_retry = 10  # Maximum number of retries per view
+    radius_increase_factor = 1.1  # Increase radius by 10% when too close to boundary
+    radius_decrease_factor = 0.9  # Decrease radius by 10% when too far from boundary
+    min_boundary_distance = 130  # Minimum distance to boundary in pixels
+    
+    for i, view in enumerate(views):
+        current_radius = view['radius']
+        retry_count = 0
+        
+        while retry_count < max_retry:
+            cam_dir = np.array([
+                np.cos(view['yaw']) * np.cos(view['pitch']),
+                np.sin(view['yaw']) * np.cos(view['pitch']),
+                np.sin(view['pitch'])
+            ])
+            init_random_lighting(cam_dir)
+            cam.location = (
+                current_radius * cam_dir[0],
+                current_radius * cam_dir[1],
+                current_radius * cam_dir[2]
+            )
+            cam.data.lens = 16 / np.tan(view['fov'] / 2)
+            
+            output_path = os.path.join(arg.cond_output_folder, f'{i:03d}.png')
+            bpy.context.scene.render.filepath = output_path
+                
+            # Render the scene
+            bpy.ops.render.render(write_still=True)
+            bpy.context.view_layer.update()
+            
+            # Check mask boundary distance
+            touches_boundary, too_far, min_dist = check_mask_boundary_distance(output_path)
+            
+            if touches_boundary:
+                # Object is too close to boundary, increase radius
+                retry_count += 1
+                old_radius = current_radius
+                current_radius *= radius_increase_factor
+                print(f'[WARNING] View {i}: Mask touches boundary (dist={min_dist}px). Increasing radius from {old_radius:.4f} to {current_radius:.4f} (retry {retry_count}/{max_retry})')
+            elif too_far:
+                # Object is too far from boundary (>80px), decrease radius
+                retry_count += 1
+                old_radius = current_radius
+                current_radius *= radius_decrease_factor
+                print(f'[INFO] View {i}: Mask too far from boundary (dist={min_dist}px > {min_boundary_distance}px). Decreasing radius from {old_radius:.4f} to {current_radius:.4f} (retry {retry_count}/{max_retry})')
+            else:
+                # Good distance, stop retrying
+                if retry_count > 0:
+                    print(f'[INFO] View {i}: Mask boundary distance OK (dist={min_dist}px) after {retry_count} retries. Final radius: {current_radius:.4f}')
+                else:
+                    print(f'[INFO] View {i}: Mask boundary distance OK (dist={min_dist}px). Radius: {current_radius:.4f}')
+                break
+        
+        if retry_count >= max_retry:
+            print(f'[WARNING] View {i}: Max retries reached. Using final radius: {current_radius:.4f} (dist={min_dist}px)')
+            
+        # Save camera parameters (with potentially updated radius)
+        metadata = {
+            "file_path": f'{i:03d}.png',
+            "camera_angle_x": view['fov'],
+            "transform_matrix": get_transform_matrix(cam),
+            "radius": current_radius,  # Save the actual radius used
+            "original_radius": view['radius'],  # Save original for reference
+            "retries": retry_count
+        }
+        to_export["frames"].append(metadata)
+    
+    # Save the camera parameters
+    with open(os.path.join(arg.cond_output_folder, 'transforms.json'), 'w') as f:
+        json.dump(to_export, f, indent=4)
+
+        
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
+    parser.add_argument('--object', type=str, help='Path to the 3D model file to be rendered.')
+    parser.add_argument('--cond_views', type=str, help='JSON string of views. Contains a list of {yaw, pitch, radius, fov} object.')
+    parser.add_argument('--cond_output_folder', type=str, default='/tmp', help='The path the output will be dumped to.')
+    parser.add_argument('--cond_resolution', type=int, default=1024, help='Resolution of the conditional images.')
+    parser.add_argument('--engine', type=str, default='CYCLES', help='Blender internal engine for rendering. E.g. CYCLES, BLENDER_EEVEE, ...')
+    argv = sys.argv[sys.argv.index("--") + 1:]
+    args = parser.parse_args(argv)
+
+    main(args)
+