import gradio as gr
from run import get_model, detect_video
from FakeVD.code_test import predict
import os
os.environ['GRADIO_TEMP_DIR'] = "../cache/"


def prepare_env():
    import subprocess
    # 定义要运行的命令
    command = [
        'pip', 'install', 'modelscope[audio]', '-f',
        'https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html'
    ]
    # 运行命令
    try:
        # 使用subprocess.run来执行命令
        result = subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        # 打印命令的输出
        print("STDOUT:", result.stdout.decode())
        print("STDERR:", result.stderr.decode())
    except subprocess.CalledProcessError as e:
        # 捕获并打印命令执行过程中的错误
        print("An error occurred while executing the command:", e)


prepare_env()

model = get_model()
FakeVD_model = predict.get_model()

def greet(video):
    print(video, type(video))
    generated_pred = detect_video(video_path=video, model=model)
    context_pred = predict.main(FakeVD_model, video)

    generated_output = f"Fake: {generated_pred*100:.2f}%" if generated_pred > 0.5 else f"Real: {(1-generated_pred)*100:.2f}%"
    context_output = f"Fake: {context_pred*100:.2f}%" if context_pred > 0.5 else f"Real: {(1-context_pred)*100:.2f}%"
    print(generated_output, '\n', context_output)
    return generated_output, context_output

with gr.Blocks() as demo:
    gr.Markdown("# Fake Video Detector")
    with gr.Tabs():
        with gr.TabItem("Video Detect"):
            with gr.Column():
                video_input = gr.Video(height=330)
                video_button = gr.Button("detect")
                detect_output = gr.Textbox(label="AI detect result")
                context_output = gr.Textbox(label="Context detect result")

    video_button.click(greet, inputs=video_input, outputs=[detect_output, context_output])

demo.launch()