clean_mesh_blender.py

import sys
import os

try:
    import bpy
except ImportError:
    print("[Error] This script must be run inside Blender's python environment.")
    sys.exit(1)

def delete_isolated_components(obj, threshold_vertices=100):
    # Store original name
    orig_name = obj.name
    
    # Deselect all, select obj, make active
    bpy.ops.object.select_all(action='DESELECT')
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    
    # Enter Edit Mode to separate
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    
    # Separate loose parts
    bpy.ops.mesh.separate(type='LOOSE')
    
    # Go back to Object Mode to process parts
    bpy.ops.object.mode_set(mode='OBJECT')
    
    # The separated parts will be named like orig_name.001, orig_name.002, etc.
    # Collect all mesh objects that are part of this separation
    separated_objs = [o for o in bpy.context.scene.objects if o.type == 'MESH' and (o.name == orig_name or o.name.startswith(orig_name + "."))]
    
    if len(separated_objs) <= 1:
        print("[Blender] Mesh is already a single component.")
        return obj
        
    print(f"[Blender] Separated into {len(separated_objs)} loose components.")
    
    max_verts = max(len(o.data.vertices) for o in separated_objs)
    
    keep_objs = []
    delete_objs = []
    
    for o in separated_objs:
        num_verts = len(o.data.vertices)
        # Keep if it has at least threshold_vertices, or is at least 5% of the largest component size
        if num_verts >= threshold_vertices or num_verts >= (max_verts * 0.05):
            keep_objs.append(o)
        else:
            delete_objs.append(o)
            
    if not keep_objs:
        largest = max(separated_objs, key=lambda o: len(o.data.vertices))
        keep_objs = [largest]
        delete_objs = [o for o in separated_objs if o != largest]
        
    if delete_objs:
        print(f"[Blender] Deleting {len(delete_objs)} isolated noise components with vertex count < {threshold_vertices}.")
        bpy.ops.object.select_all(action='DESELECT')
        for o in delete_objs:
            o.select_set(True)
        bpy.ops.object.delete()
        
    # Re-join remaining ones
    bpy.ops.object.select_all(action='DESELECT')
    for o in keep_objs:
        o.select_set(True)
    bpy.context.view_layer.objects.active = keep_objs[0]
    bpy.ops.object.join()
    
    joined_obj = bpy.context.active_object
    joined_obj.name = orig_name
    return joined_obj

def strip_gltf_extensions(glb_path):
    try:
        with open(glb_path, "rb") as f:
            data = f.read()
        if len(data) < 20 or data[:4] != b'glTF':
            return
        json_len = int.from_bytes(data[12:16], byteorder='little')
        json_bytes = data[20:20+json_len]
        import json
        gltf_json = json.loads(json_bytes.decode('utf-8', errors='ignore'))
        
        modified = False
        for key in ['extensionsRequired', 'extensionsUsed']:
            if key in gltf_json and 'EXT_texture_webp' in gltf_json[key]:
                gltf_json[key].remove('EXT_texture_webp')
                modified = True
                
        if modified:
            new_bytes = json.dumps(gltf_json).encode('utf-8')
            if len(new_bytes) <= len(json_bytes):
                new_bytes = new_bytes.ljust(len(json_bytes), b' ')
                new_data = data[:20] + new_bytes + data[20+len(json_bytes):]
                with open(glb_path, "wb") as f:
                    f.write(new_data)
                print(f"[Blender] Stripped EXT_texture_webp extension requirement from GLB.")
    except Exception as e:
        print(f"[Blender] Extension strip note: {e}")

def fix_materials_for_fbx_export():
    """
    Fixes black material issue during FBX export by:
    1. Resetting Principled BSDF Base Color tint factor to pure white (1, 1, 1, 1).
    2. Unpacking images to disk so FBX embed_textures can copy them.
    """
    print("[Blender] Fixing material properties and unpacking textures for bright FBX export...")
    for mat in bpy.data.materials:
        if not mat or not mat.use_nodes:
            continue
        for node in mat.node_tree.nodes:
            if node.type == 'BSDF_PRINCIPLED':
                try:
                    node.inputs['Base Color'].default_value = (1.0, 1.0, 1.0, 1.0)
                except Exception:
                    pass
                if 'Specular' in node.inputs:
                    try:
                        node.inputs['Specular'].default_value = 0.5
                    except Exception:
                        pass
                if 'Specular IOR Level' in node.inputs:
                    try:
                        node.inputs['Specular IOR Level'].default_value = 0.5
                    except Exception:
                        pass
                if 'Roughness' in node.inputs:
                    try:
                        node.inputs['Roughness'].default_value = 0.5
                    except Exception:
                        pass
                if 'Metallic' in node.inputs:
                    try:
                        node.inputs['Metallic'].default_value = 0.0
                    except Exception:
                        pass

    for img in bpy.data.images:
        if img and img.has_data:
            try:
                if img.packed_file:
                    img.unpack(method='WRITE_LOCAL')
            except Exception as e:
                print(f"[Blender] Image unpack note for {img.name}: {e}")

