File size: 3,106 Bytes
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 |
import gradio as gr
import os
from videocr import save_subtitles_to_file
# Define OCR function
def run_video_ocr(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 = '/persistent'
if not os.path.exists(persistent_dir):
os.makedirs(persistent_dir)
# 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(
input_video.name, # The uploaded video file
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():
input_video = gr.File(label="Upload Video", type="file")
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 what happens when the button is clicked
submit_btn.click(
fn=run_video_ocr,
inputs=[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
)
return demo
# Launch the Gradio interface
demo = video_ocr_interface()
demo.launch()
|