in03 / squawk

Automatic subtitles for DaVinci Resolve with OpenAI Whisper
MIT License
37 stars 3 forks source link

Allow importing to custom nested folder path #4

Closed github-actions[bot] closed 2 years ago

github-actions[bot] commented 2 years ago

Allow importing to custom nested folder path

https://github.com/in03/squawk/blob/c88c70eb80ccf2864fe21c924363a2abbff3d44b/squawk/app/transcribe.py#L71


import logging
import os
import time

from squawk.app.resolve import ResolveObjects
from squawk.settings.manager import SettingsManager
from squawk.utils import core

import whisper
from whisper import utils as whisper_utils

logger = logging.getLogger(__name__)

settings = SettingsManager()

class Transcribe:
    def __init__(self, translate: bool = False):

        self.translate = translate
        self.ro = ResolveObjects()
        self.media_file = ""

    def render(self):
        """
        Render active timeline in Resolve
        """

        logger.info("[yellow]Starting Resolve render")
        core.notify("Squawk", "Starting Resolve render")

        self.media_file = self.ro.render_timeline(
            settings["paths"]["working_dir"], settings["render"]["render_preset"]
        )
        logger.info("[green]Finished rendering")
        core.notify("Squawk", "Finished rendering")

    def tts(self):

        if not self.media_file:
            raise ValueError(f"Media file has not been defined! Run render first.")

        core.notify("Squawk", "Starting transcription")

        model = whisper.load_model(settings["whisper"]["model"])
        start_time = time.time()

        if self.translate:
            result = model.transcribe(self.media_file, task="translate")
        else:
            result = model.transcribe(self.media_file)

        core.notify(
            "Squawk", f"Processing finished after {time.time() - start_time} seconds"
        )

        # save SRT file to previously selected target_dir
        self.srt_path = os.path.join(
            settings["whisper"]["model"], (os.path.basename(self.media_file) + ".srt")
        )
        with open(self.srt_path, "w", encoding="utf-8") as srt:
            whisper_utils.write_srt(result["segments"], file=srt)

    def import_srt(self):
        """
        Import SRT file into Resolve
        """

        logger.info(f"[cyan]Importing SRT file: '{self.srt_path}'")

        # TODO: Allow importing to custom nested folder path
        root_folder = self.ro.media_pool.GetRootFolder()
        self.ro.media_pool.SetCurrentFolder(root_folder)
        self.ro.media_pool.ImportMedia(self.media_file)

    def run(self):

        self.render
        self.tts
        self.import_srt

dd29a9313612145e54df81df2593101cd6a30c23

github-actions[bot] commented 2 years ago

Closed in da7e7d0e43b81beb67605be3a29734ae5803d5ba