gp_render/OP_export_to_ae.py

from encodings import utf_8
import bpy
import json

# import bpy_extras
from bpy_extras.object_utils import world_to_camera_view # as cam_space
from bpy_extras.io_utils import ExportHelper
from mathutils import Vector
from pathlib import Path
from . import fn

from bpy.props import (
    StringProperty,
    BoolProperty,
    FloatProperty,
    EnumProperty,
    CollectionProperty,
)

from bpy.app.translations import pgettext_data as data_

'''
def Export_AE_2d_position_json_data():
    scn = bpy.context.scene
    cam = scn.objects.get('anim_cam')
    if not cam:
        print('Active camera not "anim_cam"')
        cam = scn.camera

    rx = scn.render
    rx, ry = rd.resolution_x, rd.resolution_y

    targets = [o for o in bpy.context.selected_objects if o.type != 'CAMERA']

    for ob in targets:
        # get an idea of the scale relative to image res ? (too complex since it's not even the right resolution...)
        pos = world_to_camera_view(scn, scn.objects['anim_cam'], ob.matrix_world.to_translation())[:-1]
        pix_pos = Vector((pos[0]*rx, (1-pos[1])*ry))
'''

def export_ae_transforms(directory, selection=None, camera=None, exposition=True, prefix='ae_', suffix='', fr=False, export_format='txt', export_cam=True):
    """
    Export After Effects transform data for selected objects and camera in Blender.

    This function exports keyframe data for position, scale, and optional exposition for
    selected objects and the camera. The data can be exported in TXT and/or JSON formats,
    which can be imported into Adobe After Effects.

    Parameters:
    directory (str): The directory path where the exported files will be saved.
    selection (list, optional): List of Blender objects or pose bones to export. Defaults to None.
    camera (bpy.types.Object, optional): The camera object to use for calculations. Defaults to the scene's active camera.
    exposition (bool, optional): Whether to export exposition data. Defaults to True.
    prefix (str, optional): Prefix for the exported filenames. Defaults to 'ae_'.
    suffix (str, optional): Suffix for the exported filenames. Defaults to ''.
    fr (bool, optional): If True, uses French labels for exposition data. Defaults to False.
    export_format (str, optional): The format(s) to export. Can be 'txt', 'json', or 'txt,json'. Defaults to 'txt'.

    Example usage:
    export_ae_transforms('/path/to/export', selection=bpy.context.selected_objects, export_format='txt,json')
    """

    scn = bpy.context.scene
    res = Vector((scn.render.resolution_x, scn.render.resolution_y))

    def get_keyframe_data(coords_2d, exposition=False):
        frame_start, org_2d, _ = coords_2d[0]
        org_distance = org_2d[-1]
        
        keyframe_data = {'position': [], 'scale': []}

        for frame, co_2d, is_animated in coords_2d:
            position = Vector((co_2d[0], 1-co_2d[1]))
            position *= res # Multiply by resolution to get pixels value
            keyframe_data['position'].append((frame-frame_start, (round(position[0]), round(position[1]))))
                                    
            scale = (org_distance / co_2d[-1]) * 100            
            keyframe_data['scale'].append((frame-frame_start, (round(scale, 5),)*3))
            
            if is_animated:
                keyframe_data.setdefault('exposition', []).append(frame-frame_start)

        return keyframe_data

    def get_keyframe_data_txt(keyframe_data, exposition=True):
        keyframe_data_txt = (
            'Adobe After Effects 8.0 Keyframe Data\n\n'
            f'\tUnits Per Second\t{scn.render.fps}\n'
            f'\tSource Width\t{scn.render.resolution_x}\n'
            f'\tSource Height\t{scn.render.resolution_y}\n'
            '\tSource Pixel Aspect Ratio\t1\n'
            '\tComp Pixel Aspect Ratio\t1\n\n'
        )

        # Position
        keyframe_data_txt += (    
            '\nTransform\tPosition\n'
            '\tFrame\tX pixels\tY pixels\tZ pixels\t\n'
        )
        for frame, position in keyframe_data['position']:
            keyframe_data_txt += f'\t{frame}\t{position[0]}\t{position[1]}\t0\t\n'

        # Scale
        keyframe_data_txt += (
            '\nTransform\tScale\n'
            '\tFrame\tX percent\tY percent\tZ percent\t\n'
        )
        for frame, scale in keyframe_data['scale']:
            keyframe_data_txt += f'\t{frame}\t{scale[0]}\t{scale[1]}\t{scale[2]}\t\n'

        if exposition:
            if fr:
                keyframe_data_txt += '\nEffects\tParamètre case #1\tCase #2\t\n'
            else:
                keyframe_data_txt += '\nEffects\tCheckbox Control #1\tCheckbox #2\t\n'
            
            keyframe_data_txt += '\tFrame\t\t\n'

            for i, frame in enumerate(keyframe_data['exposition']):
                keyframe_data_txt += f'\t{frame}\t{i%2}\t\n'

        keyframe_data_txt += '\nEnd of Keyframe Data\n' 
        
        return keyframe_data_txt


    if selection is None:
        selection = []

    if camera is None:
        camera = scn.camera
    
    use_simplify = scn.render.use_simplify
    simplify_subdivision = scn.render.simplify_subdivision
    scn.render.use_simplify = True
    scn.render.simplify_subdivision = 0
    
    # Get the camera at the first frame for reference
    scn.frame_set(scn.frame_start)
    cam_org_3d = camera.matrix_world @ Vector((0, 0, -100))
        
    camera_coords = []
    ob_coords = {}
    animation = {}
    for i in range(scn.frame_start, scn.frame_end + 1):
        scn.frame_set(i)
        
        if camera is not None:
            co_2d = world_to_camera_view(scn, camera, cam_org_3d)
            camera_coords.append((i, co_2d, False))
            
        for item in selection:
            if isinstance(item, bpy.types.PoseBone):
                co_3d = item.id_data.matrix_world @ item.matrix.to_translation()
                name = f'{item.id_data.name}_{item.name}'
                rig_transforms = [tuple(round(x, 5) for v in b.matrix for x in v) for b in item.id_data.pose.bones]
                is_animated = rig_transforms != animation.get(name)
                animation[name] = rig_transforms
            elif item.type == 'GPENCIL': # isinstance(item.data, bpy.types.GreasePencil):
                name = item.name
                co_3d = item.matrix_world.to_translation()
                ## Check if there is a GP-scene at scene-frame (i) on any visible layer
                is_animated = i in [f.frame_number for l in item.data.layers for f in l.frames if not l.hide]
                animation[name] = co_3d
            else:
                name = item.name
                co_3d = item.matrix_world.to_translation()
                is_animated = co_3d != animation.get(name)
                animation[name] = co_3d

            co_2d = world_to_camera_view(scn, camera, co_3d)
            ob_coords.setdefault(name, []).append((i, co_2d, is_animated))
    
    cam_key_frame_data = get_keyframe_data(camera_coords)
    objs_keyframe_data = {name: get_keyframe_data(co, exposition=exposition) for name, co in ob_coords.items()} 

    if 'json' in export_format or export_format == 'json':
        keyframe_data = {
                "units_per_second": scn.render.fps,
                "source_width": scn.render.resolution_x,
                "source_height": scn.render.resolution_y,
                "source_pixel_aspect_ratio": 1,
                "comp_pixel_aspect_ratio": 1,
                'objects': objs_keyframe_data
            }
        if export_cam: #camera is not None:
            keyframe_data['camera'] = cam_key_frame_data
                
        Path(directory, f'{prefix}keyframe_data{suffix}.json').write_text(json.dumps(keyframe_data, indent=4), encoding='utf-8')
    
    if 'txt' in export_format or export_format == 'txt':
        for name, keyframe_data in objs_keyframe_data.items():
            name = name.replace('.', '_')
            keyframe_data_txt = get_keyframe_data_txt(keyframe_data, exposition=exposition)
            Path(directory, f'{prefix}{name}{suffix}.txt').write_text(keyframe_data_txt, newline='\r\n')
        
        if export_cam: #camera is not None:
            keyframe_data_txt = get_keyframe_data_txt(cam_key_frame_data, exposition=False)
            Path(directory, f'{prefix}{camera.name}{suffix}.txt').write_text(keyframe_data_txt, newline='\r\n')
    
    scn.render.use_simplify = use_simplify
    scn.render.simplify_subdivision = simplify_subdivision