def run_blender_cleanup(input_path, output_path, target_faces=8000, remesh_method="cleanup"):
    print(f"[Blender] Loading model: {input_path} (target_faces={target_faces}, remesh_method={remesh_method})")
    
    # 1. Reset scene to factory default empty scene
    bpy.ops.wm.read_factory_settings(use_empty=True)
    
    # 2. Import GLTF/GLB
    strip_gltf_extensions(input_path)
    bpy.ops.import_scene.gltf(filepath=input_path)
    
    # 3. Locate the imported mesh
    mesh_objs = [obj for obj in bpy.context.scene.objects if obj.type == 'MESH']
    if not mesh_objs:
        print("[Error] No mesh found in imported GLTF.")
        sys.exit(1)
        
    dirty_obj = mesh_objs[0]
    dirty_obj.name = "DirtyMesh"
    print(f"[Blender] Found source mesh: {dirty_obj.name} with {len(dirty_obj.data.polygons)} polygons")
    
    # DIAGNOSTIC: Print material and texture details of DirtyMesh
    print("[Blender] --- DirtyMesh Material Diagnostics ---")
    for i, mat in enumerate(dirty_obj.data.materials):
        if not mat:
            print(f"  Material {i}: None")
            continue
        print(f"  Material {i}: {mat.name} (use_nodes={mat.use_nodes})")
        if mat.use_nodes:
            for node in mat.node_tree.nodes:
                print(f"    Node: {node.name} (type={node.type})")
                if node.type == 'TEX_IMAGE':
                    img = node.image
                    if img:
                        print(f"      Image: {img.name}, size={img.size[:]}, filepath='{img.filepath}', has_data={img.has_data}")
                    else:
                        print("      Image: None")
            for link in mat.node_tree.links:
                print(f"    Link: {link.from_node.name} ({link.from_socket.name}) -> {link.to_node.name} ({link.to_socket.name})")
    print("[Blender] ---------------------------------------")
    
    # Fast non-destructive cleanup option
    if remesh_method == "cleanup":
        print("[Blender] Running fast cleanup (merging duplicates and removing isolated parts)...")
        # 1. Clean up mesh (remove doubles & loose)
        bpy.ops.object.select_all(action='DESELECT')
        dirty_obj.select_set(True)
        bpy.context.view_layer.objects.active = dirty_obj
        bpy.ops.object.mode_set(mode='EDIT')
        bpy.ops.mesh.select_all(action='SELECT')
        bpy.ops.mesh.remove_doubles(threshold=0.001)
        bpy.ops.mesh.delete_loose()
        bpy.ops.mesh.tris_convert_to_quads(face_threshold=1.570796, shape_threshold=1.570796)
        bpy.ops.mesh.normals_make_consistent(inside=False)
        bpy.ops.object.mode_set(mode='OBJECT')
        
        # 2. Delete isolated components
        dirty_obj = delete_isolated_components(dirty_obj, threshold_vertices=100)
        
        # 3. Export directly (retaining original materials/textures)
        bpy.ops.object.select_all(action='DESELECT')
        dirty_obj.select_set(True)
        bpy.context.view_layer.objects.active = dirty_obj
        
        print(f"[Blender] Exporting cleaned model to: {output_path}")
        bpy.ops.export_scene.gltf(
            filepath=output_path,
            export_format='GLB',
            use_selection=True
        )
        
        # Also export FBX
        fbx_path = output_path.replace(".glb", ".fbx")
        print(f"[Blender] Exporting cleaned model as FBX: {fbx_path}")
        try:
            fix_materials_for_fbx_export()
            bpy.ops.export_scene.fbx(
                filepath=fbx_path,
                use_selection=True,
                path_mode='COPY',
                embed_textures=True
            )
            print("[Blender] ✓ FBX export completed successfully.")
        except Exception as fe:
            print(f"[Blender] Error: FBX export failed: {fe}")
            
        print("[Blender] Fast cleanup completed successfully.")
        return

    # 4. Duplicate the mesh to create the clean version
    bpy.ops.object.select_all(action='DESELECT')
    dirty_obj.select_set(True)
    bpy.context.view_layer.objects.active = dirty_obj
    
    bpy.ops.object.duplicate(linked=False)
    clean_obj = bpy.context.active_object
    clean_obj.name = "CleanMesh"
    
    # 5. Clean up initial mesh (remove doubles)
    bpy.context.view_layer.objects.active = clean_obj
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    # Merge close vertices (Merge by distance)
    bpy.ops.mesh.remove_doubles(threshold=0.001)
    bpy.ops.mesh.delete_loose()
    # Recalculate normals outwards
    bpy.ops.mesh.normals_make_consistent(inside=False)
    bpy.ops.object.mode_set(mode='OBJECT')
    
    # Delete isolated components before remeshing/decimating
    clean_obj = delete_isolated_components(clean_obj, threshold_vertices=100)
    
    # 6. Apply Decimate Modifier or QuadriFlow Remesh
    if remesh_method == "quadriflow":
        print("[Blender] Preparing mesh for QuadriFlow (running high-res Voxel Remesh to guarantee manifold geometry)...")
        try:
            bpy.context.view_layer.objects.active = clean_obj
            bbox_size = max(clean_obj.dimensions)
            voxel_size = max(0.002, bbox_size / 200.0) # clean high-resolution voxel shell
            clean_obj.data.remesh_voxel_size = voxel_size
            bpy.ops.object.voxel_remesh()
            print("[Blender] Voxel Remesh completed successfully.")
        except Exception as vre:
            print(f"[Blender] Voxel Remesh preparation warning: {vre}")

        print(f"[Blender] Running QuadriFlow Remesh to target {target_faces} faces...")
        qf_success = False
        try:
            bpy.context.view_layer.objects.active = clean_obj
            # Ensure object mode and valid selection
            if bpy.context.active_object and bpy.context.active_object.mode != 'OBJECT':
                bpy.ops.object.mode_set(mode='OBJECT')
            
            # Check if operator can run in headless background mode
            if hasattr(bpy.ops.object, "quadriflow_remesh") and bpy.ops.object.quadriflow_remesh.poll():
                bpy.ops.object.quadriflow_remesh(
                    use_preserve_sharp=True,
                    use_preserve_boundary=True,
                    target_faces=target_faces
                )
                if len(clean_obj.data.polygons) <= target_faces * 1.5:
                    qf_success = True
                    print(f"[Blender] ✓ QuadriFlow Remesh completed. New polygon count: {len(clean_obj.data.polygons)}")
        except Exception as qfe:
            print(f"[Blender] QuadriFlow Remesh notice: {qfe}")

        if not qf_success:
            print("[Blender] QuadriFlow unavailable in headless mode. Using high-precision Decimation + Tris-to-Quads fallback.")
            remesh_method = "tris_to_quads"

            
    if remesh_method == "tris_to_quads":
        poly_count_before = len(clean_obj.data.polygons)
        if poly_count_before > target_faces:
            ratio = target_faces / poly_count_before
            print(f"[Blender] Decimating mesh from {poly_count_before} to {target_faces} polygons (ratio={ratio:.4f})...")
            try:
                dec_mod = clean_obj.modifiers.new(name="AutoDecimate", type='DECIMATE')
                dec_mod.ratio = ratio
                dec_mod.use_symmetry = True
                dec_mod.symmetry_axis = 'X'
                bpy.ops.object.modifier_apply(modifier="AutoDecimate")
                print(f"[Blender] ✓ Decimation completed. New polygon count: {len(clean_obj.data.polygons)}")
            except Exception as de:
                print(f"[Blender] Warning: Decimation failed: {de}")
        else:
            print(f"[Blender] Mesh already has fewer polygons ({poly_count_before}) than target ({target_faces}). Skipping decimation.")

        # 7. Convert triangles to quads with high angle threshold (1.57 rad / 90 degrees)
        # This converts almost 100% of flat areas to quads while keeping the mesh completely watertight and sharp
        print("[Blender] Converting triangles to quads with high threshold...")
        try:
            bpy.ops.object.mode_set(mode='EDIT')
            bpy.ops.mesh.select_all(action='SELECT')
            bpy.ops.mesh.tris_convert_to_quads(face_threshold=1.570796, shape_threshold=1.570796)
            bpy.ops.object.mode_set(mode='OBJECT')
            print("[Blender] ✓ Conversion to quads completed.")
        except Exception as tq:
            print(f"[Blender] Warning: Triangles to quads conversion failed: {tq}")
            if clean_obj.mode != 'OBJECT':
                bpy.ops.object.mode_set(mode='OBJECT')
            
    # 8. Unwrap UVs on CleanMesh
    print("[Blender] Unwrapping UV coordinates...")
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.uv.smart_project(island_margin=0.002)
    bpy.ops.object.mode_set(mode='OBJECT')
    
    # 9. Setup Baking Material & Target Image Node
    print("[Blender] Setting up material for texture projection...")
    clean_mat = bpy.data.materials.new(name="CleanMaterial")
    clean_mat.use_nodes = True
    clean_obj.data.materials.clear()
    clean_obj.data.materials.append(clean_mat)
    
    bake_image = bpy.data.images.new(name="BakedTexture", width=2048, height=2048)
    
    nodes = clean_mat.node_tree.nodes
    texture_node = nodes.new(type='ShaderNodeTexImage')
    texture_node.image = bake_image
    texture_node.select = True
    nodes.active = texture_node 
    
    # 10. Configure Cycles and Bake from Selected (DirtyMesh) to Active (CleanMesh)
    print("[Blender] Configuring Cycles baking engine...")
    bpy.context.scene.render.engine = 'CYCLES'
    bpy.context.scene.cycles.device = 'CPU'
    
    print("[Blender] Bypassing shader nodes for unlit EMIT baking...")
    for mat in dirty_obj.data.materials:
        if mat and mat.use_nodes:
            output_node = next((n for n in mat.node_tree.nodes if n.type == 'OUTPUT_MATERIAL'), None)
            bsdf_node = next((n for n in mat.node_tree.nodes if n.type == 'BSDF_PRINCIPLED'), None)
            if output_node and bsdf_node:
                base_color_input = bsdf_node.inputs['Base Color']
                if base_color_input.is_linked:
                    from_node = base_color_input.links[0].from_node
                    from_socket = base_color_input.links[0].from_socket
                    mat.node_tree.links.new(from_socket, output_node.inputs['Surface'])
                    print(f"  [Bypass] Linked {from_node.name} directly to Material Output Surface on {mat.name}")
    
    bpy.context.scene.cycles.bake_type = 'EMIT'
    bpy.context.scene.render.bake.use_selected_to_active = True
    bpy.context.scene.render.bake.margin = 16
    bpy.context.scene.render.bake.margin_type = 'EXTEND'
    bpy.context.scene.render.bake.cage_extrusion = 0.02
    bpy.context.scene.render.bake.max_ray_distance = 0.05
    
    # Select DirtyMesh and CleanMesh, making CleanMesh active
    bpy.ops.object.select_all(action='DESELECT')
    dirty_obj.select_set(True)
    clean_obj.select_set(True)
    bpy.context.view_layer.objects.active = clean_obj
    
    print("[Blender] Projecting (baking) texture map using EMIT. Please wait...")
    try:
        bpy.ops.object.bake(type='EMIT')
        print("[Blender] ✓ Texture projection completed successfully.")
        bake_image.pack()
        bsdf_node = next(n for n in nodes if n.type == 'BSDF_PRINCIPLED')
        clean_mat.node_tree.links.new(texture_node.outputs['Color'], bsdf_node.inputs['Base Color'])
    except Exception as be:
        print(f"[Blender] Texture baking notice (falling back to original material & textures): {be}")
        clean_obj.data.materials.clear()
        for orig_mat in dirty_obj.data.materials:
            if orig_mat:
                clean_obj.data.materials.append(orig_mat)


    # 10. Delete the original dirty mesh before export
    bpy.ops.object.select_all(action='DESELECT')
    dirty_obj.select_set(True)
    bpy.ops.object.delete()
    
    # 11. Export CleanMesh as GLB (preserving original textures and materials)
    bpy.ops.object.select_all(action='DESELECT')
    clean_obj.select_set(True)
    bpy.context.view_layer.objects.active = clean_obj
    
    print(f"[Blender] Exporting clean model to: {output_path}")
    bpy.ops.export_scene.gltf(
        filepath=output_path,
        export_format='GLB',
        use_selection=True
    )
    
    # 12. Also export as FBX
    fbx_path = output_path.replace(".glb", ".fbx")
    print(f"[Blender] Exporting clean model as FBX: {fbx_path}")
    try:
        fix_materials_for_fbx_export()
        bpy.ops.export_scene.fbx(
            filepath=fbx_path,
            use_selection=True,
            path_mode='COPY',
            embed_textures=True
        )
        print("[Blender] ✓ FBX export completed successfully.")
    except Exception as fe:
        print(f"[Blender] Error: FBX export failed: {fe}")
            
    print("[Blender] Process completed successfully.")

if __name__ == "__main__":
    # Extract arguments passed after '--'
    try:
        args_idx = sys.argv.index("--")
        args = sys.argv[args_idx + 1:]
    except ValueError:
        args = []
        
    if len(args) < 2:
        print("[Error] Usage: blender --background --python clean_mesh_blender.py -- <input_glb> <output_glb> [target_faces] [remesh_method]")
        sys.exit(1)
        
    target_faces = 60000
    if len(args) >= 3:
        try:
            target_faces = int(args[2])
        except ValueError:
            pass
            
    remesh_method = "cleanup"
    if len(args) >= 4:
        remesh_method = args[3]
            
    run_blender_cleanup(args[0], args[1], target_faces, remesh_method)