File size: 4,374 Bytes
a38a1e2 b24d321 c64f927 b24d321 31363a4 a38a1e2 b24d321 a38a1e2 80995ef a38a1e2 b24d321 a38a1e2 b24d321 a38a1e2 b24d321 a38a1e2 31363a4 b24d321 a38a1e2 b24d321 a38a1e2 31363a4 a38a1e2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
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()
|