## Export operator without export_helper
class GPEXP_OT_export_anim_to_ae(bpy.types.Operator):
    bl_idname = "gp.export_anim_to_ae"
    bl_label = "Export AE Files"
    bl_description = "Export the animation to After Effects, 2D transform of objects, camera\
        \nand/or exposition (including greasepencil frames)"
    bl_options = {"REGISTER"}

    # filter_glob: StringProperty(default='*.*', options={'HIDDEN'})# *.jpeg;*.png;*.tif;*.tiff;*.bmp

    # filename_ext = ''

    # filepath : StringProperty(
    #     name="File Path",
    #     description="File path used for export", 
    #     maxlen= 1024)

    ## Only need directory
    directory : StringProperty(
        name="File Path",
        description="File path used for export", 
        maxlen= 1024,
        subtype='DIR_PATH'
    )
    prefix : StringProperty(
        name="Prefix",
        default='ae_',
        description="Prefix name for exported txt and json files", 
        maxlen= 1024,
    )

    use_selection: BoolProperty(
        name="Selected Objects",
        description="Export selected and visible objects only",
        default=True,
    )
    use_visible: BoolProperty(
        name='Visible Objects',
        description='Export visible objects only',
        default=False
    )
    use_active_collection: BoolProperty(
        name="Active Collection",
        description="Export only objects from the active collection (and its children)",
        default=False,
    )
    object_types: EnumProperty(
        name="Object Types",
        options={'ENUM_FLAG'},
        items=(('EMPTY', "Empty", ""),
               ('CAMERA', "Camera", ""),
               ('LIGHT', "Lamp", ""),
               ('ARMATURE', "Armature", "WARNING: not supported in dupli/group instances"),
               ('MESH', "Mesh", ""),
               ('GPENCIL', "Grease Pencil", ""),
               ('OTHER', "Other", "Other geometry types, like curve, metaball, etc. (converted to meshes)"),
               ),
        description="Which kind of object to export",
        default={'EMPTY', 'CAMERA', 'LIGHT', 'ARMATURE', 'MESH', 'GPENCIL', 'OTHER'},
    )
    exposition: BoolProperty(
        name='Exposition',
        description='Export the exposition of the keys',
        default=True
    )
    # use_object_keys: BoolProperty(
    #     name='Object Keys',
    #     description='Consider object transform keys for animated exposition',
    #     default=True
    # )
    use_active_camera: BoolProperty(
        name='Active Camera',
        description='Export active camera keys',
        default=True
    )
    data_lang: EnumProperty(
        name="AE Language",
        items=(('FR', "French", ""),
               ('EN', "English", ""),
               ),
        description="Clipboard keyframe data language",
        default='FR',
    )
    file_format: EnumProperty(
        name="File Type",
        options={'ENUM_FLAG'},
        items=(('txt', "txt", ""),
               ('json', "json", ""),
               ),
        description="File format to export (possible to select multiple choices with Shift + Click)",
        default={'txt'},
    )

    def invoke(self, context, _event):
        if not self.directory:
            blend_filepath = context.blend_data.filepath
            if blend_filepath:
                dest_folder = Path(blend_filepath).parent
                
                ## If pre-enter a specific subfolder exists
                ## (Commented, should be used with a project environment variable)
                # output_folder = dest_folder / 'render'
                # if output_folder.exists() and output_folder.is_dir():
                #     dest_folder = output_folder

                self.directory = str(dest_folder)
                
        context.window_manager.fileselect_add(self)
        return {'RUNNING_MODAL'}

    def draw(self, context):
        layout = self.layout
        layout.use_property_split = True
        layout.use_property_decorate = False  # No animation.

        # Are we inside the File browser
        is_file_browser = context.space_data.type == 'FILE_BROWSER'

        export_main(layout, self, is_file_browser)
        export_panel_format(layout, self, is_file_browser)
        export_panel_include(layout, self, is_file_browser)

    def execute(self, context):
        objects_selection = get_object_selection(use_selection=self.use_selection,
                             use_visible=self.use_visible,
                             use_active_collection=self.use_active_collection,
                             object_types=self.object_types
                             )

        ## Ensure output path is directory
        output_path = Path(self.directory)
        if not output_path.is_dir():
            output_path = output_path.parent

        print('Output directory: ', output_path)
        
        cam = None
        if context.scene.camera:
            cam = context.scene.camera
        
        if cam and cam in objects_selection:
            ## Remove active camera from objects list
            objects_selection.pop(objects_selection.index(cam))

        print('Export AE transform from objects:')
        for o in objects_selection:
            print('- ', o.name)

        export_ae_transforms(directory=output_path,
                             selection=objects_selection,
                             camera=cam,
                             exposition=self.exposition,
                             prefix=self.prefix,
                             fr=self.data_lang == 'FR',
                             export_format=self.file_format,
                             export_cam=self.use_active_camera)

        self.report({'INFO'}, f'File(s) saved in folder: {output_path}')
        return {"FINISHED"}


