vse_toolbox/sequencer_utils.py

# SPDX-License-Identifier: GPL-2.0-or-later

import re
from pathlib import Path
import os

import bpy
from bpy.app.handlers import persistent

from vse_toolbox.bl_utils import get_scene_settings, get_strip_settings, attr_set
from vse_toolbox.file_utils import install_module, parse
from vse_toolbox.constants import SOUND_SUFFIXES
#import multiprocessing
#from multiprocessing.pool import ThreadPool
import subprocess
from collections import defaultdict


def get_render_attributes():
    scn = bpy.context.scene
    return [
        (scn.view_settings, "view_transform", 'Standard'),
        (scn.render.image_settings, "file_format", 'FFMPEG'),
        (scn.render.ffmpeg, "gopsize", 8),
        (scn.render.ffmpeg, "constant_rate_factor", 'HIGH'),
        (scn.render.ffmpeg, "format", 'QUICKTIME'),
        (scn.render.ffmpeg, "audio_codec", 'MP3'),
        (scn.render.ffmpeg, "audio_mixrate", 44100),
        (scn.render.ffmpeg, "audio_bitrate", 128)
    ]

def frame_to_timecode(frame, fps):
    total_seconds = frame / fps

    hours = int(total_seconds // 3600)
    minutes = int((total_seconds % 3600) // 60)
    seconds = int(total_seconds % 60)
    frames = int((total_seconds * fps) % fps)

    timecode = f"{hours:02d}:{minutes:02d}:{seconds:02d}:{frames:02d}"
    return timecode

def new_text_strip(name='Text', channel=0, start=0, end=50, text='Text', font_size=48,
        x=0.5, y=0.5, align_x='CENTER', align_y='CENTER', select=False,
        box_color=None, box_margin=0.005):

    sequences = bpy.context.scene.sequence_editor.sequences
    strip = sequences.new_effect(name, 'TEXT', channel, frame_start=start, frame_end=end)
    strip.select = select
    strip.text = text
    strip.location.x = x
    strip.location.y = y
    strip.align_y = align_y
    strip.align_x = align_x
    strip.channel = channel
    strip.font_size = font_size

    if box_color:
        strip.use_box = True
        strip.box_color = box_color
        strip.box_margin = box_margin

    return strip

def is_strip_at(strip, frame=None):
    if frame is None:
        frame = bpy.context.scene.frame_current

    return (strip.frame_final_start <= frame < strip.frame_final_end)

def get_strips(channel=0, selected_only=False):
    scn = bpy.context.scene
    
    if isinstance(channel, str):
        channel = get_channel_index(channel)

    strips = [s for s in scn.sequence_editor.sequences if s.channel==channel]
    
    if selected_only:
        strips = [s for s in strips if s.select]
    
    return sorted(strips, key=lambda x : x.frame_final_start)

def get_strip_at(channel=0, frame=None):
    strips = get_strips(channel=channel)
    return next((s for s in strips if is_strip_at(s, frame)), None)
    
def get_channel_index(name):
    scn = bpy.context.scene

    channel_id = 0
    channel = scn.sequence_editor.channels.get(name)
    if channel:
        channel_id = scn.sequence_editor.channels.keys().index(name)
    
    return channel_id

def get_channel_name(strip):
    if not strip:
        return

    scn = bpy.context.scene
    return scn.sequence_editor.channels[strip.channel].name

def get_strip_sequence_name(strip):
    sequence_strip = get_strip_at(channel='Sequences', frame=strip.frame_final_start)
    if sequence_strip:
        return sequence_strip.name
    else:
        return 'NoSequence'

def rename_strips(strips, template, increment=10, start_number=0, padding=3, by_sequence=False):
    scn = bpy.context.scene
    settings = get_scene_settings()

    project = settings.active_project
    episode_name = ''
    if settings.active_episode:
        episode_name = settings.active_episode.name

    prev_sequence_name = None
    strip_number = 0
    
    for strip in strips: 
        channel = get_channel_name(strip)
        sequence_name = get_strip_sequence_name(strip)
        format_data = {}
        if channel == 'Shots':
            format_data = parse(project.sequence_template, sequence_name)
        else:
            format_data['sequence'] = str(strip_number*increment + start_number).zfill(padding)

        if (by_sequence and prev_sequence_name and
                sequence_name and sequence_name != prev_sequence_name):
            strip_number = 0

        format_data.update(
            sequence_strip=sequence_name,
            episode=episode_name,
            shot=str(strip_number*increment + start_number).zfill(padding))

        name = template.format(**format_data)

        existing_strip = scn.sequence_editor.sequences_all.get(name)
        if existing_strip: 
            existing_strip.name = f"{name}_tmp"
        
        print(f'Renaming {strip.name} -> {name}')
        strip.name = name
        
        prev_sequence_name = sequence_name    
        strip_number += 1

def set_channels():
    scn = bpy.context.scene
    settings = get_scene_settings()
    items = settings.rna_type.bl_rna.properties['channel'].enum_items 
    
    for i, c in enumerate(items.keys(), start=1):
        scn.sequence_editor.channels[i].name = c.title()

def get_strip_render_path(strip, template):
    scn = bpy.context.scene
    settings = get_scene_settings()
    project = settings.active_project

    suffix = Path(scn.render.frame_path()).suffix

    strip_data = parse(strip.name, template=project.shot_template)

    #print(strip_data)

    index = int(strip_data['index'])
    
    sequence = get_strip_sequence_name(strip)
    render_path = template.format(episode=project.episode_name, sequence=sequence, 
        strip=strip.name, index=index, ext=suffix[1:])

    return Path(os.path.abspath(bpy.path.abspath(render_path)))


'''
# def render_strip_background(blender_path, filepath, start, end, output):
#     cmd = [
#         blender_path, '-b', '--factory-startup', str(tmp_path), '-a',
#         '-s', str(start), '-e', str(end), 
#         '-o', str(output)
#     ]

#     print(cmd)
#     process = subprocess.call(cmd)

def render_strips(strips, template):
    from functools import partial
    scn = bpy.context.scene
    # scene_start = scn.frame_start
    # scene_end = scn.frame_end
    # render_path = scn.render.filepath

    tmp_name = Path(bpy.data.filepath).name if bpy.data.filepath else 'Untitled.blend'
    tmp_path = Path(bpy.app.tempdir, tmp_name)

    bpy.ops.wm.save_as_mainfile(filepath=str(tmp_path), copy=True)


    script_code = dedent(f"""
        import bpy

        for 

        """)

    script_path = Path(bpy.app.tempdir) / 'bundle_library.py'
    script_path.write_text(script_code)

    cmd = [bpy.app.binary_path, tmp_path, '--python', ]

    # nb_threads = min(multiprocessing.cpu_count()-2, 8)
    # print(nb_threads)

    # pool = multiprocessing.Pool(nb_threads)
    
    #arguments = [(bpy.app.binary_path, str(tmp_path), s.frame_final_start, s.frame_final_end-1, str(get_strip_render_path(s, template))) for s in strips]
    #print(arguments)
    #pool.starmap(render_strip_background, arguments)

    # def render_strip_background(index):
    #     cmd = [bpy.app.binary_path, etc]
    #     process = subprocess.Popen(cmd)

    # pool = ThreadPool(nb_threads)

    # for strip in strips:
    #     start = strip.frame_final_start
    #     end = strip.frame_final_end-1
    #     output = str(get_strip_render_path(strip, template))

    #     cmd = [
    #         bpy.app.binary_path, '-b', str(tmp_path),
    #         '-s', str(start), '-e', str(end), '-a',
    #         '-o', str(output)
    #     ]




    for strip in strips:
        #print(render_template, strip.name, path)

        scn.frame_start = strip.frame_final_start
        scn.frame_end = strip.frame_final_end - 1
        
        ## render animatic
        #strip_render_path = render_template.format(strip_name=strip.name, ext=Path(scn.render.frame_path()).suffix)

        scn.render.filepath = str(get_strip_render_path(strip, template))
        #print(scn.render.filepath)
        print(f'Render Strip to {scn.render.filepath}')
        #bpy.ops.render.render(animation=True)  
        bpy.ops.render.opengl(animation=True, sequencer=True)
    
            
    scn.frame_start = scene_start
    scn.frame_end = scene_end
    scn.render.filepath = render_path
'''


def render_sound(strip, output):
    output = os.path.abspath(bpy.path.abspath(output))
    #ext = os.path.splitext(output)[1:]

    scn = bpy.context.scene
    scene_start = scn.frame_start
    scene_end = scn.frame_end
    scene_current = scn.frame_current

    scn.frame_start = strip.frame_final_start
    scn.frame_end = strip.frame_final_end - 1

    print(f'Render Strip to {scn.render.filepath}')
    Path(output).parent.mkdir(exist_ok=True, parents=True)
    bpy.ops.sound.mixdown(filepath=output)#, container=ext.upper()
            
    scn.frame_start = scene_start
    scn.frame_end = scene_end
    scn.frame_current = scene_current

def render_scene(output, attributes=None):
    output = os.path.abspath(bpy.path.abspath(str(output)))

    scn = bpy.context.scene
    render_path = scn.render.filepath

    if attributes is None:
        attributes = []

    attributes += [
        (scn.render, "filepath", output),
    ]

    Path(output).parent.mkdir(exist_ok=True, parents=True)
    with attr_set(attributes):
        print(f'Render Strip to {scn.render.filepath}')
        bpy.ops.render.opengl(animation=True, sequencer=True)

def render_strip(strip, output, attributes=None):
    output = os.path.abspath(bpy.path.abspath(str(output)))

    scn = bpy.context.scene

    if attributes is None:
        attributes = []

    attributes += [
        (scn, "frame_start", strip.frame_final_start),
        (scn, "frame_end", strip.frame_final_end - 1),
        (scn, "frame_current"),
        (scn.render, "filepath", output),
    ]

    Path(output).parent.mkdir(exist_ok=True, parents=True)
    with attr_set(attributes):
        print(f'Render Strip to {scn.render.filepath}')
        bpy.ops.render.opengl(animation=True, sequencer=True)
            
def import_edit(filepath, adapter="cmx_3600", channel='Shots', match_by='name'):
    opentimelineio = install_module('opentimelineio')

    from opentimelineio.schema import Clip

    scn = bpy.context.scene
    sequencer = scn.sequence_editor.sequences
    strips = get_strips(channel='Shots')
    shot_channel = get_channel_index('Shots')

    # Move all strips on an empty channel
    empty_channel = max((s.channel for s in sequencer), default=0) + 1
    for s in strips:
        s.channel = empty_channel

    edl = Path(filepath)
    try:
        timeline = opentimelineio.adapters.read_from_file(
            str(edl), adapter, rate=scn.render.fps, ignore_timecode_mismatch=True)
    except:
        print("[>.] read_from_file Failed. Using read_from_string method.")
        data = edl.read_text(encoding='latin-1')
        timeline = opentimelineio.adapters.read_from_string(
            data, adapter, rate=scn.render.fps, ignore_timecode_mismatch=True)

    scn.frame_start = (
        0 if timeline.global_start_time is None else timeline.global_start_time
    )

    # Get all video clips only
    clips = []
    for track in timeline.tracks:
        for child in track.each_child(shallow_search=True):
            
            # FIXME Exclude Gaps for now. Gaps are Transitions, Blank Spaces...
            if not isinstance(child, Clip):
                continue
            
            # FIXME Exclude Audio for now
            if any(child.name.lower().endswith(ext) for ext in SOUND_SUFFIXES):
                continue

            clips.append(child)

    clips.sort(key=lambda x: x.range_in_parent().start_time)


    for i, clip in enumerate(clips):
            
        frame_start = opentimelineio.opentime.to_frames(
            clip.range_in_parent().start_time)

        frame_end = frame_start + opentimelineio.opentime.to_frames(
            clip.range_in_parent().duration)
        
        strip = None

        if match_by.lower() == 'name':
            strip = next((s for s in strips if s.vsetb_strip_settings.source_name == clip.name), None)
        elif match_by.lower() == 'index':
            try:
                strip = strips[i]
            except IndexError:
                print(f'No strip found for {clip.name}')

        if strip:
            if frame_start != strip.frame_final_start or frame_end !=strip.frame_final_end:
                print(f'The strip {strip.name} is updated with new range')
                #self.report({'INFO'}, f'The strip {strip.name} is updated with new range')

            strip.channel = shot_channel
            strip.frame_final_start = frame_start
            strip.frame_final_end = frame_end

        else:
            print('Create a new strip')
            strip = sequencer.new_effect(
                name=clip.name,
                type='COLOR',
                channel=shot_channel,
                frame_start=frame_start,
                frame_end=frame_end,
            )

        strip.blend_alpha = 0.0
        strip.select = False
        strip.vsetb_strip_settings.source_name = clip.name

        #except Exception as e:
        #    print('e: ', e)
        #    continue

    scn.frame_end = frame_end-1
    return timeline

def import_movie(filepath, frame_start=None, frame_end=None):
    scn = bpy.context.scene

    if frame_start is None:
        frame_start = scn.frame_start

    res_x = scn.render.resolution_x
    res_y = scn.render.resolution_y

    if bpy.data.is_saved:
        relpath = Path(bpy.path.relpath(str(filepath)))
        if len(relpath.as_posix()) < len(filepath.as_posix()):
            filepath = relpath
            
    strip = scn.sequence_editor.sequences.new_movie(
        name=Path(filepath).stem,
        filepath=str(filepath),
        channel=get_channel_index('Movie'),
        frame_start=frame_start
    )

    elem = strip.strip_elem_from_frame(scn.frame_current)
    src_width, src_height = elem.orig_width, elem.orig_height

    if src_width != res_x:
        strip.transform.scale_x = (res_x / src_width)
    if src_height != res_y:
        strip.transform.scale_y = (res_y / src_height)

    if frame_end is not None:
        strip.frame_final_end = frame_end

    return strip

def scale_clip_to_fit(strip):
    scn = bpy.context.scene
    res = scn.render.resolution_x, scn.render.resolution_y
    strip_res = strip.elements[0].orig_width, strip.elements[0].orig_height

    #print(strip.name, strip_res, width, height)

    strip.transform.scale_x = res[0] / strip_res[0]
    strip.transform.scale_y = strip.transform.scale_x

def import_sound(filepath, frame_start=None, frame_end=None):
    scn = bpy.context.scene

    if frame_start is None:
        frame_start = scn.frame_start

    if bpy.data.is_saved:
        relpath = Path(bpy.path.relpath(str(filepath)))
        if len(relpath.as_posix()) < len(filepath.as_posix()):
            filepath = relpath

    strip = scn.sequence_editor.sequences.new_sound(
        name=f'{filepath.stem} Audio',
        filepath=str(filepath),
        channel=get_channel_index('Audio'),
        frame_start=frame_start
    )

    if frame_end is not None:
        strip.frame_final_end = frame_end

    #strip.show_waveform = True if strip.frame_final_duration < 10000 else False

    return strip

def get_empty_channel():
    return max(s.channel for s in bpy.context.scene.sequence_editor.sequences) + 1

def clean_sequencer(edit=False, movie=False, sound=False):
    scn = bpy.context.scene
    sequences = []

    if edit:
        sequences.extend(get_strips('Shots'))
    if movie:
        sequences.extend(get_strips('Movie'))
    if sound:
        sequences.extend(get_strips('Audio'))

    for sequence in sequences:
        scn.sequence_editor.sequences.remove(sequence)

@persistent
def set_active_strip(scene):
    scn = bpy.context.scene
    if not scn.sequence_editor:
        return

    if not get_scene_settings().auto_select_strip:
        return
    
    
    #scene.sequence_editor.active_strip = None
    
    shot_strip = get_strip_at('Shots')
    if shot_strip:
        bpy.ops.sequencer.select_all(action="DESELECT")
        shot_strip.select = True
        scene.sequence_editor.active_strip = shot_strip

@persistent
def update_text_strips(scene):
    if not scene.sequence_editor or not scene.sequence_editor.sequences_all:
        return

    #print("update_text_strips")
    settings = get_scene_settings()
    project = settings.active_project
    episode = settings.active_episode


    class MissingKey(dict):
        def __missing__(self, key):
            return '{' + key + '}'


    scn = bpy.context.scene
    if not scn.sequence_editor:
        return

    format_data = {
        'scene': scene, 
        'project': '...',
        'episode': '...',
        'sequence': '...',
        'sequence_strip': '...', 
        'shot': '...', 
        'shot_strip': '...', 
        'shot_frame': 0,
        'shot_duration': 0, 
        'shot_start': 0, 
        'shot_end': 0,
        'timecode': ''
    }

    shot_strip = get_strip_at('Shots', frame=scene.frame_current)
    if shot_strip:
        format_data.update({
            'project': project.name,
            'episode': episode.name if episode else '',
            'shot_duration': shot_strip.frame_final_duration,
            'shot_frame': scene.frame_current - shot_strip.frame_final_start + 1,
            'shot_start': shot_strip.frame_final_start,
            'shot_end': shot_strip.frame_final_end,
            'timecode' : frame_to_timecode(scene.frame_current, scene.render.fps)
        })
        format_data.update(shot_strip.vsetb_strip_settings.format_data)

    for strip in scene.sequence_editor.sequences_all:
        if not strip.type == 'TEXT':
            continue

        if not is_strip_at(strip):
            continue

        if '{' in strip.text:
            strip['text_pattern'] = strip.text
        
        if 'text_pattern' in strip.keys():
            strip.text = strip['text_pattern'].format_map(MissingKey(**format_data))


def create_shot_strip(name, start, end):

    shot_strip = bpy.context.scene.sequence_editor.sequences.new_effect(
        name,
        'COLOR',
        get_channel_index('Shots'),
        frame_start=start,
        frame_end=end
    )
    shot_strip.blend_alpha = 0
    shot_strip.color = (0.5, 0.5, 0.5)

    return shot_strip