app.py

import gradio as gr
import os
from videocr import save_subtitles_to_file

# Define path for demo video
DEMO_VIDEO_PATH = "demo.mp4"

# Function to list files in /data directory
def list_files():
    files = os.listdir('/data')
    file_paths = [f"/data/{file}" for file in files]
    return file_paths

# Define OCR function
def run_video_ocr(use_demo_video, input_video, output_file_name, language_code, use_gpu, start_time, end_time, confidence_threshold, similarity_threshold, frames_to_skip, crop_x, crop_y, crop_width, crop_height):
    try:
        # Ensure the output directory exists
        persistent_dir = '/data'
        if not os.path.exists(persistent_dir):
            os.makedirs(persistent_dir)

        # Check whether to use demo video or uploaded video
        video_path = DEMO_VIDEO_PATH if use_demo_video else input_video

        # Define full path for the output file
        output_path = os.path.join(persistent_dir, output_file_name)

        # Save the subtitles to file
        save_subtitles_to_file(
            video_path,       # Use demo video or uploaded video path
            output_path,       # Output .srt file path
            lang=language_code,
            use_gpu=use_gpu,
            time_start=start_time,
            time_end=end_time,
            conf_threshold=confidence_threshold,
            sim_threshold=similarity_threshold,
            frames_to_skip=frames_to_skip,
            crop_x=crop_x,
            crop_y=crop_y,
            crop_width=crop_width,
            crop_height=crop_height
        )

        return f"Subtitle extraction completed! File saved to {output_path}"
    except Exception as e:
        return f"Error: {str(e)}"

# Define Gradio interface
def video_ocr_interface():
    with gr.Blocks() as demo:
        with gr.Row():
            # Radio button for video selection
            use_demo_video = gr.Radio(choices=["Upload Video", "Use Demo Video"], label="Choose Video", value="Upload Video")

            input_video = gr.File(label="Upload Video", type="filepath", visible=True)
            output_file_name = gr.Textbox(label="Output File Name (.srt)", value="subtitle.srt")
            language_code = gr.Textbox(label="Language Code", value="ch")
            use_gpu = gr.Checkbox(label="Use GPU", value=True)
        
        with gr.Row():
            start_time = gr.Textbox(label="Start Time (HH:MM:SS)", value="00:00:00")
            end_time = gr.Textbox(label="End Time (HH:MM:SS)", value="")
        
        with gr.Row():
            confidence_threshold = gr.Slider(label="Confidence Threshold", minimum=0, maximum=100, value=75)
            similarity_threshold = gr.Slider(label="Similarity Threshold", minimum=0, maximum=100, value=80)
        
        with gr.Row():
            frames_to_skip = gr.Slider(label="Frames to Skip", minimum=0, maximum=10, value=0)
            crop_x = gr.Number(label="Crop X", value=0)
            crop_y = gr.Number(label="Crop Y", value=0)
            crop_width = gr.Number(label="Crop Width", value=0)
            crop_height = gr.Number(label="Crop Height", value=0)

        submit_btn = gr.Button("Start OCR")

        output_label = gr.Textbox(label="Status", interactive=False)

        # Define the file explorer component
        file_list = gr.File(label="Download Extracted .srt Files", interactive=True)

        # Refresh button to update the list of files
        refresh_btn = gr.Button("Refresh File List")
        
        # Define the visibility toggle for upload based on the radio button
        def toggle_video_input(selected_option):
            return {"visible": selected_option == "Upload Video"}

        use_demo_video.change(fn=toggle_video_input, inputs=use_demo_video, outputs=input_video)

        # Define what happens when the button is clicked
        submit_btn.click(
            fn=run_video_ocr,
            inputs=[use_demo_video, input_video, output_file_name, language_code, use_gpu, start_time, end_time, confidence_threshold, similarity_threshold, frames_to_skip, crop_x, crop_y, crop_width, crop_height],
            outputs=output_label
        )
        
        # Refresh the file explorer when the refresh button is clicked
        refresh_btn.click(fn=list_files, inputs=[], outputs=file_list)

    return demo

# Launch the Gradio interface
demo = video_ocr_interface()
demo.launch()