def get_object_selection(use_selection=False, use_visible=False, use_active_collection=False, object_types=None):
    context = bpy.context
    ## determine selection based on filters4
    source_collection = None
    if use_active_collection:
        source_collection = context.view_layer.active_layer_collection.collection

    if source_collection:
        if use_selection:
            ctx_objects = tuple(obj for obj in source_collection.all_objects if obj.select_get())
        else:
            ctx_objects = source_collection.all_objects
    else:
        if use_selection:
            ctx_objects = context.selected_objects
        else:
            ctx_objects = context.view_layer.objects
    
    if use_visible:
        ctx_objects = tuple(obj for obj in ctx_objects if obj.visible_get())

    ## Filter by object type
    if object_types is None:
        object_types = {'EMPTY', 'CAMERA', 'LIGHT', 'ARMATURE', 'MESH', 'GPENCIL', 'OTHER'}

    ctx_objects = [obj for obj in ctx_objects if obj.type in object_types]
    ## Constant for compatible other object type
    # BLENDER_OTHER_OBJECT_TYPES = {'CURVE', 'SURFACE', 'FONT', 'META'}

    if 'OTHER' in object_types:
        ## Any object that is not in proposed list
        ctx_objects += [obj for obj in ctx_objects if obj.type not in object_types]
    
    return ctx_objects


