Ketengan-Diffusion-Lab commited on
Commit
b24d321
1 Parent(s): 184059f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +20 -5
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
- input_video, # The uploaded video file path
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
- input_video = gr.File(label="Upload Video", type="filepath") # Changed to 'filepath'
 
 
 
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