blender_plugin/interiorfusion_blender.py

"""Blender plugin for InteriorFusion.

Features:
- Generate 3D scene from reference image
- Import generated meshes with PBR materials
- Interactive scene editing within Blender
- Export to game engines
"""

import os
import tempfile
import bpy
import bmesh
from bpy.props import StringProperty, BoolProperty, EnumProperty, FloatProperty
from bpy.types import Operator, Panel
from mathutils import Vector, Matrix


bl_info = {
    "name": "InteriorFusion",
    "author": "InteriorFusion Research Team",
    "version": (0, 1, 0),
    "blender": (3, 6, 0),
    "location": "3D Viewport > Sidebar > InteriorFusion",
    "description": "Single image to editable 3D interior scene",
    "category": "3D View",
    "support": "COMMUNITY",
}


class INTERIORFUSION_OT_generate_scene(Operator):
    """Generate 3D interior scene from reference image."""
    bl_idname = "interiorfusion.generate_scene"
    bl_label = "Generate 3D Scene"
    bl_options = {'REGISTER', 'UNDO'}
    
    image_path: StringProperty(
        name="Image Path",
        description="Path to interior photo",
        subtype='FILE_PATH',
    )
    
    room_type: EnumProperty(
        name="Room Type",
        items=[
            ('AUTO', 'Auto Detect', 'Automatically detect room type'),
            ('LIVING_ROOM', 'Living Room', 'Living room / lounge'),
            ('BEDROOM', 'Bedroom', 'Bedroom'),
            ('KITCHEN', 'Kitchen', 'Kitchen'),
            ('DINING_ROOM', 'Dining Room', 'Dining room'),
            ('OFFICE', 'Office', 'Office / study'),
        ],
        default='AUTO',
    )
    
    style: EnumProperty(
        name="Style",
        items=[
            ('AUTO', 'Auto Detect', 'Automatically detect style'),
            ('MODERN', 'Modern', 'Modern contemporary'),
            ('SCANDINAVIAN', 'Scandinavian', 'Scandinavian minimalist'),
            ('LUXURY', 'Luxury', 'Luxury / upscale'),
            ('INDUSTRIAL', 'Industrial', 'Industrial loft'),
            ('MINIMALIST', 'Minimalist', 'Minimalist'),
        ],
        default='AUTO',
    )
    
    use_pbr: BoolProperty(
        name="Use PBR Materials",
        description="Generate metallic/roughness/normal maps",
        default=True,
    )
    
    def execute(self, context):
        from interiorfusion.pipelines import InteriorFusionPipeline
        from PIL import Image
        
        # Check image exists
        if not os.path.exists(self.image_path):
            self.report({'ERROR'}, f"Image not found: {self.image_path}")
            return {'CANCELLED'}
        
        # Load image
        image = Image.open(self.image_path).convert("RGB")
        
        # Generate scene
        self.report({'INFO'}, "Generating 3D scene...")
        
        pipeline = InteriorFusionPipeline(
            model_size="L",
            device="cuda" if bpy.app.version >= (3, 5) else "cpu",
            use_pbr=self.use_pbr,
        )
        
        output = pipeline(
            image=image,
            room_type_hint=self.room_type if self.room_type != 'AUTO' else None,
            style_hint=self.style if self.style != 'AUTO' else None,
        )
        
        # Import scene into Blender
        self.import_scene(context, output)
        
        self.report({'INFO'}, 
                    f"Scene generated: {output.room_type} ({output.processing_time:.1f}s)")
        
        return {'FINISHED'}
    
    def import_scene(self, context, output):
        """Import generated scene into Blender."""
        # Import room shell
        if output.room_shell_mesh is not None:
            self.import_mesh(output.room_shell_mesh, "Room_Shell")
        
        # Import objects
        for i, obj_mesh in enumerate(output.object_meshes):
            obj_name = f"Furniture_{i:02d}"
            self.import_mesh(obj_mesh, obj_name)
        
        # Create scene graph collection
        scene_collection = bpy.data.collections.new("InteriorFusion_Scene")
        context.scene.collection.children.link(scene_collection)
        
        # Move objects to scene collection
        for obj in context.selected_objects:
            scene_collection.objects.link(obj)
            context.scene.collection.objects.unlink(obj)
    
    def import_mesh(self, mesh, name):
        """Import a trimesh mesh into Blender."""
        try:
            import trimesh
        except ImportError:
            self.report({'WARNING'}, "trimesh not available, skipping mesh import")
            return None
        
        # Create Blender mesh
        bm = bmesh.new()
        
        # Add vertices
        verts = {}
        for i, v in enumerate(mesh.vertices):
            verts[i] = bm.verts.new(Vector(v))
        
        # Add faces
        bm.verts.ensure_lookup_table()
        for face in mesh.faces:
            try:
                face_verts = [verts[v_idx] for v_idx in face]
                bm.faces.new(face_verts)
            except Exception:
                pass
        
        # Create mesh object
        mesh_data = bpy.data.meshes.new(name)
        bm.to_mesh(mesh_data)
        bm.free()
        
        obj = bpy.data.objects.new(name, mesh_data)
        bpy.context.collection.objects.link(obj)
        
        # Apply materials if available
        if hasattr(mesh, 'materials') and mesh.materials:
            for mat_name, mat_data in mesh.materials.items():
                mat = self.create_pbr_material(mat_name, mat_data)
                mesh_data.materials.append(mat)
        
        return obj
    
    def create_pbr_material(self, name, material_data):
        """Create a Blender PBR material."""
        mat = bpy.data.materials.new(name=name)
        mat.use_nodes = True
        
        # Get principled BSDF
        bsdf = mat.node_tree.nodes["Principled BSDF"]
        
        # Set parameters
        albedo = material_data.get("albedo", [0.7, 0.7, 0.7])
        bsdf.inputs['Base Color'].default_value = (*albedo, 1.0)
        bsdf.inputs['Metallic'].default_value = material_data.get("metallic", 0.0)
        bsdf.inputs['Roughness'].default_value = material_data.get("roughness", 0.5)
        
        return mat


