blender-mask-peoples/core/compositor_setup.py

"""
Compositor Node Tree Setup for mask-based blur effect.

Creates and manages compositing node trees that apply blur
only to masked regions of a video strip.
"""

from typing import Optional, Tuple


def create_mask_blur_node_tree(
    name: str = "FaceMaskBlur",
    blur_size: int = 50,
) -> "bpy.types.NodeTree":
    """
    Create a compositing node tree for mask-based blur.
    
    Node structure:
    [Render Layers] ──┬──────────────────────────→ [Mix] → [Composite]
                      │                              ↑
                      └→ [Blur] → [Mix Factor Mask] ↗
    
    Args:
        name: Name for the node tree
        blur_size: Blur radius in pixels
        
    Returns:
        The created NodeTree
    """
    import bpy
    
    # Create new node tree or get existing
    if name in bpy.data.node_groups:
        # Return existing tree
        return bpy.data.node_groups[name]
    
    # Create compositing scene if needed
    tree = bpy.data.node_groups.new(name=name, type='CompositorNodeTree')
    tree.use_fake_user = True  # Prevent deletion
    
    nodes = tree.nodes
    links = tree.links
    
    # Clear default nodes
    nodes.clear()
    
    # Create nodes
    # Input: Image and Mask
    input_node = nodes.new('NodeGroupInput')
    input_node.location = (-400, 0)
    
    # Output
    output_node = nodes.new('NodeGroupOutput')
    output_node.location = (600, 0)
    
    # Blur node
    blur_node = nodes.new('CompositorNodeBlur')
    blur_node.location = (0, -150)
    blur_node.filter_type = 'GAUSS'
    blur_node.label = "Face Blur"
    # Note: Blender 5.0 uses 'Size' input socket instead of size_x/size_y properties
    # We'll set default_value on the socket after linking
    
    # Mix node (combines original with blurred using mask)
    mix_node = nodes.new('CompositorNodeMixRGB')
    mix_node.location = (300, 0)
    mix_node.blend_type = 'MIX'
    mix_node.label = "Mask Mix"
    
    # Set blur size via input socket (Blender 5.0 API)
    if 'Size' in blur_node.inputs:
        blur_node.inputs['Size'].default_value = blur_size / 100.0  # Size is 0-1 range
    elif 'size' in blur_node.inputs:
        blur_node.inputs['size'].default_value = blur_size / 100.0
    
    # Define interface sockets
    tree.interface.new_socket(
        name="Image",
        in_out='INPUT',
        socket_type='NodeSocketColor',
    )
    tree.interface.new_socket(
        name="Mask",
        in_out='INPUT',
        socket_type='NodeSocketFloat',
    )
    tree.interface.new_socket(
        name="Image",
        in_out='OUTPUT',
        socket_type='NodeSocketColor',
    )
    
    # Link nodes
    # Input Image → Blur
    links.new(input_node.outputs[0], blur_node.inputs['Image'])
    
    # Input Image → Mix (first color)
    links.new(input_node.outputs[0], mix_node.inputs[1])
    
    # Blur → Mix (second color)
    links.new(blur_node.outputs[0], mix_node.inputs[2])
    
    # Input Mask → Mix (factor)
    links.new(input_node.outputs[1], mix_node.inputs[0])
    
    # Mix → Output
    links.new(mix_node.outputs[0], output_node.inputs[0])
    
    return tree


def setup_strip_compositor_modifier(
    strip: "bpy.types.Strip",
    mask_strip: "bpy.types.Strip",
    node_tree: "bpy.types.NodeTree",
) -> "bpy.types.SequenceModifier":
    """
    Add a Compositor modifier to a strip using the mask-blur node tree.
    
    Args:
        strip: The video strip to add the modifier to
        mask_strip: The mask image sequence strip
        node_tree: The compositing node tree to use
        
    Returns:
        The created modifier
    """
    import bpy
    
    # Add compositor modifier
    modifier = strip.modifiers.new(
        name="FaceMaskBlur",
        type='COMPOSITOR',
    )

    # Set the node tree/group (Blender 5.0 uses node_group instead of node_tree)
    if hasattr(modifier, 'node_group'):
        # Blender 5.0+
        modifier.node_group = node_tree
    elif hasattr(modifier, 'node_tree'):
        # Blender 4.x
        modifier.node_tree = node_tree
    else:
        raise AttributeError("Compositor modifier has neither 'node_group' nor 'node_tree' attribute")

    # Configure input mapping
    # The modifier automatically maps strip image to first input
    # We need to configure the mask input

    # TODO: Blender 5.0 may have different API for this
    # This is a placeholder for the actual implementation

    return modifier


def get_or_create_blur_node_tree(blur_size: int = 50) -> "bpy.types.NodeTree":
    """
    Get existing or create new blur node tree with specified blur size.
    
    Args:
        blur_size: Blur radius in pixels
        
    Returns:
        The node tree
    """
    import bpy
    
    name = f"FaceMaskBlur_{blur_size}"
    
    if name in bpy.data.node_groups:
        return bpy.data.node_groups[name]
    
    return create_mask_blur_node_tree(name=name, blur_size=blur_size)