def export_main(layout, operator, is_file_browser):
    col = layout.column()
    col.prop(operator, 'exposition')
    col.prop(operator, 'use_active_camera')
    # col.prop(operator, 'use_object_keys')

def export_panel_format(layout, operator, is_file_browser):
    header, body = layout.panel("AE_export_format", default_closed=False)
    header.label(text="Format")
    if body:
        col = body.column()
        col.prop(operator, 'prefix')
        ## Format (language and file)

        col.row().prop(operator, 'data_lang', expand=True)
        if is_file_browser:
            col.column().prop(operator, 'file_format')

def export_panel_include(layout, operator, is_file_browser):
    header, body = layout.panel("AE_export_include", default_closed=False)
    header.label(text="Include")
    if body:
        sublayout = body.column(heading="Limit to")
        # if is_file_browser:
        sublayout.prop(operator, "use_selection")
        sublayout.prop(operator, "use_visible")
        sublayout.prop(operator, "use_active_collection")

        body.column().prop(operator, "object_types")
        # body.prop(operator, "use_custom_props")

### --- Old functions

# Unused (old func that might not be usefull at all...)
def correct_shift(vec, cam):
    resX = bpy.context.scene.render.resolution_x
    resY = bpy.context.scene.render.resolution_y
    ratio = resX/resY
    shiftX = 2*cam.data.shift_x
    shiftY = 2*cam.data.shift_y

    if ratio<1:
        return vec - Vector((shiftX*(1/ratio), shiftY, 0))
    elif ratio>1:
        return vec - Vector((shiftX, shiftY*ratio, 0))
    else:
        return vec - Vector((shiftX, shiftY, 0))