class INTERIORFUSION_OT_edit_object(Operator):
    """Edit a selected furniture object."""
    bl_idname = "interiorfusion.edit_object"
    bl_label = "Edit Selected Object"
    bl_options = {'REGISTER', 'UNDO'}
    
    action: EnumProperty(
        name="Action",
        items=[
            ('MOVE', 'Move', 'Move to new position'),
            ('REPLACE', 'Replace', 'Replace with new object'),
            ('REMOVE', 'Remove', 'Remove from scene'),
            ('SCALE', 'Scale', 'Change dimensions'),
        ],
    )
    
    def execute(self, context):
        obj = context.active_object
        if obj is None:
            self.report({'ERROR'}, "No object selected")
            return {'CANCELLED'}
        
        if self.action == 'REMOVE':
            bpy.data.objects.remove(obj, do_unlink=True)
            self.report({'INFO'}, f"Removed {obj.name}")
            
        elif self.action == 'MOVE':
            # Enter move mode
            bpy.ops.transform.translate('INVOKE_DEFAULT')
            
        elif self.action == 'SCALE':
            # Enter scale mode
            bpy.ops.transform.resize('INVOKE_DEFAULT')
        
        return {'FINISHED'}


class INTERIORFUSION_PT_panel(Panel):
    """InteriorFusion main panel."""
    bl_label = "InteriorFusion"
    bl_idname = "INTERIORFUSION_PT_panel"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = 'InteriorFusion'
    
    def draw(self, context):
        layout = self.layout
        
        # Scene generation
        box = layout.box()
        box.label(text="Scene Generation", icon='SCENE_DATA')
        
        box.prop(context.scene, "interiorfusion_image_path")
        box.prop(context.scene, "interiorfusion_room_type")
        box.prop(context.scene, "interiorfusion_style")
        box.prop(context.scene, "interiorfusion_use_pbr")
        
        box.operator("interiorfusion.generate_scene", icon='MESH_CUBE')
        
        # Object editing
        box = layout.box()
        box.label(text="Object Editing", icon='OBJECT_DATA')
        
        if context.active_object:
            box.label(text=f"Selected: {context.active_object.name}")
            row = box.row()
            row.operator("interiorfusion.edit_object", text="Move").action = 'MOVE'
            row.operator("interiorfusion.edit_object", text="Scale").action = 'SCALE'
            row = box.row()
            row.operator("interiorfusion.edit_object", text="Remove").action = 'REMOVE'
        else:
            box.label(text="Select an object to edit")
        
        # Export
        box = layout.box()
        box.label(text="Export", icon='EXPORT')
        box.operator("export_scene.gltf", text="Export GLB", icon='EXPORT')
        box.operator("wm.obj_export", text="Export OBJ", icon='EXPORT')


def register():
    # Scene properties
    bpy.types.Scene.interiorfusion_image_path = StringProperty(
        name="Image Path",
        description="Path to interior photo",
        subtype='FILE_PATH',
    )
    
    bpy.types.Scene.interiorfusion_room_type = EnumProperty(
        name="Room Type",
        items=[
            ('AUTO', 'Auto Detect', 'Auto'),
            ('LIVING_ROOM', 'Living Room', 'Living room'),
            ('BEDROOM', 'Bedroom', 'Bedroom'),
            ('KITCHEN', 'Kitchen', 'Kitchen'),
            ('DINING_ROOM', 'Dining Room', 'Dining room'),
            ('OFFICE', 'Office', 'Office'),
        ],
        default='AUTO',
    )
    
    bpy.types.Scene.interiorfusion_style = EnumProperty(
        name="Style",
        items=[
            ('AUTO', 'Auto Detect', 'Auto'),
            ('MODERN', 'Modern', 'Modern'),
            ('SCANDINAVIAN', 'Scandinavian', 'Scandinavian'),
            ('LUXURY', 'Luxury', 'Luxury'),
            ('INDUSTRIAL', 'Industrial', 'Industrial'),
            ('MINIMALIST', 'Minimalist', 'Minimalist'),
        ],
        default='AUTO',
    )
    
    bpy.types.Scene.interiorfusion_use_pbr = BoolProperty(
        name="Use PBR",
        description="Generate PBR materials",
        default=True,
    )
    
    # Register classes
    bpy.utils.register_class(INTERIORFUSION_OT_generate_scene)
    bpy.utils.register_class(INTERIORFUSION_OT_edit_object)
    bpy.utils.register_class(INTERIORFUSION_PT_panel)


def unregister():
    bpy.utils.unregister_class(INTERIORFUSION_PT_panel)
    bpy.utils.unregister_class(INTERIORFUSION_OT_edit_object)
    bpy.utils.unregister_class(INTERIORFUSION_OT_generate_scene)
    
    del bpy.types.Scene.interiorfusion_image_path
    del bpy.types.Scene.interiorfusion_room_type
    del bpy.types.Scene.interiorfusion_style
    del bpy.types.Scene.interiorfusion_use_pbr


if __name__ == "__main__":
    register()