Ketengan-Diffusion-Lab/HardSubExtract

1import gradio as gr2import os3from videocr import save_subtitles_to_file4 5# Define path for demo video6DEMO_VIDEO_PATH = "demo.mp4"7 8# Function to list files in /data directory9def list_files():10    files = os.listdir('/data')11    file_paths = [f"/data/{file}" for file in files]12    return file_paths13 14# Define OCR function15def 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):16    try:17        # Ensure the output directory exists18        persistent_dir = '/data'19        if not os.path.exists(persistent_dir):20            os.makedirs(persistent_dir)21 22        # Check whether to use demo video or uploaded video23        video_path = DEMO_VIDEO_PATH if use_demo_video else input_video24 25        # Define full path for the output file26        output_path = os.path.join(persistent_dir, output_file_name)27 28        # Save the subtitles to file29        save_subtitles_to_file(30            video_path,       # Use demo video or uploaded video path31            output_path,       # Output .srt file path32            lang=language_code,33            use_gpu=use_gpu,34            time_start=start_time,35            time_end=end_time,36            conf_threshold=confidence_threshold,37            sim_threshold=similarity_threshold,38            frames_to_skip=frames_to_skip,39            crop_x=crop_x,40            crop_y=crop_y,41            crop_width=crop_width,42            crop_height=crop_height43        )44 45        return f"Subtitle extraction completed! File saved to {output_path}"46    except Exception as e:47        return f"Error: {str(e)}"48 49# Define Gradio interface50def video_ocr_interface():51    with gr.Blocks() as demo:52        with gr.Row():53            # Radio button for video selection54            use_demo_video = gr.Radio(choices=["Upload Video", "Use Demo Video"], label="Choose Video", value="Upload Video")55 56            input_video = gr.File(label="Upload Video", type="filepath", visible=True)57            output_file_name = gr.Textbox(label="Output File Name (.srt)", value="subtitle.srt")58            language_code = gr.Textbox(label="Language Code", value="ch")59            use_gpu = gr.Checkbox(label="Use GPU", value=True)60        61        with gr.Row():62            start_time = gr.Textbox(label="Start Time (HH:MM:SS)", value="00:00:00")63            end_time = gr.Textbox(label="End Time (HH:MM:SS)", value="")64        65        with gr.Row():66            confidence_threshold = gr.Slider(label="Confidence Threshold", minimum=0, maximum=100, value=75)67            similarity_threshold = gr.Slider(label="Similarity Threshold", minimum=0, maximum=100, value=80)68        69        with gr.Row():70            frames_to_skip = gr.Slider(label="Frames to Skip", minimum=0, maximum=10, value=0)71            crop_x = gr.Number(label="Crop X", value=0)72            crop_y = gr.Number(label="Crop Y", value=0)73            crop_width = gr.Number(label="Crop Width", value=0)74            crop_height = gr.Number(label="Crop Height", value=0)75 76        submit_btn = gr.Button("Start OCR")77 78        output_label = gr.Textbox(label="Status", interactive=False)79 80        # Define the file explorer component81        file_list = gr.File(label="Download Extracted .srt Files", interactive=True)82 83        # Refresh button to update the list of files84        refresh_btn = gr.Button("Refresh File List")85        86        # Define the visibility toggle for upload based on the radio button87        def toggle_video_input(selected_option):88            return {"visible": selected_option == "Upload Video"}89 90        use_demo_video.change(fn=toggle_video_input, inputs=use_demo_video, outputs=input_video)91 92        # Define what happens when the button is clicked93        submit_btn.click(94            fn=run_video_ocr,95            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],96            outputs=output_label97        )98        99        # Refresh the file explorer when the refresh button is clicked100        refresh_btn.click(fn=list_files, inputs=[], outputs=file_list)101 102    return demo103 104# Launch the Gradio interface105demo = video_ocr_interface()106demo.launch()107