import os
import gradio as gr
from pipelines.pipeline import InferencePipeline

pipelines = {
    "VSR(mediapipe)": InferencePipeline("./configs/LRS3_V_WER19.1.ini", device="cpu", face_track=True, detector="mediapipe"),
    "ASR": InferencePipeline("./configs/LRS3_A_WER1.0.ini", device="cpu", face_track=True, detector="mediapipe"),
    "AVSR(mediapipe)": InferencePipeline("./configs/LRS3_AV_WER0.9.ini", device="cpu", face_track=True, detector="mediapipe")
}
print("Step 0. Model has been loaded.")

def fn(pipeline_type, filename):
    print("Step 0. Video has been uploaded.")
    selected_pipeline_instance = pipelines[pipeline_type]
    print("Step 1. Video has been converted.")
    landmarks = selected_pipeline_instance.process_landmarks(filename, landmarks_filename=None)
    print("Step 2. Landmarks have been detected.")
    data = selected_pipeline_instance.dataloader.load_data(filename, landmarks)
    print("Step 3. Data has been preprocessed.")
    transcript = selected_pipeline_instance.model.infer(data)
    print("Step 4. Inference has been done.")
    print(f"transcript: {transcript}")
    return transcript

demo = gr.Blocks()

with demo:
    gr.HTML(
        """
            <div style="text-align: center; max-width: 700px; margin: 0 auto;">
              <div
                style="
                  display: inline-flex;
                  align-items: center;
                  gap: 0.8rem;
                  font-size: 1.75rem;
                "
              >
                <h1 style="font-weight: 900; margin-bottom: 7px; line-height: normal;">
                  Auto-AVSR: Audio-Visual Speech Recognition with Automatic Labels
                </h1>
              </div>
              <p style="margin-bottom: 10px; font-size: 94%">
              </p>
            </div>
        """
    )
    dropdown_list = gr.inputs.Dropdown(["ASR", "VSR(mediapipe)", "AVSR(mediapipe)"], label="model")
    video_file = gr.Video(label="INPUT VIDEO", include_audio=True)
    text = gr.Textbox(label="PREDICTION")
    btn = gr.Button("Submit").style(full_width=True)

    btn.click(fn, inputs=[dropdown_list, video_file], outputs=text)

    with gr.Accordion("Additional information", open=False):
        gr.HTML(
            """
            <div class="acknowledgments">
                <p> We used retinaface for training, but for the demo we used mediapipe </p>
                <p> We share this demo only for non-commercial purposes. </p>
            </div>
            """
        )

demo.launch()