Ketengan-Diffusion-Lab
commited on
Commit
•
b24d321
1
Parent(s):
184059f
Update app.py
Browse files
app.py
CHANGED
@@ -2,6 +2,9 @@ import gradio as gr
|
|
2 |
import os
|
3 |
from videocr import save_subtitles_to_file
|
4 |
|
|
|
|
|
|
|
5 |
# Function to list files in /data directory
|
6 |
def list_files():
|
7 |
files = os.listdir('/data')
|
@@ -9,19 +12,22 @@ def list_files():
|
|
9 |
return file_paths
|
10 |
|
11 |
# Define OCR function
|
12 |
-
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):
|
13 |
try:
|
14 |
# Ensure the output directory exists
|
15 |
persistent_dir = '/data'
|
16 |
if not os.path.exists(persistent_dir):
|
17 |
os.makedirs(persistent_dir)
|
18 |
-
|
|
|
|
|
|
|
19 |
# Define full path for the output file
|
20 |
output_path = os.path.join(persistent_dir, output_file_name)
|
21 |
|
22 |
# Save the subtitles to file
|
23 |
save_subtitles_to_file(
|
24 |
-
|
25 |
output_path, # Output .srt file path
|
26 |
lang=language_code,
|
27 |
use_gpu=use_gpu,
|
@@ -44,7 +50,10 @@ def run_video_ocr(input_video, output_file_name, language_code, use_gpu, start_t
|
|
44 |
def video_ocr_interface():
|
45 |
with gr.Blocks() as demo:
|
46 |
with gr.Row():
|
47 |
-
|
|
|
|
|
|
|
48 |
output_file_name = gr.Textbox(label="Output File Name (.srt)", value="subtitle.srt")
|
49 |
language_code = gr.Textbox(label="Language Code", value="ch")
|
50 |
use_gpu = gr.Checkbox(label="Use GPU", value=True)
|
@@ -74,10 +83,16 @@ def video_ocr_interface():
|
|
74 |
# Refresh button to update the list of files
|
75 |
refresh_btn = gr.Button("Refresh File List")
|
76 |
|
|
|
|
|
|
|
|
|
|
|
|
|
77 |
# Define what happens when the button is clicked
|
78 |
submit_btn.click(
|
79 |
fn=run_video_ocr,
|
80 |
-
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],
|
81 |
outputs=output_label
|
82 |
)
|
83 |
|
|
|
2 |
import os
|
3 |
from videocr import save_subtitles_to_file
|
4 |
|
5 |
+
# Define path for demo video
|
6 |
+
DEMO_VIDEO_PATH = "/data/demo.mp4"
|
7 |
+
|
8 |
# Function to list files in /data directory
|
9 |
def list_files():
|
10 |
files = os.listdir('/data')
|
|
|
12 |
return file_paths
|
13 |
|
14 |
# Define OCR function
|
15 |
+
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):
|
16 |
try:
|
17 |
# Ensure the output directory exists
|
18 |
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 video
|
23 |
+
video_path = DEMO_VIDEO_PATH if use_demo_video else input_video
|
24 |
+
|
25 |
# Define full path for the output file
|
26 |
output_path = os.path.join(persistent_dir, output_file_name)
|
27 |
|
28 |
# Save the subtitles to file
|
29 |
save_subtitles_to_file(
|
30 |
+
video_path, # Use demo video or uploaded video path
|
31 |
output_path, # Output .srt file path
|
32 |
lang=language_code,
|
33 |
use_gpu=use_gpu,
|
|
|
50 |
def video_ocr_interface():
|
51 |
with gr.Blocks() as demo:
|
52 |
with gr.Row():
|
53 |
+
# Radio button for video selection
|
54 |
+
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)
|
|
|
83 |
# Refresh button to update the list of files
|
84 |
refresh_btn = gr.Button("Refresh File List")
|
85 |
|
86 |
+
# Define the visibility toggle for upload based on the radio button
|
87 |
+
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 clicked
|
93 |
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_label
|
97 |
)
|
98 |
|