def export_AE_objects_position_keys():
    '''Export keys as paperclip to paste in after'''

    scn = bpy.context.scene
    result = {}
    print(f'Exporting 2d position (scene range: {scn.frame_start} - {scn.frame_end})')
    for fr in range(scn.frame_start,scn.frame_end + 1):

        print(f'frame: {fr}')
        scn.frame_set(fr)

        for o in bpy.context.selected_objects:
            if not result.get(o.name):
                result[o.name] = []
            proj2d = world_to_camera_view(scn, scn.camera, o.matrix_world.to_translation()) # + Vector((.5,.5,0))

            # proj2d = correct_shift(proj2d, scn.camera) # needed ?
            x = (proj2d[0]) * scn.render.resolution_x
            y = -(proj2d[1]) * scn.render.resolution_y + scn.render.resolution_y

            result[o.name].append((fr,x,y))

    for name,value in result.items():

        txt = fn.get_ae_keyframe_clipboard_header(scn)

        for v in value:
            txt += '\t%s\t%s\t%s\t0\t\n'%(v[0],v[1],v[2]) # add 0 for Z (probably not needed)

        txt += '\n\nEnd of Keyframe Data\n' # keyframe txt footer

        blend = Path(bpy.data.filepath)
        keyfile = blend.parent / 'render' / f'pos_{name}.txt'
        keyfile.parent.mkdir(parents=False, exist_ok=True)

        print(f'exporting keys for {name} at {keyfile}')

        ## save forcing CRLF terminator (DOS style, damn windows)
        ## in case it's exported from linux
        with open(keyfile, 'w', newline='\r\n') as fd:
            fd.write(txt)


class GPEXP_OT_export_keys_to_ae(bpy.types.Operator):
    bl_idname = "gp.export_keys_to_ae"
    bl_label = "Export 2D Position To AE"
    bl_description = "Export selected objects positions as text file containing key paperclip for AfterEffects layers"
    bl_options = {"REGISTER"}

    @classmethod
    def poll(cls, context):
        return context.selected_objects

    def execute(self, context):
        export_AE_objects_position_keys()
        return {"FINISHED"}


def export_anim_cam_position(camera=None, context=None):
    context = context or bpy.context
    scn = context.scene

    camera = camera or bpy.data.objects.get('anim_cam')
    if not camera:
        return 'Abort: No "anim_cam" found!'

    text = fn.get_ae_keyframe_clipboard_header(scn)
    for i in range(scn.frame_start, scn.frame_end + 1):
        scn.frame_set(i)
        center = fn.get_cam_frame_center_world(camera)
        coord = fn.get_coord_in_cam_space(scn, scn.camera, center, ae=True)
        text += f'\t{i}\t{coord[0]}\t{coord[1]}\t\n'

    text += '\n\nEnd of Keyframe Data\n' # Ae Frame ending

    blend = Path(bpy.data.filepath)
    keyfile = blend.parent / 'render' / f'anim_cam_pos.txt'
    keyfile.parent.mkdir(parents=False, exist_ok=True)

    print(f'Exporting anim cam positions keys at: {keyfile}')
    with open(keyfile, 'w') as fd:
        fd.write(text)

class GPEXP_OT_export_cam_keys_to_ae(bpy.types.Operator):
    bl_idname = "gp.export_cam_keys_to_ae"
    bl_label = "Export Camera 2D Position To AE"
    bl_description = "Export anim cam positions as text file containing key paperclip for AfterEffects layers"
    bl_options = {"REGISTER"}

    @classmethod
    def poll(cls, context):
        return context.selected_objects

    def execute(self, context):
        cam = None
        if context.object and context.object.type == 'CAMERA' and context.object != context.scene.camera:
            cam = context.object

        err = export_anim_cam_position(camera=cam, context=context)
        if err:
            self.report({'ERROR'}, err)
            return {"CANCELLED"}
        return {"FINISHED"}

