main/adjust.py

from ui.ui_component import *
from tool.file_name import *
from pydub import AudioSegment
from tool.json_tool import load_groups_json
import os
import gradio as gr
import tool.text_file_tool as text_file_tool


min_duration_ms = 10000
current_pos = 0

speaker_to_name = {"Speaker": "Name"}
speaker_to_sample_file = {"Speaker": "File"}

sample_groups = []


def get_current_sample_file():
    global speaker_to_name
    global speaker_to_sample_file

    print(f"GET CURRENT speaker_to_name: {speaker_to_name}")
    print(f"GET CURRENT speaker_to_sample: {speaker_to_sample_file}")

    name = list(speaker_to_name.values())[current_pos]
    sample = list(speaker_to_sample_file.values())[current_pos]
    print(f"CURRENT: {name} {sample}")

    return [name, sample]


def get_sample_file_for_speakers():
    global sample_groups
    global speaker_to_name
    global speaker_to_sample_file

    sample_groups, _ = load_groups_json()
    print(f"SAMPLE GROUP: {sample_groups}")
    speaker_to_name = {}
    speaker_to_sample_file = {}

    for speaker in sample_groups:
        for suffix in range(1, 100):
            file_path = f"{speaker}-{suffix}.wav"
            if os.path.exists(file_path):
                audio_segment = AudioSegment.from_file(file_path)
                if len(audio_segment) > min_duration_ms:
                    print(f"Found file: {file_path}")
                    print(f"Duration: {len(audio_segment) / 1000} seconds")
                    break

        temp_file_name = f"{speaker}-sample.wav"
        audio_segment[:10 * 1000].export(temp_file_name, format="wav")
        speaker_to_sample_file[speaker] = temp_file_name
        speaker_to_name[speaker] = speaker
        print(f"speaker_to_name: {speaker_to_name}")
        print(f"speaker_to_sample: {speaker_to_sample_file}")

    return get_current_sample_file()


def prepare_output(input_file):
    if input_file is None or not os.path.exists(input_file):
        raise gr.Error(current_ui_lang["input_video_warning"])

    speakers = get_sample_file_for_speakers()

    download_video_update = download_video_subtitle_button.update(
        interactive=True)
    adjust_speaker_update = adjust_speaker.update(
        speakers[0], interactive=True)
    adjust_audio_update = adjust_audio.update(
        speakers[1], interactive=True)
    prev_button_update = prev_button.update(
        interactive=True)
    next_button_update = next_button.update(
        interactive=True)
    adjust_button_update = adjust_button.update(
        interactive=True)
    summary_button_update = summary_button.update(
        interactive=True
    )

    # Return download button, adjust speaker, adjust audio, previous, next, adjust button
    return [download_video_update, adjust_speaker_update, adjust_audio_update, prev_button_update, next_button_update, adjust_button_update, summary_button_update]


def change_name(to_name):
    global sample_groups
    global speaker_to_name
    global current_pos

    current_speaker = sample_groups[current_pos]
    speaker_to_name[current_speaker] = to_name
    print(str(get_current_sample_file()))


def get_speakers_next(to_name):
    change_name(to_name)
    global sample_groups
    global current_pos
    if (current_pos < len(sample_groups) - 1):
        current_pos += 1
    return get_current_sample_file()


def get_speakers_previous(to_name):
    change_name(to_name)
    global current_pos
    if (current_pos > 0):
        current_pos -= 1
    return get_current_sample_file()


def start_adjust(to_name, progress=gr.Progress()):
    change_name(to_name)

    # Replacing texts
    progress(0.4, desc=current_ui_lang["progress_adjust_speaker"])
    transcribe_txt_list, subtitle_txt_list = text_file_tool.read_transcribe_subtitle_file(
        False)
    modified_transcribe = replace_text(transcribe_txt_list)
    modified_subtitle = replace_text(subtitle_txt_list)
    text_file_tool.write_transcribe_subtitle_file(
        modified_transcribe, modified_subtitle, True)

    # Get complete transcribe into string
    transcribe_txt_list, _ = text_file_tool.read_transcribe_subtitle_file(True)
    print(line for line in transcribe_txt_list)
    transcribe_txt = "\n".join(transcribe_txt_list)

    # Return to output textbox, output files, and output video
    return [
        transcribe_txt,
        [dir_adjusted_transcribe_file, dir_adjusted_subtitle_file],
        [dir_cut_video_file, dir_adjusted_subtitle_file]
    ]


def replace_text(lines):
    modified_lines = []
    for line in lines:
        for key, value in speaker_to_name.items():
            line = line.replace(key, value)
            print(f"Replacing {key} with {value}")
        modified_lines.append(line)
    print(modified_lines)
    return modified_lines