class GPEXP_OT_fix_overscan_shift(bpy.types.Operator):
    bl_idname = "gp.fix_overscan_shift"
    bl_label = "Fix Cam Shift Value With Overscan"
    bl_description = "(Gp render operator) change shift values to re-center overscan"
    bl_options = {"REGISTER"}

    # @classmethod
    # def poll(cls, context):
    #     return context.object

    init_rx : bpy.props.IntProperty(name='pre-overscan res x')
    init_ry : bpy.props.IntProperty(name='pre-overscan res y')

    def invoke(self, context, event):
        # if not context.object:
        #     self.report({'ERROR'}, 'No object selected')
        #     return {'CANCELLED'}

        self.use_selection = False
        if context.active_object and context.active_object.type == 'CAMERA' and context.active_object != context.scene.camera:
            self.use_selection = True
            self.cam_ob = context.active_object
        else:
            self.cam_ob = context.scene.camera

        self.init_rx = context.scene.render.resolution_x
        self.init_ry = context.scene.render.resolution_y

        return context.window_manager.invoke_props_dialog(self)

    def draw(self, context):
        layout = self.layout

        if self.use_selection:
            col = layout.column()
            col.label(text=f'Camera "{self.cam_ob.name}" selected', icon='INFO')
            col.label(text='Change in shifts will apply on this one', icon='BLANK1')

        col = layout.column()
        col.label(text='Overscan res (current)')
        col.label(text=f'{context.scene.render.resolution_x} x {context.scene.render.resolution_y}')

        col = layout.column()
        col.label(text='Enter Initial (pre-overscan) resolution:')
        row = col.row(align=True)
        row.prop(self, 'init_rx', text='')
        row.prop(self, 'init_ry', text='')

    def execute(self, context):
        cam = self.cam_ob.data

        ratio_x = self.init_rx / context.scene.render.resolution_x
        ratio_y = self.init_ry / context.scene.render.resolution_y

        if ratio_x == 1 and ratio_y == 1:
            self.report({'ERROR'}, 'Same init and overscan resolution, nothing to change')
            return {'CANCELLED'}

        if ratio_x != 1:
            if fn.has_keyframe(cam, 'shift_x'):
                fcu = cam.animation_data.action.fcurves.find('shift_x')
                for k in fcu.keyframe_points:
                    k.co.y = k.co.y * ratio_x
            else:
                if cam.shift_x != 1:
                    cam.shift_x = cam.shift_x * ratio_x

        if ratio_y != 1:
            if fn.has_keyframe(cam, 'shift_y'):
                fcu = cam.animation_data.action.fcurves.find('shift_y')
                for k in fcu.keyframe_points:
                    k.co.y = k.co.y * ratio_y
            else:
                if cam.shift_y != 1:
                    cam.shift_y = cam.shift_y * ratio_y

        return {"FINISHED"}


# ui panel
class GPEXP_PT_extra_gprender_func(bpy.types.Panel):
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'
    bl_context = "output"
    bl_label = "GP Render Extras"
    bl_parent_id = "RENDER_PT_format" if (3, 0, 0) <= bpy.app.version else "RENDER_PT_dimensions"
    bl_options = {'DEFAULT_CLOSED'}
    # COMPAT_ENGINES = {'CYCLES', 'BLENDER_EEVEE', 'BLENDER_WORKBENCH'}

    # def draw_header(self, context):
    #     overscan = context.scene.camera_overscan
    #     self.layout.prop(overscan, "RO_Activate", text="")

    def draw(self, context):
        layout = self.layout
        col = layout.column()
        col.operator("gp.fix_overscan_shift")
        col.operator("gp.export_keys_to_ae")
        col.operator("gp.export_cam_keys_to_ae")

# def overcan_shift_fix_ui(self, context):
#     layout = self.layout
#     layout.operator("gp.fix_overscan_shift")

def export_ae_anim_menu(self, context):
    row = self.layout.row(align=False)
    row.operator('gp.export_anim_to_ae', text='After Effects Keyframe Data (.txt/.json)')
    # row.operator('gp.export_anim_to_ae', text='', icon='COPYDOWN')


classes=(
GPEXP_OT_export_keys_to_ae,
GPEXP_OT_export_cam_keys_to_ae,
GPEXP_OT_fix_overscan_shift,
GPEXP_PT_extra_gprender_func,
GPEXP_OT_export_anim_to_ae,
)

def register():
    for cls in classes:
        bpy.utils.register_class(cls)

    bpy.types.TOPBAR_MT_file_export.append(export_ae_anim_menu)
    # if hasattr(bpy.types, 'RENDER_PT_overscan'):
    #     bpy.types.RENDER_PT_overscan.append(overcan_shift_fix_ui)


def unregister():
    # if hasattr(bpy.types, 'RENDER_PT_overscan'):
    #     bpy.types.RENDER_PT_overscan.remove(overcan_shift_fix_ui)

    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)
    bpy.types.TOPBAR_MT_file_export.remove(export_ae_anim_menu)